;;
;; Definition and code generators for a simple applicative language
;; for numerical simulation.
;;
;; Copyright 2010-2011 Ivan Raikov and the Okinawa Institute of
;; Science and Technology.
;;
;; This program is free software: you can redistribute it and/or
;; modify it under the terms of the GNU General Public License as
;; published by the Free Software Foundation, either version 3 of the
;; License, or (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful, but
;; WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;; General Public License for more details.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.
;;

(module flsim 

	(
	 value?
	 V:C V:Var V:Rec V:Sel V:Vec V:Sub V:Fn V:Op V:Ifv

	 expr?
	 E:Val E:Ife E:Let E:Ret E:Noop

	 name/scheme prelude/scheme expr->scheme value->scheme
	 name/ML prelude/ML expr->ML value->ML
	 name/Octave prelude/Octave expr->Octave value->Octave
	)

	(import scheme chicken)
	
	(require-extension extras data-structures srfi-1 datatype)

(define nl "\n")


(define-datatype value value?
  (V:C       (v (lambda (x) (or (symbol? x) (number? x) ))))
  (V:Var     (name symbol?))
  (V:Rec     (flds (lambda (xs) (and (pair? xs) (every (lambda (x) (and (symbol? (car x)) (value? (cadr x)))) xs)))))
  (V:Sel     (field (lambda (x) (or (symbol? x) (integer? x)))) (v value?))
  (V:Vec     (vals (lambda (x) (every value? x))))
  (V:Sub     (index (lambda (x) (and (integer? x) (or (zero? x) (positive? x))))) (v value?))
  (V:Fn      (args pair?) (body expr?))
  (V:Op      (name symbol?)  (args (lambda (x) (every value? x))))
  (V:Ifv     (test value?)  (ift value?) (iff value?))
  )

  
(define-record-printer (value x out)
  (fprintf out "~A"
	   (cases value x
		  (V:C  (v)    (sprintf "V:C ~A" v))
		  (V:Var (n)   (sprintf "V:Var ~A " n))
		  (V:Rec (lst) (sprintf "V:Rec ~A" lst))
		  (V:Sel (f v)  (sprintf "V:Sel ~A ~A" f v))
		  (V:Vec (lst)  (sprintf "V:Vec ~A" lst))
		  (V:Sub (i v)  (sprintf "V:Sub ~A ~A" i v))
		  (V:Ifv (test tv fv)   "V:Ifv ~A ~A ~A" test tv fv)
		  (V:Fn  (args body) (sprintf "V:Fn ~A = ~A" args body))
		  (V:Op (name args) (sprintf "(V:Op ~A ~A)" name args)))))
		  

(define-datatype expr expr?
  (E:Val     (name symbol?)  (v value?))
  (E:Ife     (test value?)   (ift expr?) (iff expr?))
  (E:Let     (bnds pair?)    (body expr?))
  (E:Ret     (v value?))
  (E:Noop)
  )

  
(define-record-printer (expr x out)
  (fprintf out "~A"
	   (cases expr x
		  (E:Val (name v)       (sprintf "E:Val ~A = ~A" name v))
		  (E:Ife (test ift iff) (sprintf "E:Ife ~A ~A ~A" test ift iff))
		  (E:Let (bnds body) (sprintf "E:Let ( ~A ) ~A" bnds body))
		  (E:Ret (v)         (sprintf "E:Ret ~A" v))
		  (E:Noop ()         (sprintf "E:Noop"))
		  )))



(define (name/scheme s)
  (let ((cs (string->list (->string s))))
    (let loop ((lst (list)) (cs cs))
      (if (null? cs) (string->symbol (list->string (reverse lst)))
	  (let* ((c (car cs))
		 (c1 (cond ((or (char-alphabetic? c) (char-numeric? c) 
				(char=? c #\_) (char=? c #\-)) c)
			   (else #\-))))
	    (loop (cons c1 lst) (cdr cs)))))))


(define (expr->scheme x . rest)
  (let-optionals rest ((bnd? #t))
    (cases expr x
	 (E:Val     (name v)
		    (list "(" (name/scheme name) " " (value->scheme v) ")" nl))

	 (E:Ife     (test ift iff)
		    (list "(cond (" (value->scheme test) " " nl
			  " " (expr->scheme ift ) ")" nl
			  "(else " (expr->scheme iff) "))" nl))

	 (E:Let     (bnds body)
		    (list "(let* (" nl
			  (map (lambda (x) (expr->scheme x #t)) bnds) nl
			  ") " nl
			  (expr->scheme body #f) nl
			  ")" nl))
			  
	 (E:Ret     (v)  (value->scheme v))
		    
	 (E:Noop    () (list "(void)"))
	 )))


(define (value->scheme v)
  (let ((result
  (cases value v
	 (V:C       (v) v)
	 (V:Var     (name) (name/scheme name))
	 (V:Rec     (lst) 
		    (list "`(" (intersperse (map (lambda (nv) (list "(" (name/scheme (car nv)) " . ,"  
								    (value->scheme (cadr nv)) ")")) lst) " ") ")"))
	 (V:Sel     (field v) 
		    (if (number? field)
			(list "(list-ref " (value->scheme v) " " (- field 1) ")")
			(list "(alist-ref '" (name/scheme field) " " (value->scheme v) ")")))
	 (V:Vec     (lst) 
		    (list "(list " (intersperse (map value->scheme lst) " ") ")"))
	 (V:Sub     (index v) 
		    (list "(list-ref " (value->scheme v) " " index ")"))
	 (V:Fn      (args body) 
		    (list "(lambda (" (intersperse (map name/scheme args) " ") ") "
			  (expr->scheme body #f) ")"))
	 (V:Op     (name args)
		    (let* ((fp? (case name
				     ((+ - * / >= > < <= neg)  #t)
				     (else #f)))
			   (op (if fp? (conc "fp" name) name)))
		      (cond ((null? args) 
			     (case name 
			       ((NONE)  (list "#f"))
			       (else    (list "(" name ")"))))
			    
			    (fp?
			     (if (pair? (cdr args))
				 (fold-right (lambda (x ax) (list "(" op " " (value->scheme x) " " ax ")"))
					     (list "(" op " " (value->scheme (car args)) " " (value->scheme (cadr args)) ")")
					     (cddr args))
				 (list "(" op " " (value->scheme (car args)) ")")
				 ))

			    (else
			     (list "(" op " " (intersperse (map value->scheme args) " ") ")")))))
	 (V:Ifv     (test ift iff)
		    (list "(if " (value->scheme test) " " 
			  (value->scheme ift) " " 
			  (value->scheme iff) ")"))

	 )))
    result))
  
(define (prelude/scheme #!key (solver 'rk4b))

`(
 #<<EOF
(use runge-kutta mathh)

;; Variant types 

(define-syntax define-datatype
  (syntax-rules ()
    [(_ type (name field ...) ...)
     (begin
       (define-constructors type ((name field ...) ...)))]))


(define-syntax define-constructors
  (syntax-rules ()
    [(define-constructors type ((name field ...) ...))
     (define-constructors type ((name field ...) ...) (name ...))]
    [(define-constructors type ((name field ...) ...) names)
     (begin
       (define-constructor type (name field ...) names)
       ...)]))


(define-syntax define-constructor
  (syntax-rules ()
    [(_ type (name field ...) names)
     (define (name field ...)
       (cons 'type
             (lambda names
               (name field ...))))]))


(define-syntax cases
  (syntax-rules ()
    [(_ type x [(name field ...) exp]
          ...)
     ((cdr x) (lambda (field ...) exp)
              ...)]))

(define-datatype trs (TRSA a) (TRSB b))
(define-datatype trc (TRC f fk e))

(define (tsCase fa fb x) (cases trs x ((TRSA a) (fa a)) ((TRSB b) (fb b))))
(define (trfOf x) (cases trc x ((TRC f fk e) f)))
(define (trfkOf x) (cases trc x ((TRC f fk e) fk)))
(define (treOf x) (cases trc x ((TRC f fk e) e)))

(define-datatype option (NONE) (SOME a))
(define (swap x v) (cases option v ((NONE) x) ((SOME v) v)))

(define false #f)
(define true  #t)

(define equal equal?)

(define (signalOf v) (if (not v) (error 'signalOf "empty signal" v) v))


EOF
,(if (not solver)
     `((";; dummy solver; returns only the computed derivatives")
       ("(define (integrate1d f x y h) (f x y))" ,nl))
     (case solver
       ((cvode) 
	`(("(use sundials lolevel)" ,nl)
	  ( #<<EOF
(define (integrate1d f x y h) 
  (let ((solver
	 (or (procedure-data f)
	     (let ((solver (cvode-create-solver
			    x 0.0 (list->f64vector y) 
			    (lambda (x y) (list->f64vector (f x (f64vector->list y))))
			    abstol: 1e-1
			    reltol: 1e-2
			    lmm: cvode-lmm/bdf)))
	       (extend-procedure f solver)
	       solver))))
    (cvode-solve solver (+ x h))
    (let ((yy (cvode-yy solver)))
      (list (f64vector-ref yy 0))
      )))
EOF
)))
       (else
	`(("(define (scaler x a) (map (lambda (k) (fp* x k)) a))" ,nl)
	  ("(define (summer a b) (map fp+ a b))" ,nl)
	  ("(define " ,solver " (make-" ,solver "))" ,nl)
	  ("(define (make_stepper deriv) (" ,solver " scaler summer deriv))" ,nl)
	  ("(define (integrate1d f x y h) (((make_stepper f) h) x y))" ,nl)))
       ))
))


(define (name/ML s)
  (let ((cs (string->list (->string s))))
    (let loop ((lst (list)) (cs cs))
      (cond ((null? cs) (string->symbol (list->string (reverse lst))))
	    ((null? (cdr cs))
	     (let ((c (car cs)))
	       (if (or (char-alphabetic? c) (char-numeric? c))
		   (loop (cons c lst) (cdr cs))
		   (loop (append (reverse (string->list (->string (gensym 't)))) lst) (cdr cs))
		   )))
	    (else
	     (let* ((c (car cs))
		    (c1 (cond ((or (char-alphabetic? c) (char-numeric? c) 
				   (char=? c #\_) (char=? c #\#)) c)
			      (else #\_))))
	       (loop (cons c1 lst) (cdr cs))))))))
			    

(define (expr->ML x . rest)
  (let-optionals rest ((bnd? #t))
    (cases expr x
	 (E:Val     (name v)
		    (list "val " (name/ML name) " = " (value->ML v) nl))

	 (E:Ife     (test ift iff)
		    (list "if (" (value->ML test) ") " nl
			  "then " (expr->ML ift ) nl
			  "else " (expr->ML iff) nl))

	 (E:Let     (bnds body)
		    (list "let " nl
			  (map (lambda (x) (expr->ML x #t)) bnds) nl
			  "in " nl
			  (expr->ML body #f) nl
			  "end" nl))
			  
	 (E:Ret     (v)  (value->ML v))
		    
	 (E:Noop    () (list "()"))
	 )))


(define (value->ML v)
  (cases value v
	 (V:C       (v) (cond ((and (number? v) (negative? v))
			       (list "~" (abs v)))
			      (else v)))
	 (V:Var     (name) (name/ML name))
	 (V:Rec     (lst) 
		    (list "{" (intersperse (map (lambda (nv) (list (name/ML (first nv)) " = " 
								   (value->ML (cadr nv)))) lst) ", ") "}"))
	 (V:Sel     (field v) 
		    (list "(#" (name/ML field) "(" (value->ML v) "))"))
	 (V:Vec     (lst) 
		    (let ((n (length lst)))
		      (list "([" (intersperse (map (lambda (v) (value->ML v)) lst) ", ") "])")))
	 (V:Sub     (index v) 
		    (list "List.nth (" (value->ML v) ", " index ")"))
	 (V:Fn      (args body) 
		    (list "(fn (" (intersperse (map name/ML args) ",") ") => "
			  (expr->ML body #f) ")"))
	 (V:Op    (name args)
		    (let* ((infix? (case name
				     ((+ - * / >= > < <=)  #t)
				     (else #f)))
			   (op (if infix? (list "(op " name ")") name)))
		      (cond ((null? args) 
			     (case name 
			       ((NONE)  (list name))
			       (else    (list name "()"))))
			    ((null? (cdr args))
			     (list "(" op " " (value->ML (car args)) ")"))
			    ((and infix? (pair? (cddr args)))
			     (list "(foldr " op "(" (value->ML (V:Op name (list (car args) (cadr args)))) ")"
				   "[" (intersperse (map value->ML (cddr args)) ",") "])"))
			    (else
			     (list "(" op "(" (intersperse (map value->ML args) ",") "))")))))
	 (V:Ifv     (test ift iff)
		    (list "(if (" (value->ML test) ") " 
			  "then " (value->ML ift ) " " 
			  "else " (value->ML iff) ")"))

	 ))


(define (prelude/ML  #!key (solver 'rk4b) )
`(
 #<<EOF
structure Model = 
struct

open Real
open Math
open RungeKutta

datatype ('b,'c) trs = TRSA of 'b | TRSB of 'c
datatype ('a,'b,'c) trc = TRC of ((('a -> (('b,'c) trs))) * 
                                  (('a -> (('b,'c) trs))) * 
                                  ((('b,'c) trs) -> bool))
	 
fun tsCase (fa,fb,x) = case x of TRSA a => (fa a) | TRSB b => (fb b)
fun trfOf x = case x of TRC (f,fk,e) => f
fun trfkOf x = case x of TRC (f,fk,e) => fk
fun treOf x = case x of TRC (f,fk,e) => e

fun putStrLn str = 
  (TextIO.output (TextIO.stdOut, str);
   TextIO.output (TextIO.stdOut, "\n"))

fun putStr str = (TextIO.output (TextIO.stdOut, str))

fun showReal n = 
let open StringCvt
in
(if n < 0.0 then "-" else "") ^ (fmt (FIX (SOME 12)) (abs n))
end

exception EmptySignal
val neg = (op ~)
val swap = fn (x,v) => (case v of NONE => x | SOME v => v) 
val equal = fn (x,y) => (x = y) 
val signalOf = fn (v) => (case v of NONE => raise EmptySignal | SOME v => v) 

EOF

,(if (not solver)
     `(("(* dummy solver; returns only the computed derivatives *)")
       ("fun integrate1d (f,x: real,y: real list,h) = (f (x,y))" ,nl))
     `(("val summer = fn (a,b) => (ListPair.map (fn (x,y) => x+y) (a,b))" ,nl)
       ("val scaler = fn(a,lst) => (map (fn (x) => a*x) lst)" ,nl)
       ("val " ,solver ": (real list) stepper1 = make_" ,solver "()" ,nl)
       ("fun make_stepper (deriv) = " ,solver " (scaler,summer,deriv)" ,nl)
       ("fun integrate1d (f,x: real,y: real list,h) = ((make_stepper f) h) (x,y)" ,nl)))

))


(define (name/Octave s)
  (let ((cs (string->list (->string s))))
    (let loop ((lst (list)) (cs cs))
      (cond ((null? cs) (string->symbol (list->string (reverse lst))))
	    ((null? (cdr cs))
	     (let ((c (car cs)))
	       (if (or (char-alphabetic? c) (char-numeric? c))
		   (loop (cons c lst) (cdr cs))
		   (loop (append (reverse (string->list (->string (gensym 't)))) lst) (cdr cs))
		   )))
	    (else
	     (let* ((c (car cs))
		    (c1 (cond ((or (char-alphabetic? c) (char-numeric? c) 
				   (char=? c #\_)) c)
			      (else #\_))))
	       (loop (cons c1 lst) (cdr cs))))))))


(define (expr->Octave x . rest)
  (let-optionals rest ((inner? #f))
    (cases expr x
	 (E:Val     (name v)
		    (if inner?
			(list (name/Octave name) " = " (value->Octave v) )
			(list (name/Octave name) " = " (value->Octave v) ";" nl)))

	 (E:Ife     (test ift iff)
		    (list "if (" (value->Octave test) ") " nl
			  (expr->Octave ift ) nl
			  " else " (expr->Octave iff) " endif" nl))

	 (E:Let     (bnds body)
		    (if inner?
			(list #\{ (intersperse (map (lambda (x) (list "(" (expr->Octave x #t) ")")) 
						    (append bnds (list body))) ", ") #\}
			      "{" (+ 1 (length bnds)) "}" nl)
			(list (map (lambda (x) (list (expr->Octave x) nl)) bnds) nl
			      (expr->Octave body) nl)))
			  
	 (E:Ret     (v)  (value->Octave v))
		    
	 (E:Noop    () (list))
	 )))


(define (value->Octave v)
  (cases value v
	 (V:C       (v) v)
	 (V:Var     (name) (name/Octave name))
	 (V:Rec     (lst) 
		    (list "struct (" (intersperse (map (lambda (nv) (list #\" (name/Octave (first nv)) #\" ", " 
									  (value->Octave (cadr nv)))) lst) ", ") ")"))
	 (V:Sel     (field v) 
		    (list (value->Octave v) "." (name/Octave field)))
	 (V:Vec     (lst) 
		    (let ((n (length lst)))
		      (list "([" (intersperse (map (lambda (v) (value->Octave v)) lst) ", ") "])")))
	 (V:Sub     (index v) 
		    (list (value->Octave v) "((" index ")+1" ")"))
	 (V:Fn      (args body) 
		    (list "(@ (" (intersperse (map name/Octave args) ",") ") "
			  (expr->Octave body #t) ")"))
	 (V:Op    (name args)
		    (let* ((infix? (case name
				     ((+ - * / >= > < <=)  #t)
				     (else #f)))
			   (op name))
		      (cond ((null? args) 
			     (case name 
			       ((NONE)  (list name))
			       (else    (list name "()"))))
			    ((null? (cdr args))
			     (list op "(" (value->Octave (car args)) ")"))
			    ((and infix? (null? (cddr args)))
			     (list "(" (value->Octave (car args)) ")" op "(" (value->Octave (cadr args)) ")"))
			    (infix?
			     (let ((op (case op ((+) 'sum) ((*) 'prod) (else op))))
			       (list  op "([" (intersperse (map value->Octave args) ",") "])")))
			    (else
			     (list  op "(" (intersperse (map value->Octave args) ",") ")")))))
	 (V:Ifv     (test ift iff)
		    (list "(ifv (" (value->Octave test) ", " 
			  (value->Octave ift ) ", " 
			  (value->Octave iff) "))"))

	 ))

  
(define (prelude/Octave #!key (solver 'lsode))
`(
#<<EOF

function res = TRSA(v)
  res = struct ("TRSA",v);  
end

function res = TRSB(v)
  res = struct ("TRSB",v);  
end

function res = tsCase(fa,fb,v)
   if (isfield (v, "TRSA"))
     res = fa(getfield(v,"TRSA"));
   else
     res = fb(getfield(v,"TRSB"));
   endif
end

function res = TRSC(f,fk,e)
  res = struct ("TRSC",[f,fk,e]);  
end

function res = trfOf(x)
   res = getfield(x,"TRC")(1);
end

function res = trfkOf(x)
   res = getfield(x,"TRC")(2);
end

function res = treOf(x)
   res = getfield(x,"TRC")(3);
end

function res = NONE()
  res = struct ();  
end

function res = SOME(v)
  res = struct ("SOME",v);  
end

function res = swap (x,v)
   if (isfield (v, "SOME"))
     res = getfield(v,"SOME");
   else
     res = x;
   endif
end

function res = equal (x, y)
   res = (x==y);
end

function res = signalOf (v)
  if (not(v))
   error ("empty signal")
 else
   res = v;
 endif
end

function res = neg (x)
  res = -x;
end

function res = ifv (b,x,y)
  if (b) 
    res = x;
  else
    res = y;
  endif
end

EOF
,(if (not solver)
     `("function res = integrate1d (f,x,y,h)" ,nl 
       "  res = f(x,y); " ,nl end ,nl)
     `((
#<<EOF

EOF
)
       ("global " ,solver #\; ,nl)
       (,solver = #\@ ,(case solver
			 ((rk3) 'ode23) ((rk4b rk4a) 'ode45) ((rkf45) 'ode54) 
			 (else 'lsode)) #\; ,nl)

       ,(cond

	 ((member solver '(rk3 rk4b rk4a rkf45))

	       `(
#<<EOF
global reltol abstol dt;
reltol = 0.1;
abstol = 0.01;
dt = 0.001;

global P;
P = odeset ("RelTol" , reltol , "AbsTol" , abstol , "MaxStep" , 1 , "InitialStep" , dt) ;

EOF
       "function res = integrate1d (f,x,y,h)" ,nl 
       "  global P " ,solver ";" ,nl
       "  [t,y] = " ,solver "(f,[x,x+h],y,P); " ,nl 
       "  res = y(2);" ,nl end ,nl
       ))
	      (else

	       `(
#<<EOF
lsode_options("relative tolerance",1e-1);
lsode_options("absolute tolerance",1e-2);


EOF
       "function res = integrate1d (f,x,y0,h)" ,nl 
       "  global " ,solver ";" ,nl
       "  y = " ,solver "(@(yvec,t) f(t,yvec),y0,[x,x+h]); " ,nl 
       "  res = y(2);" ,nl end ,nl
       )
   ))
 ))
))

	 

)