;;       
;; 
;; An extension for generating voltage clamp scripts from NEMO models.
;;
;; Copyright 2008-2013 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 (escape-str x) (s+ "\"" x "\""))

(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 (hoc-prelude indent)


(pp indent (#<<EOF

vgraphbox.intercept(0)

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

vce=new SEClamp(0)

dt = 0.001

proc vcrun() {
	
	vce.dur1 = vchdur
	vce.dur2 = vcbdur
	vce.dur3 = vchdur

	tstop=vce.dur1 + vce.dur2 + vce.dur3
	
	vce.amp1 = vchold
	vce.amp3 = vchold
	
	for j=0, vcsteps-1 {
		x=vcbase+j*vcincrement
		vce.amp2=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 (matlab-vclamp-run indent indent+ env i
                           sysname hboxn name vchold vcbase vcinc vcsteps vchdur vcbdur)
                                           
  (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 
                   ,(escape-str (s+ name ".dat")) 
                   ,(s+ name "_log") #\;))
  )
                                         

(define (matlab-cnvclamp-run indent indent+ env i
                             sysname hboxn name vchold vcbase vcinc vcsteps vchdur vcbdur
                             mod-ion-name cnhold cnbase cninc cnsteps)
                                           
  (pp indent (,(s+ sysname "_cnilog") = ,(s+ sysname "_cnvclamp_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
                       (append
                        (list-tabulate 
                          (inexact->exact (const-val (hash-table-ref env cnsteps)))
                          (lambda (i)
                            (concatenate
                             (list-tabulate
                              (inexact->exact (const-val (hash-table-ref env vcsteps)))
                              (lambda (j) 
                                `(,(s+ sysname "_cnilog{" (+ 1 i) "," (+ 1 j) "}(:,1)") 
                                  ,(s+ (s+ sysname "_cnilog{" (+ 1 i) "," (+ 1 j) "}")
                                       (s+ "(:,i_" name "_index)"))
                                  ))
                              ))
                            )
                          ))
                       #\, )
                    )
                  #\;))

  
  (pp indent (,(s+ name "_log") = vertcat ,(intersperse 
                                            (append
                                             (list-tabulate 
                                              (inexact->exact (const-val (hash-table-ref env cnsteps)))
                                              (lambda (i)
                                                (list-tabulate
                                                 (inexact->exact (const-val (hash-table-ref env vcsteps)))
                                                 (lambda (j) 
                                                   (s+ #\[ 
                                                       (s+ sysname "_cnilog{" (+ 1 i) "," (+ 1 i) "}(:,1)") #\,
                                                       (s+ sysname "_cnilog{" (+ 1 i) "," (+ 1 i)  "}")
                                                       (s+ "(:,i_" name "_index)")
                                                       #\]))
                                                 ))
                                              ))
                                            #\, )
              #\;))
  
  (pp indent (save -ascii 
                   ,(escape-str (s+ name "_" mod-ion-name  ".dat")) 
                   ,(s+ name "_" mod-ion-name "_log") #\;))
  )
                                         

(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 ".hoc")))))

		   (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))
                   (mod-ions             (lookup-def 'mod-ions gate-complex-info))
		   )
	      
	      (if (not (null? vc-components))

		  
		  (case target 
		    
                    
                    ((matlab) 
                     (with-output-to-file filename
                       (lambda ()
                         (let* (
                                (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 = ,(escape-str (s+ sysname ".m")) #\;)
                               (autoload (,(escape-str sysname ) #\, ,defs ) #\;)
                               (autoload (,(escape-str init) #\, ,defs ) #\;))
                           
                           (pp indent (,init #\;))
                           (pp indent (y0 = ,init (-65) #\;))
                           
                           (pp indent (global N #\;)
                               (N = length(y0) #\;))
                           
                           ;; TODO: use the logistic function for voltage stepping 
                           (pp indent (function dy = ,(s+ sysname "_dy_vclamp_odepkg") (t #\, y)))
                           
                           (pp indent+ (global vc #\;))
                           
                           (pp indent+ (dy = ,sysname  (t #\, y) #\;))
                           
                           (pp indent+ (i_elec = (vc - y(1)) #\;))
                           
                           (pp indent+ (dy(1) = dy(1) + i_elec #\;))
                           
                           (pp indent (endfunction))
                           
                           (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 vc))
                           
                           (pp indent+ (P = odeset (,(escape-str "RelTol") #\, reltol #\, 
                                                    ,(escape-str "AbsTol") #\, abstol #\, 
                                                    ,(escape-str "MaxStep") #\, 1 #\, 
                                                    ,(escape-str "InitialStep") #\, dt) #\;))
                           
                           (pp indent+ (t0 = 0.0 #\;)
                               (t1 = t0+vchdur #\;)
                               (t2 = t1 #\;)
                               (t3 = t2+vcdur #\;)
                               (t4 = t3 #\;)
                               (t5 = t4+vchdur #\;))
                           
                           (pp indent+ (vc = vchold #\;)
                               (vinit = -65 #\;)
                               (y0 = ,(s+ sysname "_init") (vinit))
                               (sol = ode2r (,(s+ "@" sysname "_dy_vclamp_odepkg") #\, #\[ t0 t1 #\] #\, y0 #\, P) #\;)
                               (ys = #\[ sol.x sol.y #\] #\;))
                           
                           (pp indent+ (vc = vcbase #\;)
                               (y0 = sol.y(size(sol.y)(1) #\, #\:)#\' #\; )
                               (sol = ode2r(,(s+ "@" sysname "_dy_vclamp_odepkg") #\, #\[ t2 t3 #\] #\, y0 #\, P) #\;)
                               (ys = vertcat (ys #\, #\[ sol.x sol.y #\]) #\;))
                           
                           (pp indent+ (vc = vchold #\;)
                               (y0 = sol.y(size(sol.y)(1) #\, #\:) #\' #\; )
                               (sol = ode2r(,(s+ "@" sysname "_dy_vclamp_odepkg") #\, #\[ 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 vc))
                           
                           (pp indent+ (lsode_options (,(escape-str "absolute tolerance") #\, abstol) #\;)
                               (lsode_options (,(escape-str "relative tolerance") #\, reltol) #\;)
                               (lsode_options (,(escape-str "integration method") #\, "\"bdf\"") #\;)
                               (lsode_options (,(escape-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) #\;))
                           
                           (pp indent+ (vc = vchold #\;) 
                               (vinit = -65 #\;) 
                               (y0 = ,(s+ sysname "_init") (vinit))
                               (y = lsode(,(s+ "@" sysname "_dy_vclamp") #\, y0 #\, thold0) #\;)
                               (ys = #\[ "thold0'" y #\] #\;))
                           
                           (pp indent+ (vc = vcbase #\;)
                               (y0 = y(size(y)(1) #\, #\:) #\' #\; )
                               (y = lsode(,(s+ "@" sysname "_dy_vclamp") #\, y0 #\, tclamp) #\;)
                               (ys = vertcat (ys #\, #\[ "tclamp'" y #\]) #\;))
                           
                           (pp indent+ (vc = vchold #\;)
                               (y0 = y(size(y)(1) #\, #\:) #\' #\; )
                               (y = lsode(,(s+ "@" sysname "_dy_vclamp") #\, y0 #\, thold1) #\;))
                           
                           (pp indent+ (ys = vertcat (ys #\, #\[ "thold1'" y #\]) #\;))
                           
                           (pp indent (endfunction))
                           

                           (pp indent (function ilog = ,(s+ sysname "_cnvclamp_fn") 
                                                (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur #\, vcinc #\, vcsteps #\, 
                                                   cname #\, cnhold #\, cnbase cninc #\, cnsteps )))
                           
                           (pp indent+ 
                               (ys = cell(cnsteps #\, vcsteps) #\;)
                               (cn = cnbase #\;))

                           (pp indent+ (for i = 1:cnsteps))
                           (pp (+ 2 indent+) 
                               (y = ,(s+ sysname "_vclamp_fn") (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur #\, vcinc #\, vcsteps)
                                  (ys(i,#\:) = y #\;)
                                  (cn = cn + cninc #\;)
                                  ))
                           (pp indent+ (endif))

                           (pp indent (endfunction))


                           (pp indent (function ilog = ,(s+ sysname "_vclamp_fn") (N #\, dt #\, vchold #\, vchdur #\, vcbase #\, vcdur #\, vcinc #\, vcsteps)))
                           
                           (pp indent+ 
                               (ys = cell(1 #\, vcsteps) #\;)
                               (vc = vcbase #\;))
                           
                           (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 . rest) 
                                             ((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)) )
                                           
                                           (if (not (null? mod-ions))
                                               (if (not (= (length mod-ions) (* 4 (length rest))))
                                                   (error 'nemo:vclamp-translator "voltage clamp component " name 
                                                          " does not have enough entries for modulating ions: " 
                                                          (map first mod-ions))
                                                   (let recur ((rest rest) (mod-ions mod-ions))
                                                     (if (not (null? rest))
                                                         (let ((mod-ion (car mod-ions)))
                                                           (match-let (((cnhold cnbase cninc cnsteps . rest1) rest))
                                                                      (begin
                                                                        (pp indent (global
                                                                                    ,(matlab-name cnhold) 
                                                                                    ,(matlab-name cnbase) 
                                                                                    ,(matlab-name cninc)
                                                                                    ,(matlab-name cnsteps)))
                                                                        (matlab-cnvclamp-run indent indent+ env i sysname
                                                                                             hboxn name vchold vcbase vcinc vcsteps vchdur vcbdur
                                                                                             (first mod-ion) cnhold cnbase cninc cnsteps)
                                                                        )
                                                                      (recur rest1 (cdr mod-ions))
                                                                      ))
                                                         ))
                                                   )

                                               (matlab-vclamp-run indent indent+ env i sysname hboxn name vchold vcbase vcinc vcsteps vchdur vcbdur)
                                               ))
                                ))
                            vc-components (list-tabulate vcn (lambda (i) (+ 1 i))))
                           ))
                       ))

	  ((hoc)
	   
	   (with-output-to-file filename
	     (lambda ()
	       (let* ((vcn   (length vc-components))
		      (hboxn (inexact->exact (ceiling (/ vcn 3)))))
		 
		 (pp indent
		     ("load_file(\"nrngui.hoc\")")
		     ("load_file(" ,(s+ #\" sysname ".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)))
		   )
		 
		 (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()"))
		 
		 ))
	     ))

  )))))
)))
)