;;
;;  This module implements combinators that are used to build 
;;  algebraic system functions out of pure functions.
;;
;;  Based on ideas from:
;;  1) Opis - Reliable Distributed Systems in OCaml
;;  (Copyright (C) 2008-2009 Pierre-Evariste DAGAND)
;;
;;  2) Yampa: Functional Reactive Programming with Arrows
;;  Developed by the Yale Haskell Group.
;;
;; Copyright 2013-2014 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 algebraic-system 

	(PRIM EQUATION RELATION UNION

	 function? make-function function-formals function-body
	 prim? make-prim prim-states prim-formals prim-body prim-init

	 signal? signal-name signal-value

         enum-freevars
         
	 construct dataflow codegen/Octave codegen/scheme codegen/ML
	 )

	(import scheme chicken)

	(require-extension extras data-structures srfi-1 datatype flsim)
	(require-library lolevel srfi-13 signal-diagram)
	(import (only srfi-13 string-concatenate string<)
		(only lolevel extended-procedure? procedure-data extend-procedure )
                (only signal-diagram prelude/ML prelude/scheme prelude/Octave)
		)

(include "expr-utils")

(define nl "\n")
(define (s+ . rst) (string-concatenate (map ->string rst)))

;; based on SRV:send-reply by Oleg Kiselyov
(define (print-fragments b)
  (let loop ((fragments b) (result #f))
    (cond
      ((null? fragments) result)
      ((not (car fragments)) (loop (cdr fragments) result))
      ((null? (car fragments)) (loop (cdr fragments) result))
      ((eq? #t (car fragments)) (loop (cdr fragments) #t))
      ((pair? (car fragments))
        (loop (cdr fragments) (loop (car fragments) result)))
      ((procedure? (car fragments))
        ((car fragments))
        (loop (cdr fragments) #t))
      (else
       (display (car fragments))
       (loop (cdr fragments) #t)))))


(define (symbol-pair? x)
  (and (pair? x) 
       (and (symbol? (car x)) 
	    (or (null? (cdr x)) 
		(symbol-pair? (cdr x))))))


(define (symbol-list? x)
  (and (list? x) (every symbol? x)))



(define make-signal cons)
(define signal-name car)
(define signal-value cdr)
(define signal? pair?)


(define-values (cgenenv-empty cgenenv-add cgenenv-find cgenenv-union )
  (letrec (
	   (empty      '())
	   (add        (lambda (s v env) 
			 (if (and (symbol? s) (symbol? v))
			     (cons (cons s v) env)
			     (error 'cgenenv-add "invalid arguments to add" s v))))
	   (find       (lambda (loc s env)
			 (let ((v (alist-ref s env)))
			   (if (not v) (error loc "label not found" s))
			   v)))
	   (union      (lambda (x y) (lset-union (lambda (x y) (eq? (first x) (first y))) x y)))
	   )
    (values empty add find union )))


(define (list->cgenenv loc slst source-cgenenv)
  (fold (lambda (s env) (cgenenv-add s (cgenenv-find loc s source-cgenenv) env))
	cgenenv-empty slst))




;;
;;   An arrow is an object with:
;;
;;   * dfe method, which produces dataflow information
;;   * codegen method, which generates pseudo-imperative code
;;

(define-record-type larrow
  (make-larrow  dfe codegen sig children relations)
  larrow?
  (dfe         larrow-dfe)
  (codegen     larrow-codegen)
  (sig         larrow-sig)
  (children    larrow-children)
  (relations   larrow-relations)
  )


(define-record-type dfe
  (make-dfe gen kill in out)
  dfe?
  (gen       dfe-gen )
  (kill      dfe-kill )
  (in        dfe-in )
  (out       dfe-out )
  )


(define-record-type function
  (make-function formals body)
  function?
  (formals function-formals)
  (body    function-body))


(define-record-type prim
  (make-prim states formals outputs body init-outputs init)
  prim?
  (states  prim-states)
  (formals prim-formals)
  (outputs prim-outputs)
  (body    prim-body)
  (init    prim-init)
  (init-outputs prim-init-outputs)
  )


(define (function-list? x)
  (and (list? x) (every function? x)))


(define (relation? r)
  (and (pair? r) (symbol? (car r)) 
       (symbol-list? (cadr r))
       (function? (caddr r))))


(define-datatype alsys alsys?
  (PRIM         (f prim?) (name symbol?))
  (RELATION     (r relation?) (f alsys?))
  (UNION        (f alsys?) (g alsys?))
  (EQUATION     (s symbol?) (f function?))
  )

(define (select-signal loc s env)
  (let ((v (cgenenv-find loc s env)))
    (if (eq? s v) (V:Var s) (V:Sel s (V:Var v)))))


(define-record-type codegen
  (make-codegen0 rv renv expr)
  codegen?
  (rv          codegen-rv)
  (renv        codegen-renv)
  (expr        codegen-expr)
  )

(define (make-codegen rv renv expr)
  (if (not (symbol? rv)) (error 'make-codegen "invalid return variable"))
  (make-codegen0 rv renv expr))


(define codegen-state (make-parameter '()))
    
					   
;;
;;  The arrow combinators are used to build signal functions upon pure
;;  functions.
;;  

;;  [ls f] encapsulates a pure function into a signal function.


(define (ls f name)
   (let* (
          (fd (and (extended-procedure? f) (procedure-data f)))
          (formals (or (and (prim? fd) (prim-outputs fd)) 
                       (and (function? fd) (function-formals fd))
                       '()))
          (outputs (or (and (prim? fd) (prim-outputs fd)) '()))
          (states (or (and (prim? fd) (prim-states fd)) '()))
          )

     (make-larrow 

      ;; dataflow equations
      (make-dfe
       ;; gen
       (lambda (s) (if (prim? fd) outputs (list name)))
       ;; kill
       (lambda (s) (if (prim? fd) outputs (list name)))
       ;; in
       (lambda (s) formals)
       ;; out
       (lambda (s) (if (prim? fd) outputs (list name))))

      ;; codegen
      (lambda (s env dfe) 
	(let ((in   ((dfe-in dfe) s))
              (out  ((dfe-out dfe) s))
	      (rv1  (gensym 'rv))
	      (rv2  (gensym 'rv))
              (init-name (and (prim? fd) (gensym (string->symbol (string-append (->string name) "init")))))
              )

	  (make-codegen
	   rv2
           (fold (lambda (name env) (cgenenv-add name rv2 env)) cgenenv-empty out)
	   (append
            
	    (cond ((function? fd)
                   (list (function->expr name fd)))
                  ((prim? fd)
                   (list (prim->expr name fd) ))
                  (else '()))

	    (cond ((function? fd)
                   (if (null? (function-formals fd))
                       (list (B:Val rv2 (V:Rec `((,name ,(V:Var name))))))
                       (list (B:Val rv1 (V:Op name (map (lambda (s) (select-signal 'ls s env)) in)))
                             (B:Val rv2 (V:Rec `((,name ,(V:Var rv1))))))))

                  ((prim? fd)

                     (codegen-state
                      (append
                       (list (prim->init init-name fd))
                       (codegen-state)))

                     (list (B:Val rv1 (V:Op name (append
                                                  (map (lambda (s) (select-signal 'ls s env)) in)
                                                  (map (lambda (x) (V:Sel x (V:Var init-name))) 
                                                       (lset-difference eq? states in)))))
                           (B:Val rv2 (V:Rec (map (lambda (s) `(,s ,(V:Sel s (V:Var rv1)))) outputs)))
                           )
                     )
                  
                  (else '())
                  )
	    ))
          ))
      ;; signature
      `(EQUATION ,name ,states ,outputs)
      ;; children
      `(EQUATION)
      ;; relations
      `()
      ))
     )
  



(define (ls-equation name f)
  (let* ((f0   (cond ((function? f)  (lambda () `(,name ,(function-formals f) ,(function-body f))))
                     ((procedure? f) f)
                     (else (error 'ls-equation "invalid function" f))))
         (f1 (if (function? f) (extend-procedure f0 f) f0)))
    (ls f1 name)))


(define (ls-prim f name)
  (let* ((f0   (cond ((prim? f)  (lambda () `(,name ,(append (prim-formals f) (prim-states f)) ,(prim-body f))))
                     (else (error 'ls-prim "invalid primitive" f))))
         (f1 (if (prim? f) (extend-procedure f0 f) f0)))
    (ls f1 name)))


(define (ls-relation r ls)

  (define (relation-vars r) (function-formals (caddr r)))

  (define (relations-inputs ls)
    (let recur ((ls ls) (inputs '()))
      (let ((inputs (append
		     (concatenate (map relation-vars (larrow-relations ls)))
		     inputs)))
	(let ((ls-children (filter-map larrow? (larrow-sig ls))))
	  (if (null? ls-children) inputs
	      (fold recur inputs ls-children)
	      )))
      ))

  (if (relation? r)
      (let* ((dfe (larrow-dfe ls))
	     (dfe1 (make-dfe (dfe-gen dfe) (dfe-kill dfe)
                             (lambda (s) (delete-duplicates
                                          (lset-difference eq? 
                                              (append ((dfe-in dfe) s) 
                                                      (relations-inputs ls))
                                              (relation-vars r))))
			     (dfe-out dfe))))
	(make-larrow dfe1 
                      (larrow-codegen ls) 
		      (larrow-sig ls) (larrow-children ls)
		      (cons r (larrow-relations ls))))
      (error 'ls-relation "invalid relation" r)))


(define (relations-codegen ls env)

  (let ((kons (map (lambda (x) (car x)) (larrow-relations ls))))

    (codegen-state
     (append (codegen-state)
	     (reverse 
	      (map
	       (lambda (r k)
		 (let ((name (car r)) 
		       (fd (caddr r)))
		   (function->expr k fd)))
	       (larrow-relations ls) kons))
	     ))
    '()
    ))






;; [union f g], applies [f] and [g] to the input signal in parallel.

(define (ls-union f g)

  (define (flatten-union u)
    (let ((uc (larrow-children u)))
      (case (car uc)
	((UNION)  (append (flatten-union (cadr uc))
			  (flatten-union (caddr uc))))
	(else     (list u)))))
    

   (let* ((fe      (larrow-dfe f))
	  (ge      (larrow-dfe g))

	  (fe-in   (dfe-in fe))	  
	  (fe-out  (compose (dfe-out fe) fe-in))
	  (fe-gen  (compose (dfe-gen fe) fe-in))
	  (fe-kill (compose (dfe-kill fe) fe-in))

	  (ge-in   (dfe-in ge))
	  (ge-out  (compose (dfe-out ge) ge-in))
          (ge-gen  (compose (dfe-gen ge) ge-in))
          (ge-kill (compose (dfe-gen ge) ge-in))
          
          (flst (flatten-union f))
          (glst (flatten-union g))
	  )

     (make-larrow

      ;; dataflow equations
      (make-dfe
       ;; gen
       (lambda (s) (lset-union eq? (ge-gen s) (fe-gen s)))
       
       ;; kill
       (lambda (s) (lset-union eq? (fe-kill s) (ge-kill s)))
       
       ;; in
       (lambda (s) (lset-union eq? (ge-in s) (fe-in s)))
       
       ;; out
       (lambda (s) (lset-union eq? (ge-out s) (fe-out s)))
       
       )

      ;; codegen
      (lambda (s env dfe) 
	(let* (
	       (fgx      (lset-intersection eq? (fe-gen s) (ge-gen s)))
	       
	       (codegen (lambda (ls)
			  (let ((codegen (larrow-codegen ls))
				(dfe (larrow-dfe ls)))
			    (let ((env (list->cgenenv 'union1 ((dfe-in dfe) s) env)))
			      (codegen ((dfe-in dfe) s) env dfe)))))

	       (fld  (lambda (codegen dfe)
		       (let ((renv (codegen-renv codegen)))
			 (map (lambda (x) (list x (select-signal 'union2 x renv)))
			      ((dfe-out dfe) s)))))
	       
	       )

	  (if (not (null? fgx)) (error 'ls-union "union arguments output overlapping signals" fgx))

	  (let ((rv (gensym 'union))
                (fcodegen-lst (map codegen flst))
                (gcodegen-lst (map codegen glst))
                )

	      (let* ((renv-lst (map codegen-renv (append fcodegen-lst gcodegen-lst)))
		     (expr-lst (map codegen-expr (append fcodegen-lst gcodegen-lst)))
		     (renv (list->cgenenv 'union3 ((dfe-out dfe) s)
				(let recur ((renv-lst renv-lst) (env '()))
				  (if (null? renv-lst) env
				      (recur (cdr renv-lst) (cgenenv-union (car renv-lst) env)))))))

		(make-codegen 

		 rv
		 
		 (fold (lambda (name env) (cgenenv-add name rv env)) cgenenv-empty (map car renv))
		 
		 (let ((fflds  (map fld fcodegen-lst (map larrow-dfe flst)))
		       (gflds  (map fld gcodegen-lst (map larrow-dfe glst))))
		   (append 
		    (concatenate (map (lambda (f) (relations-codegen f env)) (append flst glst)))
		    (concatenate expr-lst)
		    (list (B:Val rv (V:Rec (map (lambda (s) `(,s ,(select-signal 'union4 s renv))) (map car renv)))))
		    ))
	       
	       )))
	    ))
      ;; signature
      `(UNION ,(larrow-sig f) ,(larrow-sig g))
      ;; children
      `(UNION  ,f ,g)
      ;; relations
      (append (larrow-relations f) (larrow-relations g))
      ))
)



(define (construct l)
  (cases alsys l
	 (EQUATION (x f)           (ls-equation x f))
	 (PRIM (f name)            (ls-prim f name))
	 (RELATION (r f)           (ls-relation r (construct f)))
	 (UNION (f g)              (ls-union (construct f) (construct g)))
	 ))


(define (dataflow f input)
  (let ((dfe (larrow-dfe f)))
    `((gen  . ,((dfe-gen dfe) input))
      (kill . ,((dfe-kill dfe) input))
      (in   . ,((dfe-in dfe) input))
      (out  . ,((dfe-out dfe) input)))))


(define (codegen/Octave name f #!key (initial #f) (pre #t))

  (let ((dfe (larrow-dfe f)))

    (codegen-state '())

    (let* ((input    (or (and initial (map car initial)) ((dfe-in dfe) '())))
	   (fenv     (map (lambda (s) (cons s 'input)) input))
	   (fcodegen ((larrow-codegen f) input fenv dfe ))
	   (relations-expr (relations-codegen f input))
	   (globals   (filter-map
		       (lambda (x) 
			 (cases binding x
				(B:Val  (name v) (name/Octave name))
				(else  #f)))
		       (codegen-state))))

      (if pre (print-fragments (prelude/Octave)))
      (print-fragments (list "global " (intersperse globals " ") ";" nl))

      (print-fragments (list (map (lambda (x) 
                                    (if (binding? x)
                                        (cases binding x
                                               (B:Val  (name v) (list (name/Octave name) " = " (value->Octave v) ";" nl)))
                                        (expr->Octave x)))
				  (reverse (codegen-state))) nl))

      (print-fragments
       (list
	"function " (name/Octave (codegen-rv fcodegen)) " = " (name/Octave name) " (input)" nl
	(list "global " (intersperse globals " ") ";" nl)
	(map (lambda (x) (list (name/Octave x) " = " (value->Octave (V:Sel x (V:Var 'input))) "; " nl)) input)
	nl
	(map binding->Octave (append relations-expr (codegen-expr fcodegen))) nl 
	"endfunction" nl))

      
      (if initial
	      (print-fragments
	       (list (name/Octave name) "_initial = " 
		     (value->Octave
		      (V:Rec (map (lambda (x) (let ((n (car x)) (v (cadr x)))
						(let ((v (cond ((and (or (number? v) (symbol? v))) v)
							       ((boolean? v) (if v 'true 'false))
							       (else v))))
						  (list n (V:C v)))))
				  initial)))
                     nl))
	      )


      ))
  )


(define (codegen/scheme name f #!key (initial #f) (pre #t))

  (let ((dfe (larrow-dfe f)))

    (codegen-state '())

    (let* ((input    (or (and initial (map car initial)) ((dfe-in dfe) '())))
	   (fenv     (map (lambda (s) (cons s 'input)) input))
	   (fcodegen ((larrow-codegen f) input fenv dfe ))
	   (relations-expr (relations-codegen f input)))

      (if pre (print-fragments (prelude/scheme)))

      (print-fragments (list (map (lambda (x) 
                                    (if (binding? x)
                                        (cases binding x
                                               (B:Val  (name v)
                                                       (list "(define " (name/scheme name) " " (value->scheme v) ")" nl)))
                                        (expr->scheme x)))
				  (reverse (codegen-state))) nl))

      (print-fragments
       (list
	"(define (" (name/scheme name) " input)" nl
	"(let (" (intersperse (map (lambda (x) (binding->scheme (B:Val x (V:Sel x (V:Var 'input))))) input) " ")  ")"  nl
	"(let* (" (map binding->scheme (append relations-expr (codegen-expr fcodegen))) nl ")" nl
	(codegen-rv fcodegen) nl
	")))" nl))

      (if initial
	  (print-fragments
	   (list "(define " (name/scheme name) "_initial " 
		 (value->scheme
		  (V:Rec (map (lambda (x) (let ((n (car x)) (v (cadr x)))
					    (let ((v (cond ((and (or (number? v) (symbol? v))) v)
							   ((boolean? v) (if v 'true 'false))
							   (else v))))
					      (list n (V:C v)))))
			      initial))) ")" nl)))


      )))

(define (codegen/ML name f #!key (initial #f) (pre #t) (post #t) (random #f))

  (let ((dfe (larrow-dfe f)))

    (codegen-state '())

    (let* ((input    (or (and initial (lset-intersection eq? (map car initial) ((dfe-in dfe) '())))
			 ((dfe-in dfe) '())))
	   (fenv     (map (lambda (s) (cons s 'input)) input))
	   (fcodegen ((larrow-codegen f) input fenv dfe ))
	   (relations-expr (relations-codegen f input)))


      (if pre (print-fragments (prelude/ML random: random)))

      (print-fragments (list (map (lambda (x) 
                                    (if (binding? x)
                                        (cases binding x
                                               (B:Val  (name v)
                                                       (list "val " (name/ML name) " = " (value->ML v) nl)))
                                        (expr->ML x)))
				  (reverse (codegen-state))) nl))


      (print-fragments
       (list
	"fun " (name/ML name) "(input as {" (intersperse (map name/ML input) ",")  "})" " = " nl
	"let" nl
	(map binding->ML (append relations-expr (codegen-expr fcodegen))) nl
	"in" nl
	(codegen-rv fcodegen)	nl
	"end" nl))

      (if initial
	  (print-fragments
	   (list "val " (name/ML name) "_initial = " 
		 (value->ML (V:Rec (map (lambda (x) 
					  (let ((n x) (v (car (alist-ref x initial))))
					    (list n 
						  (cond ((and (or (number? v) (symbol? v))) 
							 (V:C v))
							((boolean? v) 
							 (V:C (if v 'true 'false)))
							(else (V:C v))))
						    ))
					input))) nl)))

      (if post (print-fragments (list "end" nl)))
      
      )))




)