;;       
;; 
;; An extension for generating NEURON HOC voltage clamp scripts from NEMO models.
;;
;; Copyright 2008-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 nemo-vclamp

	(nemo:vclamp-translator)

	(import scheme chicken utils data-structures lolevel srfi-1 srfi-13 srfi-69)

	(require-extension lolevel datatype matchable strictly-pretty 
			   varsubst datatype 
			   nemo-core nemo-utils nemo-gate-complex)

(define-syntax pp
  (syntax-rules ()
    ((pp indent val ...) (ppf indent (quasiquote val) ...))))

(define (hoc-cname sys v)
  (s+ (hoc-name v) "_" (hoc-name sys)))

(define (hoc-name 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 #\_)) c)
			   (else #\_))))
	    (loop (cons c1 lst) (cdr cs)))))))
			    
(define (matlab-name 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 #\_)) c)
			   (else #\_))))
	    (loop (cons c1 lst) (cdr cs)))))))
	    
(define (const-val v)
  (and (nemo:quantity? v)
       (cases nemo:quantity v
	      (CONST (name value) value)
	      (else (error 'const-val "invalid constant" v)))))

(define (print-hoc-prelude indent)

(pp indent (#<<EOF

vgraphbox.intercept(0)

vec_sizes = tstop/dt + 1	// recorded traces are all this size

soma vce=new VClamp(0)

dt = 0.001

proc vcrun() {
	
	vce.dur[0]=vchdur
	vce.dur[1]=vcbdur
	vce.dur[2]=vchdur

	tstop=vce.dur[0]+vce.dur[1]+vce.dur[2]
	
	vce.amp[0]=vchold
	vce.amp[2]=vchold
	
	for j=0, vcsteps-1 {
		x=vcbase+j*vcincrement
		vce.amp[1]=x

		init()

		run()
		$o2.line(g[$1], dt)

                for i=0,$o2.size()-1 {
                    $o4.printf("%g %g\n", $o3.x[i], $o2.x[i])
                    }

		if (stoppedrun()) {
			break
		}
	}
}
EOF
))
)

(define (nemo:vclamp-translator sys . rest)

  (define (cid x)  (second x))
  (define (cn x)   (first x))
  
  (let-optionals rest ((target 'hoc) (filename #f))

    (cases nemo:quantity (hash-table-ref sys (nemo-intern 'dispatch))  

	 (DISPATCH  (dis) 

	  (let ((imports  ((dis 'imports)  sys))
		(exports  ((dis 'exports)  sys)))
	    
	    (let* ((indent        0)
		   (indent+       (+ 2 indent ))
		   
		   (sysname       (hoc-name ((dis 'sysname) sys)))
		   (filename      (or filename (s+ sysname "_vclamp" (case target ((matlab octave) ".m") (else ".ses")))))

		   (consts        ((dis 'consts)  sys))
		   (asgns         ((dis 'asgns)   sys))
		   (states        ((dis 'states)  sys))
		   (reactions     ((dis 'reactions) sys))
		   (rates         ((dis 'rates) sys))
		   (defuns        ((dis 'defuns)  sys))
		   (components    ((dis 'components) sys))
		   
		   (vc-components (filter-map (lambda (x) (and (eq? 'voltage-clamp (second x)) x)) components))
		   (g             (match-let (((state-list asgn-list g) ((dis 'depgraph*) sys))) g))
		   (poset         (vector->list ((dis 'depgraph->bfs-dist-poset) g)))
		   
		   (gate-complex-info    (nemo:gate-complex-query sys))
		   (gate-complexes       (lookup-def 'gate-complexes gate-complex-info))
		   
		   (i-gates       (lookup-def 'i-gates gate-complex-info))
		   
		   )
	      
	      (if (not (null? vc-components))

		  
		  (case target 
		    
		    ((hoc)

		     (with-output-to-file filename
		       (lambda ()
			 (let* ((vcn   (length vc-components))
				(hboxn (inexact->exact (ceiling (/ vcn 3)))))
		       
			   (pp indent
			       ("load_file(\"nrngui.hoc\")")
			       ("objref vce	// voltage clamp)")
			       ("objref g[" ,vcn "] // graph objects"))

			   (pp indent
			       ("objref vgraphbox, hgraphbox[" ,hboxn "]")
			       ("vgraphbox=new VBox()")
			       ("vgraphbox.intercept(1)"))
			   
			   (let recur ((a 0) (b (min (- vcn 1) 2)) (hbi 0))
			     (pp indent
				 ("hgraphbox[" ,hbi "]=new HBox()")
				 ("hgraphbox[" ,hbi "].intercept(1)"))
			     
			     (if (< a b)
				 (pp indent ("for i=" ,a "," ,b " {"))
				 (pp indent ("i = " ,a)))
			     (pp (if (< a b) indent+ indent)
				 ("g[i]=new Graph()")
				 ("g[i].exec_menu(\"Keep Lines\")"))
			     (if (< a b) (pp indent ("}")))
			     (pp indent
				 ("hgraphbox[" ,hbi "].intercept(0)")
				 ("hgraphbox[" ,hbi "].map()"))
			     (if (> hboxn (+ 1 hbi))
				 (recur (+ b 1) (min (- vcn 1) (+ b 3)) (+ 1 hbi)))
			     )

			   (print-hoc-prelude indent)

			   (for-each 
			    (lambda (comp i) 
			      
			      (let ((name (first comp))
				    (sym  (third comp)))
				
				(match-let (((vchold vcbase vcinc vcsteps vchdur vcbdur) ((dis 'component-exports) sys sym)))

				    (pp indent ("print \"generating " ,name "\""))

				    (pp indent
					("vchold"      = ,(s+ "soma." (hoc-cname sysname vchold)))
					("vcbase"      = ,(s+ "soma." (hoc-cname sysname vcbase)))
					("vcincrement" = ,(s+ "soma." (hoc-cname sysname vcinc)))
					("vcsteps"     = ,(s+ "soma." (hoc-cname sysname vcsteps)))
					("vchdur"      = ,(s+ "soma." (hoc-cname sysname vchdur)))
					("vcbdur"      = ,(s+ "soma." (hoc-cname sysname vcbdur)))
					)
				    
				    (pp indent
					("objref " ,name)
					(,name = "new Vector(vec_sizes)")
					(,(s+ name ".record") ,(s+ "(&soma.i_" name "_" sysname "( 0.5 ))"))
					("objref tlog")
					("tlog = new Vector(vec_sizes,0)")
					("tlog.record (&t)"))
				    
				    (pp indent ("objref logfile")
					("logfile = new File()")
					("logfile.wopen (" ,(s+ "\"" name ".dat\"") ")"))
				    
				    (pp indent ("vcrun(" ,i ", " ,name  ", tlog, logfile)")
					("g[" ,i "].label(.5,.85,\"" ,name "\")")
					"logfile.close()")
				    )))
			    vc-components 
			    (list-tabulate (length vc-components) (lambda (i) i)))
			   
			   (if (> vcn 1) (pp indent ("for i=0, " ,(- vcn 1) " {")) (pp indent "i=0"))
			   (pp (if (> vcn 1) indent+ indent) ("g[i].exec_menu(\"View = plot\")"))
			   (if (> vcn 1) (pp indent ("}")))
			   (pp indent ("vgraphbox.map()"))
			   
			   ))
		       ))

	  ((matlab) 
           (with-output-to-file filename
            (lambda ()
              (let* ((str          (lambda (x) (s+ "\"" x "\"")))
		     (defs         (s+ sysname "_defs"))
		     (init         (s+ sysname "_init"))
		     (current-vars (map (lambda (comp) (s+ "i_" (first comp))) vc-components))
		     (vcn          (length vc-components))
		     (hboxn        (inexact->exact (ceiling (/ vcn 3))))
		     )

           (pp indent (,defs = ,(str (s+ sysname ".m")) #\;)
	       (autoload (,(str sysname ) #\, ,defs ) #\;)
	       (autoload (,(str init) #\, ,defs ) #\;))

           (pp indent (,init #\;))
           (pp indent (y0 = ,init (-65) #\;))


           (pp indent (global N #\;)
                      (N = length(y0) #\;))

           (pp indent (function is = ,(s+ sysname "_currents_odepkg (y)")))
           (pp indent+ (global N #\;)
                       (global ,@(intersperse current-vars #\space) #\;))

           (pp indent+ (t = y(1) #\;)
                       (s = y(2:N+1) #\;)
                       (y1 = ,sysname (t #\, s) #\;))

           (pp indent+ (is = #\[ ,@(cons 't current-vars) #\] #\;))

           (pp indent (endfunction))

           (pp indent (function is = ,(s+ sysname "_currents_lsode (y)")))
           (pp indent+ (global N #\;)
                       (global ,@(intersperse current-vars #\space) #\;))

           (pp indent+ (t = y(1) #\;)
                       (s = y(2:N+1) #\;)
                       (y1 = ,sysname (s #\, t) #\;))

           (pp indent+ (is = #\[ ,@(cons 't current-vars) #\] #\;))

           (pp indent (endfunction))

           (pp indent (function ys = ,(s+ sysname "_vclamp_run1_odepkg") (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur)))

           (pp indent+ (global reltol abstol))

           (pp indent+ (P = odeset (,(str "RelTol") #\, reltol #\, 
				    ,(str "AbsTol") #\, abstol #\, 
				    ,(str "MaxStep") #\, 1 #\, 
				    ,(str "InitialStep") #\, dt) #\;))

           (pp indent+ (t0 = 0.0 #\;)
                       (t1 = t0+vchdur #\;)
		       (t2 = t1 #\;)
		       (t3 = t2+vcdur #\;)
		       (t4 = t3 #\;)
		       (t5 = t4+vchdur #\;))

           (pp indent+ (vinit = vchold #\;)
	               (y0 = ,(s+ sysname "_init") (vinit))
		       (sol = ode2r (,(s+ "@" sysname) #\, #\[ t0 t1 #\] #\, y0 #\, P) #\;)
		       (ys = #\[ sol.x sol.y #\] #\;))

           (pp indent+ (vinit = vcbase #\;)
	       (y0 = sol.y(size(sol.y)(1) #\, :)#\' #\; y0(1) = vinit)
	       (sol = ode2r(,(s+ "@" sysname) #\, #\[ t2 t3 #\] #\, y0 #\, P) #\;)
	       (ys = vertcat (ys #\, #\[ sol.x sol.y #\]) #\;))

           (pp indent+ (vinit = vchold #\;)
	               (y0 = sol.y(size(sol.y)(1) #\, :) #\' #\; y0(1) = vinit)
		       (sol = ode2r(,(s+ "@" sysname) #\, #\[ t4 t5 #\] #\, y0 #\, P) #\;)
		       (ys = vertcat (ys #\, #\[ sol.x sol.y #\]) #\;))

           (pp indent (endfunction #\;))

           (pp indent (function ys = ,(s+ sysname "_vclamp_run1_lsode") (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur)))

           (pp indent+ (global reltol abstol))

           (pp indent+ (lsode_options (,(str "absolute tolerance") #\, abstol) #\;)
	               (lsode_options (,(str "relative tolerance") #\, reltol) #\;)
		       (lsode_options (,(str "integration method") #\, "\"bdf\"") #\;)
		       (lsode_options (,(str "initial step size") #\, dt) #\;))

           (pp indent+ (t0 = 0.0 #\;)
	               (t1 = t0+vchdur #\;)
		       (thold0 = linspace(t0 #\, t1 #\, (t1-t0)/dt) #\;)
		       (t2 = t1 #\;)
		       (t3 = t2+vcdur #\;)
		       (tclamp = linspace(t2 #\, t3 #\, (t3-t2)/dt) #\;)
		       (t4 = t3 #\;)
		       (t5 = t4+vchdur #\;)
		       (thold1 = linspace(t4 #\, t5 #\, (t5-t4)/dt) #\;)
		       (vinit = vchold #\;))

            (pp indent+ (y0 = ,(s+ sysname "_init") (vinit))
		        (y = lsode(,(s+ "@" sysname) #\, y0 #\, thold0) #\;)
			(ys = #\[ "thold0'" y #\] #\;))

            (pp indent+ (vinit = vcbase #\;)
		        (y0 = y(size(y)(1) #\, :) #\' #\; y0(1) = vinit)
			(y = lsode(,(s+ "@" sysname) #\, y0 #\, tclamp) #\;)
			(ys = vertcat (ys #\, #\[ "tclamp'" y #\]) #\;))

            (pp indent+ (vinit = vchold #\;)
		        (y0 = y(size(y)(1) #\, :) #\' #\; y0(1) = vinit)
			(y = lsode(,(s+ "@" sysname) #\, y0 #\, thold1) #\;))

            (pp indent+ (ys = vertcat (ys #\, #\[ "thold1'" y #\]) #\;))

            (pp indent (endfunction))

            (pp indent (function ilog = ,(s+ sysname "_vclamp_fn") (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur #\, vcinc #\, vcsteps)))

            (pp indent+ (vc = vcbase #\;)
		        (ys = cell(1 #\, vcsteps) #\;))

            (pp indent+ (for i = 1:vcsteps))
            (pp (+ 2 indent+) 
		(y = ,(s+ sysname "_vclamp_run1_odepkg") (N #\, dt #\, vchold #\, vchdur #\, vc #\, vcdur) #\;)
		(ys(i) = y #\;)
		(vc = vc + vcinc #\;))
            (pp indent+ (endfor #\;))
  
            (pp indent+ (ilog = cell(vcsteps #\, 1) #\;))
            (pp indent+ (for i = 1:vcsteps))
            (pp (+ 2 indent+)
		(ilogv = #\[#\] #\;)
		(n = size(ys "{i}" )(1) #\;))
            (pp (+ 2 indent+) (for j = 1:n))
            (pp (+ 4 indent+) 
		(next = ,(s+ sysname "_currents_odepkg") (ys "{i}"(j #\, :)) #\;)
		(ilogv = vertcat (ilogv #\, next) #\;))
            (pp (+ 2 indent+) (endfor #\; ))
            (pp (+ 2 indent+) (ilog "{i}" = ilogv #\;))
            (pp indent+ (endfor #\;))
     
            (pp indent (endfunction))

            (pp indent (global reltol abstol))

            (pp indent (reltol = 1e-6 #\;))
            (pp indent (abstol = 1e-6 #\;))
            (pp indent (dt = 0.001 #\;))

	   (for-each
	    (lambda (comp i) 
	      (let* ((name (first comp))
		     (sym  (third comp))
		     (env  ((dis 'component-env) sys sym)))

		(match-let (((vchold vcbase vcinc vcsteps vchdur vcbdur) ((dis 'component-exports) sys sym)))
			   
		   (pp indent ("##" ,(s+ "i_" name) plot))
		   (pp indent (,(s+ "i_" name "_index") = ,(+ 1 i) #\;))

		   (pp indent (global
			       ;; vchold,vchdur,vcbase,vcdur,vcinc,vcsteps
			       ,(matlab-name vchold) 
			       ,(matlab-name vchdur) ,(matlab-name vcbase) 
			       ,(matlab-name vcbdur) ,(matlab-name vcinc) 
			       ,(matlab-name vcsteps)) )

		   (pp indent (,(s+ sysname "_ilog") = ,(s+ sysname "_vclamp_fn") 
			       ;; vchold,vchdur,vcbase,vcdur,vcinc,vcsteps
			       (N #\, dt #\, ,(matlab-name vchold) #\, 
				  ,(matlab-name vchdur) #\, ,(matlab-name vcbase) #\,
				  ,(matlab-name vcbdur) #\, ,(matlab-name vcinc) #\, 
				  ,(matlab-name vcsteps)) #\;))

		   (pp indent (subplot(,hboxn #\, 3 #\, ,i) #\;))
                   (pp indent (plot( ,@(intersperse
					(concatenate
					 (list-tabulate (inexact->exact (const-val (hash-table-ref env vcsteps)))
  				          (lambda (i) 
					   `(,(s+ sysname "_ilog{" (+ 1 i) "}(:,1)") 
					     ,(s+ (s+ sysname "_ilog{" (+ 1 i) "}")
						  (s+ "(:,i_" name "_index)"))))))
				        #\, )
				    )
				   #\;))


                   (pp indent (,(s+ name "_log") = vertcat ,(intersperse 
				       (list-tabulate (inexact->exact (const-val (hash-table-ref env vcsteps)))
				         (lambda (i) 
					   (s+ #\[ 
					       (s+ sysname "_ilog{" (+ 1 i) "}(:,1)") #\,
					       (s+ sysname "_ilog{" (+ 1 i) "}")
					       (s+ "(:,i_" name "_index)")
					       #\])))
				       #\, )
				     #\;))


		   (pp indent (save -ascii ,(str (s+ name ".dat")) ,(s+ name "_log") #\;))

		   )))
		vc-components (list-tabulate vcn (lambda (i) (+ 1 i))))
          ))))
  )))))
)))
)