; Eric Lavigne
; Kinetics
; Assignment 4

(defpackage :rkf (:export solve-system *a +a))
(in-package :rkf)

(defun solve-system (fp x time final-time global-tolerance)
  (let ((tolerance (/ global-tolerance (- final-time time) 1d0))
        (dt (/ (- final-time time) 1d1))
        (next-x (make-array (length x) :element-type 'double-float)))
    (loop
      (when (> time final-time) (return))
      (when (< dt 1d-9) (throw 'sub-nanosecond-timescale-requested nil))
      (catch 'error-greater-than-tolerance
             (setf next-x (increment-system fp x time dt tolerance))
             (format t "~A" time)
             (setf time (+ time dt))
             (dotimes (i (length x))
               (format t "     ~A" (elt x i)))
             (format t "~%")
             (setf x next-x)
             (setf dt (* dt 2.2d0)))
      (setf dt (* dt 0.5d0)))))

(defun increment-system (fp x time dt tolerance)
  (let* ((k1 (*a dt (funcall fp x time)))
         (k2 (*a dt (funcall fp (+a x (*a 1/4 k1)) (+ time (/ dt 4)))))
         (k3 (*a dt (funcall fp (+a x (*a 3/32 k1) (*a 9/32 k2)) 
                             (+ time (* 3/8 dt)))))
         (k4 (*a dt (funcall fp (+a x (*a 1932/2197 k1) (*a -7200/2197 k2) 
                                    (*a 7296/2197 k3)) 
                             (+ time (* 12/13 dt)))))
         (k5 (*a dt (funcall fp (+a x (*a 439/216 k1) (*a -8 k2) 
                                    (*a 3680/513 k3) (*a -845/4104 k4))
                             (+ time dt))))
         (k6 (*a dt (funcall fp (+a x (*a -8/27 k1) (*a 2 k2) 
                                    (*a -3544/2565 k3) (*a 1859/4104 k4) 
                                                        (*a -11/40 k5))
                             (+ time (/ dt 2)))))
         (error (abs-array 
                 (*a (/ 1 dt) (+a (*a 1/360 k1) (*a -128/4275 k3) 
                                  (*a -2197/75240 k4) (*a 1/50 k5) 
                                  (*a 2/55 k6))))))
    (when (> error (* 5d-1 tolerance))
      (throw 'error-greater-than-tolerance nil))
    (+a (*a 1d0 x) (*a 25/216 k1) (*a 1408/2565 k3) 
               (*a 2197/4104 k4) (*a -1/5 k5))))

(defun +a (&rest x)
  (reduce #'(lambda (x1 x2) (mapcar #'+ x1 x2)) x))

(defun *a (x1 x2) ; (scalar * array) or (array * array)
  (if (numberp x1) (mapcar #'(lambda (x) (* x1 x)) x2)
    (mapcar #'* x1 x2)))

(defun abs-array (x)
  (sqrt (reduce #'+ (*a x x))))

(defparameter *lambda-u235* 
  (mapcar #'(lambda (half-life) (/ (log 2d0) half-life)) 
    '(55.7 22.7 6.22 2.3 0.61 0.23)))

(defparameter *beta-fractions-u235* 
  (*a 1d0 '(0.0330667 0.2190095 0.1959397 0.3949553 0.1150415 0.041987080)))

(defun steady-state-values (lifetime beta lambda)
  (cons 1d0 (mapcar #'/ (*a (/ 1d0 lifetime) beta) lambda)))

(defun kinetics-equations (reactivity lifetime beta lambda)
  #'(lambda (x time)
      (let ((delayed-part (+a (*a (/ (first x) lifetime) beta)
                              (*a -1d0 (*a lambda (rest x))))))
        (cons (- (* (first x) reactivity (/ 1d0 lifetime))
                 (reduce #'+ delayed-part))
              delayed-part))))

; p285 example 1 (for testing the algorithm)
(defun circuit-fp (x time)
  (list (+ (* -4 (first x)) (* 3 (second x)) 6)
        (+ (* -2.4 (first x)) (* 1.6 (second x)) 3.6)))
(solve-system #'circuit-fp '(0d0 0d0) 0d0 0.5d0 1d-8)

; Exercise 1: p288 example 2 
(solve-system #'(lambda (x time)
                  (list (second x)
                        (+ (* (exp (* 2 time)) (sin time))
                           (* -2 (first x))
                           (* 2 (second x)))))
              '(-0.4d0 -0.6d0) 0d0 1d0 1d-7)

;Exercise 2 (and 3)
(defparameter *beta* (*a 6.4d-3 *beta-fractions-u235*))
(defparameter *lifetime* 1d-3)

; part a: reactivity = 0.001
(solve-system 
 (kinetics-equations 1d-3 *lifetime* *beta* *lambda-u235*)
 (steady-state-values *lifetime* *beta* *lambda-u235*)
 0d0 10d0 1d-5)

; part b: reactivity = -0.0063
(solve-system 
 (kinetics-equations -6.3d-3 *lifetime* *beta* *lambda-u235*)
 (steady-state-values *lifetime* *beta* *lambda-u235*)
 0d0 10d0 1d-5)