;;
;;  This module implements signal diagram combinators for differential
;;  and differential-algebraic equations.
;;
;; Copyright 2010-2012 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 signal-diagram-dynamics

	(
	 (ASSIGN make-assign-system) 
	 (ODE make-ode-system) 
	 (DAE make-dae-system) 
	 make-assign-system 
	 make-ode-system 
	 make-dae-system 
	 )

	(import scheme chicken)

	(require-extension extras data-structures srfi-1 signal-diagram )


(define (make-union rhs-list)
  (let ((n (length rhs-list)))
    (cond ((= n 1)  (car rhs-list))
	  ((= n 2)  (UNION (car rhs-list) (cadr rhs-list)))
	  (else     (UNION (UNION (car rhs-list) (cadr rhs-list)) 
			   (make-union (cddr rhs-list)))))))


(define (make-relation relation-list sf)
  (if (null? relation-list) sf
      (RELATION (car relation-list) 
		(make-relation (cdr relation-list) sf))))

  
(define (make-dae-system h indep eqs)

  (define (rewrite-relations expr aqs)

    (cond ((pair? expr)

	   (case (car expr)

	     ((if) (let ((es (cdr expr))) 
		     (map (lambda (x) (rewrite-relations x aqs)) es)))
 
	     ((let)
	      (let ((lbnds (cadr expr))
		    (body  (caddr expr)))
		(let ((lbnds1  (map (lambda (x) (list (car x) (rewrite-relations (cadr x) aqs))) lbnds))
		      (body1   (rewrite-relations body aqs)))
		  (list 'let lbnds1 body1))))

	     (else
	      (let ((s (car expr)) (es (cdr expr)))
		(cond ((and (symbol? s) (assoc s aqs)) =>
		       (lambda (x) (cons s (append es (drop (function-formals (cadr x)) (length es))))))
		      (else (cons s (map (lambda (x) (rewrite-relations x aqs)) es))))))

	     ))

	  (else  expr)))

  (let ((dqs (filter-map (lambda (x) (cond ((= 2 (length x)) (car x))
					    (else #f)))
			  eqs))

	(aqs (filter-map (lambda (x) (cond ((= 3 (length x)) (car x))
					   (else #f)))
			 eqs))

	(ads (filter-map (lambda (x) (cond ((= 3 (length x)) (cadr x))
					   (else #f)))
			 eqs)))

    (let* ((afs (filter-map
		 (lambda (eq)
		   (let ((rhs  (cond ((= 3 (length eq)) (caddr eq))
				     (else #f)))
			 (args (cond ((= 3 (length eq)) (cadr eq))
				     (else #f))))
		     (and rhs args
			  (let ((vars (enum-freevars rhs (append args symbolic-constants) '())))
			    (make-function (delete-duplicates (append args vars)) rhs)))))
		 eqs))

	   (afs (filter-map
		 (lambda (eq)
		   (let ((rhs  (cond ((= 3 (length eq)) (caddr eq))
				     (else #f)))
			 (args (cond ((= 3 (length eq)) (cadr eq))
				     (else #f))))
		     (let ((rhs (and rhs (rewrite-relations rhs (zip aqs afs)))))
		       (and rhs args
			    (let ((vars (enum-freevars rhs (append args symbolic-constants) '())))
			      (make-function (delete-duplicates (append args vars)) rhs))))))
		 eqs))
	   
	   (dfs (filter-map
		 (lambda (eq)
		   (let* ((rhs0 (cond ((= 2 (length eq)) (cadr eq))
				      (else #f)))
			  (rhs1 (and rhs0 (rewrite-relations rhs0 (zip aqs afs)))))
		    (and rhs1
			 (let ((vars (delete-duplicates (enum-freevars rhs1 symbolic-constants '()))))
			   (make-function vars rhs1)))))
		 eqs)))
      
      (let ((du (make-union (map (lambda (f d) (ACTUATE (list d) (INTEGRALH indep d h f))) dfs dqs))))

	(make-relation (zip aqs ads afs) (make-union (list du (make-assign-system `((,indep (+ ,indep ,h)))))))

      ))
    ))


(define-syntax DAE
  (syntax-rules ()
    [(_ h indep eqn ...)
     (make-dae-system (quote h) (quote indep) (quote (eqn ...)))]
    ))

  
(define (make-ode-system h indep eqs)
  (let ((deps (map car eqs))
	(rhss (map cadr eqs)))
    (let ((fs (map
	       (lambda (rhs) 		       
		 (let ((vars (delete-duplicates (enum-freevars rhs symbolic-constants '()))))
		   (make-function vars rhs)))
	       rhss)))

      (let ((u (cond ((or (null? fs) (null? deps)) (error 'make-ode-system "empty list of equations"))
		     ((null? (cdr fs))
		      (let ((d (car deps))) 
			(ACTUATE (list d) (INTEGRALH indep d h (car fs)))
			))
		     (else
		      (make-union (map (lambda (f d) (ACTUATE (list d) (INTEGRALH indep d h f))) fs deps)))
		     )))

	(make-union (list u (make-assign-system `((,indep (+ ,indep ,h))))))

	))))


(define-syntax ODE
  (syntax-rules ()
    [(_ h indep (dep rhs) ...)
     (make-ode-system (quote h) (quote indep) (quote ((dep rhs) ...)))]
    ))


(define (make-assign-system eqs)
  (let ((vars (map car eqs))
	(rhss (map cadr eqs)))
    (let ((fs (map
	       (lambda (x rhs) 		       
		 (let ((vars (delete-duplicates (enum-freevars rhs symbolic-constants '()))))
		   (ACTUATE (list x)
			    (if (pair? vars)
				(SENSE vars (PURE (make-function vars rhs)))
				(PURE (make-function vars rhs))))))
	       vars rhss)))
      (make-union fs))))


(define-syntax ASSIGN
  (syntax-rules ()
    [(_ eqn ...)
     (make-assign-system (quote (eqn ...)))]
    ))

)