;; Copyright (c) Camden Dixie O'Brien
;; SPDX-License-Identifier: AGPL-3.0-only

(in-package :ham)

(eval-when (:load-toplevel :execute)
  (maxima-interface:maxima-eval
   `((maxima::$load)
     ,(namestring (asdf:system-relative-pathname :ham "ham.mac")))))

(defun enmaxima-symbol (sym)
  (intern (format nil "$~a" (symbol-name sym))
		  (find-package :maxima)))

(defun enmaxima-lagrangian (expr)
  (defun prepare (expr)
	(cond
	  ((symbolp expr) (enmaxima-symbol expr))
	  ((and (listp expr) (eq 'dot (car expr)))
	   (append '(maxima::%derivative)
			   (mapcar #'prepare (cdr expr))
			   '(maxima::$t)))
	  ((and (listp expr) (eq '^ (car expr)))
	   (cons 'expt
			 (mapcar #'prepare (cdr expr))))
	  ((listp expr)
	   (cons (car expr)
			 (mapcar #'prepare (cdr expr))))
	  (t expr)))
  (maxima-interface::lisp-to-maxima (prepare expr)))

(defun enmaxima-coords (coords)
  (cons '(maxima::mlist) (mapcar #'enmaxima-symbol coords)))

(defun hamilton-eqns (lagrangian coords)
  (maxima-interface:maxima-eval
   `((maxima::$ham)
	 ,(enmaxima-lagrangian lagrangian)
	 ,(enmaxima-coords coords))))

(defun unmaxima-expression (expr)
  (defun unmaxima-symbol (sym)
	(let ((name (symbol-name sym)))
	  (if (and (eq (symbol-package sym) (find-package :maxima))
			   (> (length name) 1)
			   (char= (char name 0) #\$))
		  (intern (subseq name 1)) sym)))
  (cond
	((symbolp expr) (unmaxima-symbol expr))
	((listp expr) (mapcar #'unmaxima-expression expr))
	(t expr)))

(defun momentum-symbol (q)
  (intern (format nil "P_~a" (symbol-name q))))

(defun next-symbol (s)
  (intern (format nil "NEXT-~a" (symbol-name s))))

(defun update-logic (result coords)
  (defun coord-update (q exprs)
	`(,(next-symbol q)
	  (+ ,q (* dt ,(unmaxima-expression (cadr exprs))))))
  (defun momentum-update (q exprs)
	(let ((p (momentum-symbol q)))
	  `(,(next-symbol p)
		(+ ,p (* dt ,(unmaxima-expression (caddr exprs)))))))
  (loop for exprs in (cdr (maxima-interface::maxima-to-lisp result))
        for q in coords
		collect (coord-update q exprs)
		collect (momentum-update q exprs)))

(defun keys (plist)
  (loop for key in plist by #'cddr collect key))

(defun update-body (result coords params)
  (let ((vars (append coords (mapcar #'momentum-symbol coords))))
    `(let (,@(mapcar (lambda (s) `(,s (getf params ',s))) (keys params))
           ,@(mapcar (lambda (s) `(,s (getf state ',s))) vars))
	   (let ,(update-logic result coords)
		 (list ,@(loop for s in vars
					   collect `',s
					   collect (next-symbol s)))))))