(defparameter *iteration-count-inner* 0
  "keeps track of (fixed-source) flux iterations to measure convergence speed")
(defparameter *iteration-count-outer* 0
  "keeps track of source iterations to measure convergence speed")
(defparameter *iteration-count-mesh* 0
  "keeps track of number of mesh refinements to measure convergence speed")
(defparameter *iteration-count-length* 0
  "keeps track of number of length guesses to measure convergence speed")

; diverges if lower, 1e-5 too much lower than 1e-3
(defparameter *k-tolerance* 1e-4 
  "acceptable uncertainty of criticality constant")
; needs to be lower for accuracy, but too slow
(defparameter *flux-tolerance* 1e-3 
  "acceptable uncertainty of flux")
(defparameter *left-albedo* 0
  "ratio of return current to outgoing current at left edge")
(defparameter *right-albedo* 0
  "ratio of return current to outgoing current at right edge")

;; I copied this function from Vladimir Zolotykh, 
;; who in turn copied it from Barry Margolin.
(defun copy-array (array) 
  (let ((new (make-array (array-dimensions array) :displaced-to array))) 
    (adjust-array new (array-dimensions array) :displaced-to nil))) 

; creates new array identical to input
; uses the same material because material is constant
; much more efficient to copy that
; but flux changes and will not always be the same
; between the two arrays later
; since I create a new node and flux is a small
; object (just a number) it will probably be copied
; instead of referenced *crosses fingers*
(defun copy-flux (array)
  (let ((result (make-array (array-dimensions array))))
    (dotimes (i (array-dimension array 0))
      (dotimes (j (array-dimension array 1))
        (dotimes (k (array-dimension array 2))
          (setf
           (aref result i j k)
           (make-node
            :material (node-material (aref array i j k))
            :flux (node-flux (aref array i j k)))))))
    result))

(defstruct material 
  D
  v-sigma-f
  sigma-a)

(let ((xo 0) (dx 1) (slab-length 1))
  (defun simple-material-selector (i j k)
    (let ((x (+ xo (* i dx))))
      (make-material :D .9 :v-sigma-f .07 :sigma-a .066)) )
  (defun halved-material-selector (i j k)
    (let ((x (+ xo (* i dx))))
      (if (> x (/ slab-length 2.0))
          (make-material :D .9 :v-sigma-f .07 :sigma-a .066)
          (make-material :D .9 :v-sigma-f .01 :sigma-a .066))))
  ; thinking of removing xo and just using 0 always
  (defun set-scale (dx-new xo-new slab-length-new)
    (setf xo xo-new)
    (setf dx dx-new)
    (setf slab-length slab-length-new)))

(defun print-node (n stream depth) 
  (format stream "n~a" (node-flux n)))

(defstruct (node (:print-function print-node))
  material
  flux)      

(defun phi-1d (phiE phiW source dx material)
  (let ((D (material-D material))
        (sigma-a (material-sigma-a material)))
    (min
     (max (node-flux phiE) (node-flux phiW))
     (/
      (+ (/ (+ (node-flux phiW) (node-flux phiE)) dx dx) 
         (/ source D))
      (+ (/ 2 dx dx) (/ sigma-a D))))))

(defun phi-edge (phiOpp albedo source dx material)
  (let ((D (material-D material))
        (sigma-a (material-sigma-a material)))
    (make-node 
     :flux
     (/
      (+ (/ (* 2 (node-flux phiOpp)) dx dx) (/ source D))
      (+
       (/ 2 dx dx)
       (/ (- 1 albedo) (* (+ 1 albedo) D dx))
       (/ sigma-a D)))
     :material material)))

;; calculates flux at one node
(defun phi-calc (source dx material
                 x1 x2
                 y1 y2
                 z1 z2)
  (declare (single-float source dx x1 x2 y1 y2 z1 z2))
  (declare (optimize (speed 3) (compilation-speed 0) 
                     (debug 0) (safety 0)))
  ; n is the number of no-return boundaries
  (let ((x-contrib 0.0) (y-contrib 0.0) 
        (z-contrib 0.0) (n 0)
        (D (material-D material)) 
        (sigma-a (material-sigma-a material)))
    (declare (single-float 
              x-contrib y-contrib z-contrib D sigma-a))
    (declare (fixnum n))
    ; if flux very small on one side, treat that side
    ; as no-return boundary
    (if (< (min x1 x2)
           (* (max x1 x2)                
              (1- (/ 4 (+ 2 (/ dx D))))))
        (progn
          (setq x-contrib (* 2 (max x1 x2)))
          (setf n  (1+ n)))
        (setq x-contrib  (+ x1 x2)))
    (if (< (min y1 y2) 
           (* (max y1 y2)
              (- (/ 4 (+ 2 (/ dx D))) 1)))
        (progn
          (setq y-contrib (* 2 (max y1 y2)))
          (setf n (1+ n)))
        (setq y-contrib (+ y1 y2)))
    (if (< (min z1 z2) 
           (* (max z1 z2)
              (- (/ 4 (+ 2 (/ dx D))) 1)))
        (progn
          (setq z-contrib (* 2 (max z1 z2)))
          (setf n (1+ n)))
        (setq z-contrib (+ z1 z2)))
    (let ((temp (/ dx D))
	  (sqrt-n (cond ((zerop n) 0.0)
			((eql n 1) 1.0)
			((eql n 2) 1.4)
			(t 1.7))))
      (declare (single-float temp sqrt-n))
      (the single-float
      (/ (+ x-contrib y-contrib z-contrib 
            (* dx temp source))
         (+ (* dx temp sigma-a)
            6 (* sqrt-n temp)))))))

(defun multiply-array-scalar (inputarray scalar)
  (let ((resultarray (copy-array inputarray)))
    (dotimes (i (array-total-size resultarray))
      (setf (row-major-aref resultarray i)
            (* scalar (row-major-aref resultarray i))))
    resultarray))

(defun multiply-flux-scalar (flux scalar)
  (let ((result (copy-array flux)))
    (dotimes (i (array-total-size result))
      (setf (row-major-aref result i)
            (make-node
             :flux      
             (* scalar (node-flux (row-major-aref result i)))
             :material (node-material (row-major-aref flux i)))))
    result))

(defun sum-array-flux (flux)
  (let ((result 0))
    (dotimes (i (array-dimension flux 0))
      (dotimes (j (array-dimension flux 1))
        (dotimes (k (array-dimension flux 2))
          (setf result 
                (+ result 
                   (node-flux
                    (aref flux i j k)))))))
    result))

(defun set-source (flux crit)
  (let ((source (make-array (array-dimensions flux))))
    (dotimes (i (array-dimension flux 0))
      (dotimes (j (array-dimension flux 1))
        (dotimes (k (array-dimension flux 2))
          (setf (aref source i j k)
                (* (node-flux (aref flux i j k))
                   (material-v-sigma-f 
                    (node-material (aref flux i j k)))
                   (/ crit))))))
    source))

(defun gauss-seidel (source flux dx)
  ;(declare fixnum dx)
  (declare (optimize (speed 3) (compilation-speed 0) 
                     (debug 0) (safety 0)))
  (let ((flux^ nil)
        (done t))
    (do () (nil)
      (incf *iteration-count-inner*)
      (setf flux^ (copy-flux flux))
      (setf done t)
      (dotimes (i (array-dimension flux 0))
        (dotimes (j (array-dimension flux 1))
          (dotimes (k (array-dimension flux 2))
            (setf (node-flux (aref flux i j k))
                  (phi-calc
                   (aref source i j k)
                   dx 
                   (node-material (aref flux i j k))
                   (if (eql (1+ i) (array-dimension flux 0))
                       0.0
                       (node-flux (aref flux (1+ i) j k)))
                   (if (eql i 0)
                       0.0
                       (node-flux (aref flux (1- i) j k)))
                   (if (eql (1+ j) (array-dimension flux 1))
                       0.0
                       (node-flux (aref flux i (1+ j) k)))
                   (if (eql j 0)
                       0.0
                       (node-flux (aref flux i (1- j) k)))
                   (if (eql (1+ k) (array-dimension flux 2))
                       0.0
                       (node-flux (aref flux i j (1+ k))))
                   (if (eql k 0)
                       0.0
                       (node-flux (aref flux i j (1- k))))))
            (when done
              (when (> (node-flux (aref flux^ i j k)) 1e-3)
                (when (< (* (node-flux (aref flux^ i j k)) *flux-tolerance*)
                         (abs (- (node-flux (aref flux i j k))
                                 (node-flux (aref flux^ i j k)))))
                  (setf done nil)))))))
      (when done
        (return-from gauss-seidel flux)))))

(defun criticality-fixed (flux dx)
  (let ((k 0.1)
        (k^ nil)
        (k-low 0)
        (k-high nil)
        (flux-sum (sum-array-flux flux))
        (source nil))
    (do () (nil)
      (incf *iteration-count-outer*)
      ;normalize flux to prevent underflow
      ;this line prevents criticality-fixed from being destructive
      (setf flux (multiply-flux-scalar 
                  flux (/ (array-total-size flux) flux-sum)))
      (setf source (set-source flux k))
      ;flux-solver modifies its flux parameter
      (setf flux (gauss-seidel source flux dx))
      (setf flux-sum (sum-array-flux flux))
      (setf k^ k)
      (setf k (* k^ flux-sum (/ (array-total-size flux))))
      (if (> k k^)
          (setf k-low k^)
          (setf k-high k^))
      (if (null k-high)
          (setf k (* 2 k^))
          (setf k (* 0.5 (+ k-high k-low))))
      (when (not (null k-high))
        (when (< (- k-high k-low) (* k-high *k-tolerance*))
          (return-from criticality-fixed (values k flux)))))))

(defun expand-flux (flux dx material-selector k)
  (let ((result-flux 
         (make-array (1- (* 2 (length flux))) :initial-element nil)))
    (dotimes (i (array-total-size result-flux))
      (when (evenp i) (setf (aref result-flux i) (aref flux (/ i 2)))))
    (dotimes (i (array-total-size result-flux))
      (when (null (aref result-flux i))
        (let ((material (funcall material-selector i)))
          (setf (aref result-flux i)
                (phi-center
                 (aref result-flux (1+ i)) (aref result-flux (1- i))
                 (* (material-v-sigma-f material) (/ k) 0.5 
                    (+ (node-flux (aref result-flux (1+ i)))
                       (node-flux (aref result-flux (1- i)))))
                 dx 
                 material)))))
    result-flux))

(defun number-odd (i j k)
  (let ((n 0))
    (when (oddp i) (incf n))
    (when (oddp j) (incf n))
    (when (oddp k) (incf n))
    n))    

; note, the argument dx is the new dx for the bigger array
(defun expand-flux-3d (flux dx material-selector crit)
  ; let result-flux be an array that is twice as long
  ; in each direction as flux
  (let ((result-flux
         (make-array
          (mapcar 
           #'(lambda (x) (1- (* 2 x))) 
           (array-dimensions flux)))))
    ; fill in the gaps based on a 1-D estimate
    ; start with nodes that have the fewest
    ; odd indices to guarantee that there
    ; will be enough flux information
    ; in nearby nodes
    (dotimes (n 4) ; n is the number of odd indices
      (dotimes (i (array-dimension result-flux 0))
        (dotimes (j (array-dimension result-flux 1))
          (dotimes (k (array-dimension result-flux 2))
            (when (eql n (number-odd i j k)) 
              (let ((material (funcall material-selector i j k)))
                (setf
                 (aref result-flux i j k)
                 (make-node
                  :material material
                  :flux
                  (cond
                    ((oddp i) (phi-1d 
                               (aref result-flux (1+ i) j k) 
                               (aref result-flux (1- i) j k)
                               (* (material-v-sigma-f material) 
                                  (/ crit) 0.5 
                                  (+ (node-flux (aref result-flux (1+ i) j k))
                                     (node-flux (aref result-flux (1- i) j k))))
                               dx 
                               material))
                    ((oddp j) (phi-1d 
                               (aref result-flux i (1+ j) k) 
                               (aref result-flux i (1- j) k)
                               (* (material-v-sigma-f material) 
                                  (/ crit) 0.5 
                                  (+ (node-flux (aref result-flux i (1+ j) k))
                                     (node-flux (aref result-flux i (1- j) k))))
                               dx 
                               material))
                    ((oddp k) (phi-1d 
                               (aref result-flux i j (1+ k)) 
                               (aref result-flux i j (1- k))
                               (* (material-v-sigma-f material) 
                                  (/ crit) 0.5 
                                  (+ (node-flux (aref result-flux i j (1- k)))
                                     (node-flux (aref result-flux i j (1+ k)))))
                               dx 
                               material))
                    ; no odd indices? that's the easy case
                    ; just copy the value from flux
                    (t (node-flux (aref flux (/ i 2) (/ j 2) (/ k 2)))))))))))))
    result-flux))

(defun criticality-refine
    (universe-length material-selector)
  (let ((crit 0)
        (crit^ 0)
        (flux (make-array '(5 5 5)))
        (dx (/ universe-length 4)))
    (set-scale dx 0 universe-length)
    (dotimes (i (array-dimension flux 0))
      (dotimes (j (array-dimension flux 1))
        (dotimes (k (array-dimension flux 2))
          (setf (aref flux i j k)
                (make-node
                 :flux 1
                 :material (funcall material-selector i j k))))))
    (do () (nil)
      (incf *iteration-count-mesh*)
      (setf crit^ crit)
      (setf dx (/ universe-length (1- (array-dimension flux 0))))
      (multiple-value-setq (crit flux) 
        (criticality-fixed flux dx))
      (print "k ... array-size")
      (print crit)
      (print (array-dimension flux 0))
      (print " ")
      (when (< (abs (- crit crit^)) (* crit *k-tolerance*))
        (return-from criticality-refine (values crit flux)))
      (set-scale (/ dx 2) 0 universe-length)
      (setf flux (expand-flux-3d flux (/ dx 2) material-selector crit)))))

(defun criticality-search (flux-solver material-selector albedoL albedoR)
  (setf *iteration-count-inner* 0)
  (setf *iteration-count-outer* 0)
  (setf *iteration-count-mesh* 0)
  (setf *iteration-count-length* 0)
  (let ((low 0) 
        (k-low 0) 
        (high nil) 
        (k-high nil) 
        (slab-length 1) 
        (k nil))
    (do () (nil)
      (incf *iteration-count-length*)
      (setf k (criticality-refine flux-solver slab-length material-selector
                                  albedoL albedoR))
      (print "length...k")
      (print (list slab-length k))
      (when (< k 1) 
        (setf k-low k)
        (setf low slab-length))
      (when (> k 1) 
        (setf k-high k)
        (setf high slab-length))
      (when (null high)
        (setf slab-length (* 2 slab-length)))
      (when (numberp high)
        (setf slab-length (* 0.5 (+ low high))))
      (when (numberp high)
        (when (< (- k-high k-low) *k-tolerance*)
          (return-from criticality-search (* 0.5 (+ low high))))))))

(defparameter *timing-data* ())

(defmacro with-timing (label &body body)
  (let ((start (gensym)))
    `(let ((,start (get-internal-run-time)))
      (unwind-protect (progn ,@body)
        (push (list ',label ,start (get-internal-run-time)) *timing-data*)))))

(defun clear-timing-data ()
  (setf *timing-data* ()))

(defun show-timing-data ()
  (loop for (label time count time-per %-of-total) in (compile-timing-data) do
        (format t "~3d% ~a: ~d ticks over ~d calls for ~d per.~%" 
                %-of-total label time count time-per)))

(defun compile-timing-data () 
  (loop with timing-table = (make-hash-table)
        with count-table = (make-hash-table)
        for (label start end) in *timing-data*
        for time = (- end start)
        summing time into total
        do
        (incf (gethash label timing-table 0) time)
        (incf (gethash label count-table 0))
        finally 
        (return
          (sort
           (loop for label being the hash-keys in timing-table collect
                 (let  ((time (gethash label timing-table))
                        (count (gethash label count-table)))
                   (list label time count (round (/ time count)) 
                         (round (* 100 (/ time total))))))
           #'> :key #'fifth))))

(clear-timing-data)

;(criticality-search #'gauss-seidel #'simple-material-selector 
;                    *left-albedo* *right-albedo*)