(in-package :user)

(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")
(defparameter *flux* (make-array 1000 :initial-element 1))

(defparameter *D* 0.9
  "neutron diffusion coefficient (cm)")
(defparameter *sigma-a* 0.066
  "absorption cross section (1/cm)")
(defparameter *v-sigma-f* 0.07
  "fission cross section times neutrons per fission (1/cm)")
(defparameter *k-tolerance* 1e-4 ;diverges if lower, 1e-5 too much lower than 1e-3
  "acceptable uncertainty of criticality constant")
(defparameter *flux-tolerance* 1e-3 ;needs to be lower for accuracy, but too slow
  "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))) 

(defun phi-center (phiE phiW source dx D sigma-a)
  (/
    (+ (/ (+ phiW phiE) dx dx) (/ source D))
    (+ (/ 2 dx dx) (/ sigma-a D)))) 

(defun phi-edge (phiOpp albedo source dx D sigma-a)
  (/
    (+ (/ (* 2 phiOpp) dx dx) (/ source D))
    (+
      (/ 2 dx dx)
      (/ (- 1 albedo) (* (+ 1 albedo) D dx))
      (/ sigma-a D))))

(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 set-source (flux v-sigma-f k)
  (multiply-array-scalar flux (* v-sigma-f (/ k))))

(defun gauss-seidel (source flux dx D sigma-a albedoL albedoR)
  (let ((flux^ nil)
        (done t))
    (do () (nil)
      (incf *iteration-count-inner*)
      (setf flux^ (copy-array flux))
      (setf (aref flux 0) 
        (phi-edge (aref flux 1) albedoL (aref source 0) dx D sigma-a))
      (setf (aref flux (1- (length flux)))
        (phi-edge (aref flux (- (length flux) 2))
          albedoR (aref source (1- (length flux))) dx D sigma-a))
      (do ((i 1 (1+ i)))
          ((>= i (- (length flux) 1)))
        (setf (aref flux i) 
          (phi-center 
            (aref flux (1+ i)) (aref flux (1- i)) 
            (aref source i) dx D sigma-a)))
      (setf done t)
      (do ((i 0 (1+ i)))
          ((>= i (length flux)))
        (when (> (aref flux^ i) 1e-3)
          (when (< (* (aref flux^ i) *flux-tolerance*) (abs (- (aref flux i) (aref flux^ i))))
            (setf done nil))))
      (when done
        (return-from gauss-seidel flux)))))

(defun jacobian (source flux dx D sigma-a albedoL albedoR)
  (let ((flux^ nil)
        (done t))
    (do () (nil)
      (incf *iteration-count-inner*)
      (setf flux^ (copy-array flux))
      (setf (aref flux 0) 
        (phi-edge (aref flux^ 1) albedoL (aref source 0) dx D sigma-a))
      (setf (aref flux (1- (length flux)))
        (phi-edge (aref flux^ (- (length flux) 2))
          albedoR (aref source (1- (length flux))) dx D sigma-a))
      (do ((i 1 (1+ i)))
          ((>= i (- (length flux) 1)))
        (setf (aref flux i) 
          (phi-center 
            (aref flux^ (1+ i)) (aref flux^ (1- i)) 
            (aref source i) dx D sigma-a)))
      (setf done t)
      (do ((i 0 (1+ i)))
          ((>= i (length flux)))
        (when (> (aref flux^ i) 1e-3)
          (when (< (* (aref flux^ i) *flux-tolerance*) (abs (- (aref flux i) (aref flux^ i))))
            (setf done nil))))
      (when done
        (return-from jacobian flux)))))

(defun criticality-fixed (flux-solver flux dx D v-sigma-f sigma-a albedoL albedoR)
  (let ((k 0.1)
        (k^ nil)
        (k-low 0)
        (k-high nil)
        (flux-sum (loop for x across flux sum x))
        (source nil))
    (do () (nil)
      (incf *iteration-count-outer*)
      ;normalize flux to prevent underflow
      ;this line prevents changes to flux from being seen outside function
      (setf flux (multiply-array-scalar flux (/ (length flux) flux-sum)))
      (setf source (set-source flux v-sigma-f k))
      ;flux-solver modifies its flux parameter
      (funcall flux-solver source flux dx D sigma-a albedoL albedoR)
      (setf flux-sum (loop for x across flux sum x))
      (setf k^ k)
      (setf k (* k^ flux-sum (/ (length 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))))
      ;(print k)
      (defparameter *flux* flux)
      (when (not (null k-high))
        (when (< (- k-high k-low) (* k-high *k-tolerance*))
          (return-from criticality-fixed k))))))

(defun expand-flux (flux dx D v-sigma-f sigma-a k)
  (setf 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))
      (setf (aref result-flux i)
        (phi-center
          (aref result-flux (1+ i)) (aref result-flux (1- i))
          (* v-sigma-f (/ k) 0.5 (+ (aref result-flux (1+ i)) (aref result-flux (1- i))))
          dx D sigma-a))))
  result-flux)

(defun criticality-refine (flux-solver slab-length D v-sigma-f sigma-a albedoL albedoR)
  (let ((k 0)
    (flux (make-array 3 :initial-element 1)))
    (do () (nil)
      (incf *iteration-count-mesh*)
      (setf dx (/ slab-length (1- (length flux))))
      (setf k^ k)
      (setf k (criticality-fixed flux-solver flux dx D v-sigma-f sigma-a albedoL albedoR))
      (when (< (abs (- k k^)) (* k *k-tolerance*))
        (return-from criticality-refine k))
      (setf flux *flux*)
      ;(print "mesh.....k")
      ;(print (list (length flux) k))
      (setf flux (expand-flux flux (/ dx 2) D v-sigma-f sigma-a k)))))

(defun criticality-search (flux-solver D v-sigma-f sigma-a albedoL albedoR)
  (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")
  (setf low 0)
  (setf k-low 0)
  (setf high nil)
  (setf k-high nil)
  (setf slab-length 1)
  (do () (nil)
    (incf *iteration-count-length*)
    (setf k (criticality-refine flux-solver slab-length D v-sigma-f sigma-a 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)))))))

; this is the destruction operator
; represented in Duderstadt by matrix M
; was meant to be used in direct method
(defun set-destruction (nodes dx D sigma-a albedoL albedoR)
  (setf M (make-array (list nodes nodes) :initial-element 0))
  ;set diagonal
  (dotimes (i nodes)
    (setf 
      (aref M i i)
      (+ sigma-a (* 2 D (/ dx) (/ dx)))))
  ;set next to diagonals (except first and last rows)
  (dotimes (i (- nodes 2))
    (setf
      (aref M (1+ i) i)
      (* -1 D (/ dx) (/ dx)))
    (setf 
      (aref M (+ i 1) (+ i 2))
      (* -1 D (/ dx) (/ dx))))
  ;add albedo to corners
  (setf
    (aref M 0 0)
    (+ 
      (aref M 0 0)
      (/ (- 1 albedoL) dx (+ 1 albedoL))))
  (setf
    (aref M (1- nodes) (1- nodes))
    (+ 
      (aref M (1- nodes) (1- nodes))
      (/ (- 1 albedoR) dx (+ 1 albedoR))))
  ;set next to corners
  (setf
    (aref M 0 1)
    (* -2 D (/ dx) (/ dx)))
  (setf 
    (aref M (- nodes 1) (- nodes 2))
    (* -2 D (/ dx) (/ dx)))
  M)

;criticality search could be called like this:
;(criticality-search #'gauss-seidel *D* *v-sigma-f* *sigma-a* *left-albedo* *right-albedo*)