val zeropid = Posix.Process.wordToPid (SysWord.fromInt 0)


fun putStrLn out str = 
    (TextIO.output (out, str);
     TextIO.output (out, "\n"))
    
fun putStr out str = 
    (TextIO.output (out, str))
    
fun showBoolean b = (if b then "1" else "0")

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

fun foldl1 f lst = let val v = List.hd lst
                       val lst' = List.tl lst
                   in
                       List.foldl f v lst'
                   end

fun realRandomTensor randarray shape =
    let 
        val length = Index.length shape
    in
       RTensor.fromArray (shape, randarray length)
    end

exception Index        
 
val order = 2500     (* scales size of network (total 5*order neurons) *)


val Ne = 4 * order  (* number of excitatory neurons *)
val Ni = 1 * order  (* number of inhibitory neurons *)
val N  = Ni+Ne

val epsilon = 0.1     (* connectivity probability *)

val Ce = epsilon * (Real.fromInt Ne) (* total number of excitatory synapses *)
val Ci = epsilon * (Real.fromInt Ni) (* total number of inhibitory synapses *)
val C  = Real.+(Ce,Ci) (* total number of internal synapses *)


val C       = 1.0 (* neuron membrane capacitance  *) 
val R       = 20.0 (* neuron membrane resistance  *) 
val tau     = Real.* (R, C) (* neuron membrane time constant [ms] *) 
val tau_rp  = 2.0 (* refractory time [ms] *) 
val theta   = 20.0 (* spike threshold [mV] *)
val Vreset  = 10.0 (* reset potential [mV] *)

val  delay  = 1.5 (* synaptic delay, all connections [ms]  *)
val  J      = 0.1 (* synaptic weight [mV] *)


(*
 Ci, Ce: number of source neurons for inh and exc connections, respectively 
 Ji, Je: connection weights for inh, exc connections
 *)

val seed_e = RandomMTZig.fromEntropy() (* seed for excitatory connections *)
val seed_i = RandomMTZig.fromEntropy() (* seed for inhibitory connections *)

val h = 0.1
val I0 = RTensor.new([N,1],0.0)

(* delay expressed as # time steps *)
val D  = (List.tabulate (Real.round (Real./(delay,h)), fn i => I0))

val seed_v = RandomMTZig.fromEntropy() (* seed for initial membrane potential *)
val zt_v   = RandomMTZig.ztnew()

fun v_init () = Real.* (Vreset, RandomMTZig.randNormal(seed_v,zt_v))


val BrunelIaF_Inhibitory_initial_vector = 
    Vector.tabulate (Ni, fn (i) => {rpstate=(false),rpend=(false),spike=(false),h=h,t=(0.0),tau_rp=tau_rp,tspike=(0.0),Vreset=Vreset,theta=theta,V=v_init(),Isyn=(0.0),tau=tau,R=R})


val BrunelIaF_Excitatory_initial_vector = 
    Vector.tabulate (Ne, fn (i) => {rpstate=(false),rpend=(false),spike=(false),h=h,t=(0.0),tau_rp=tau_rp,tspike=(0.0),Vreset=Vreset,theta=theta,V=v_init(),Isyn=(0.0),tau=tau,R=R})

val BrunelIaF_Inhibitory_f = BrunelIaF.TestBrunelIaFdeltaivp146

fun BrunelIaF_Inhibitory_run (Wnet,n0) (i,input as {t,h,V,tau,spike,tspike,theta,Vreset,tau_rp,rpend,rpstate,Isyn,R}) =
    let 
        val Isyn_i  = RTensor.sub(Wnet,[i+n0,0])
        (*val _ = putStrLn TextIO.stdOut ("# Inh: Isyn_i = " ^ (showReal Isyn_i) ^ " V = " ^ (showReal V))*)
        val nstate = BrunelIaF_Inhibitory_f {t=t,h=h,Isyn=Isyn_i,V=V,tau=tau,spike=spike,tspike=tspike,theta=theta,Vreset=Vreset,tau_rp=tau_rp,rpend=rpend,rpstate=rpstate,R=R}
        val nstate1 = {rpstate=(#rpstate(nstate)),rpend=(#rpend(nstate)),spike=(#spike(nstate)),h=h,t=(#t(nstate)),tau_rp=tau_rp,tspike=(#tspike(nstate)),Vreset=Vreset,theta=theta,V=(#V(nstate)),tau=tau,Isyn=Isyn,R=R}
    in 
        nstate1
    end

val BrunelIaF_Excitatory_f = BrunelIaF.TestBrunelIaFdeltaivp146
        
fun BrunelIaF_Excitatory_run (Wnet,n0) (i,input as {t,h,V,tau,spike,tspike,theta,Vreset,tau_rp,rpend,rpstate,Isyn,R}) =
    let 
        val Isyn_i = RTensor.sub(Wnet,[i+n0,0])
        (*val _ = putStrLn TextIO.stdOut ("# Exc: Isyn_i = " ^ (showReal Isyn_i) ^ " V = " ^ (showReal V))*)
        val nstate = BrunelIaF_Excitatory_f {t=t,h=h,Isyn=Isyn_i,V=V,tau=tau,spike=spike,tspike=tspike,theta=theta,Vreset=Vreset,tau_rp=tau_rp,rpend=rpend,rpstate=rpstate,R=R}
        val nstate1 = {rpstate=(#rpstate(nstate)),rpend=(#rpend(nstate)),spike=(#spike(nstate)),h=h,t=(#t(nstate)),tau_rp=tau_rp,tspike=(#tspike(nstate)),Vreset=Vreset,theta=theta,V=(#V(nstate)),tau=tau,Isyn=Isyn,R=R}
    in 
        nstate1
    end
        
val Ext_f = PoissonNeuron.Poissonfn        

fun Ext_run (i,input) =
    let 
        val nstate1 = Ext_f input
    in 
        nstate1
    end


fun paramset (J_gain, Ce_gain) =
    [
     (let
          (* case A : almost full synchronization *)
          val g    =  3.0 (* rel strength, inhibitory synapses *)
          val eta  =  2.0 (* nu_ext / nu_thr *)
      in
          ("caseAdelta" ^ "_Jg" ^ (Real.toString J_gain) ^ "_Cg" ^ (Real.toString Ce_gain),g,eta,J_gain,Ce_gain)
      end),
     
     
     (let
          (* case B : fast oscillation, irregular activity *)
          val g    =  6.0 (* rel strength, inhibitory synapses *)
          val eta  =  4.0 (* nu_ext / nu_thr *)
      in
          ("caseBdelta" ^ "_Jg" ^ (Real.toString J_gain) ^ "_Cg" ^ (Real.toString Ce_gain),g,eta,J_gain,Ce_gain)
      end),
     
     (let
          (* case C : slow oscillations *)
          val g    =  5.0 (* rel strength, inhibitory synapses *)
          val eta  =  2.0 (* nu_ext / nu_thr *)
      in
          ("caseCdelta" ^ "_Jg" ^ (Real.toString J_gain) ^ "_Cg" ^ (Real.toString Ce_gain),g,eta,J_gain,Ce_gain)
      end),
     
     (let
          (* case D : slow oscillations *)
          val g    =  4.5 (* rel strength, inhibitory synapses *)
          val eta  =  0.9 (* nu_ext / nu_thr *)
      in
          ("caseDdelta" ^ "_Jg" ^ (Real.toString J_gain) ^ "_Cg" ^ (Real.toString Ce_gain),g,eta,J_gain,Ce_gain)
      end)
    ]


fun randomseed (i) = RandomMTZig.fromInt(i+1)


fun dispatch (label,g,eta,J_gain,Ce_gain) =
    let

        val Cext = Real.* (Ce, Ce_gain) (* total number of external synapses *)
        val Np   = Real.round (Cext) (* total number of external poisson generators *)
        
        (* excitatory weights *)
        val Je =  Real.* (J_gain, J)
                 
        (* inhibitory weights *)
        val Ji = Real.* (Real.~ g, Je)
                 
        (* external weights *)
        val Jext = Je
                   
        (* threshold rate*)
        val nu_thresh = Real./ (theta, (foldl1 Real.* [Ce, Je, tau]))
                        
        (* external rate per synapse *)
        val nu_ext    = foldl Real.* 1.0 [eta, nu_thresh]
                        
        (* mean input spiking rate [Hz] *)
        val input_rate    = Real.* (1000.0, nu_ext)

        val Ext_initial_vector = Vector.tabulate
                                     (Np, fn (i) => {spike=(false),h=h,t=(0.0),spikecount=0.0,
                                                     rate=input_rate,
                                                     st=randomseed i,
                                                     zt=RandomMTZig.ztnew()})


        val initial = (BrunelIaF_Inhibitory_initial_vector,
                       BrunelIaF_Excitatory_initial_vector,
                       Ext_initial_vector)


        val BrunelIaF_Inhibitory_n0 = 0
        val BrunelIaF_Excitatory_n0 = BrunelIaF_Inhibitory_n0 + (Vector.length BrunelIaF_Inhibitory_initial_vector)
        val Ext_n0 = BrunelIaF_Excitatory_n0 + (Vector.length BrunelIaF_Excitatory_initial_vector)
                
        val Pn = [BrunelIaF_Inhibitory_n0, BrunelIaF_Excitatory_n0, Ext_n0]

        fun frun (I,[Inh_n0,Exc_n0,Ext_n0])
                 (BrunelIaF_Inhibitory_state_vector,
                  BrunelIaF_Excitatory_state_vector,
                  Ext_state_vector) =
            let
                val BrunelIaF_Inhibitory_state_vector' = 
                    Vector.mapi (BrunelIaF_Inhibitory_run (I,Inh_n0)) 
                                BrunelIaF_Inhibitory_state_vector
                val BrunelIaF_Excitatory_state_vector' = 
                    Vector.mapi (BrunelIaF_Excitatory_run (I,Exc_n0))
                                BrunelIaF_Excitatory_state_vector
                val Ext_state_vector' = 
                    Vector.mapi Ext_run Ext_state_vector
            in
                (BrunelIaF_Inhibitory_state_vector',
                 BrunelIaF_Excitatory_state_vector',
                 Ext_state_vector') 
            end
            | frun (I,_) _ = raise Index
                                        

        fun ftime (BrunelIaF_Inhibitory_state_vector,
                   BrunelIaF_Excitatory_state_vector,
                   Ext_state_vector) =

            (#t (Vector.sub (BrunelIaF_Inhibitory_state_vector,0)))

                 
        fun fspikes (BrunelIaF_Inhibitory_state_vector,
                     BrunelIaF_Excitatory_state_vector,
                     Ext_state_vector) =
            let
                val BrunelIaF_Inhibitory_spike_i = 
                    Vector.foldri (fn (i,v,ax) => (if (#spike(v)) then ((i+BrunelIaF_Inhibitory_n0,1.0))::ax else ax)) 
                                  [] BrunelIaF_Inhibitory_state_vector
                val BrunelIaF_Excitatory_spike_i = 
                    Vector.foldri (fn (i,v,ax) => (if (#spike(v)) then ((i+BrunelIaF_Excitatory_n0,1.0)::ax) else ax)) 
                                  [] BrunelIaF_Excitatory_state_vector
                val Ext_spike_i = 
                    Vector.foldri (fn (i,v,ax) => (if (#spike(v)) then ((i+Ext_n0),(#spikecount(v)))::ax else ax)) 
                                  [] Ext_state_vector
                    
                val neuron_spike_i = List.concat [BrunelIaF_Inhibitory_spike_i, BrunelIaF_Excitatory_spike_i]
                val all_spike_i    = List.concat [neuron_spike_i, Ext_spike_i]
            in
                (all_spike_i, neuron_spike_i)
            end


        fun finfo ((BrunelIaF_Inhibitory_state_vector,
                    BrunelIaF_Excitatory_state_vector,
                    Ext_state_vector),out) = ()

                                     
        val pid = Posix.Process.fork()
                  
    in
        
        case pid of
            SOME n => n
          | NONE =>
            (let

                 val Si  = RTensor.*> Ji
                                   (RTensor.map (fn (x) => (if Real.> (epsilon,x) then 1.0 else 0.0))
                                                (realRandomTensor 
                                                     (fn (len) => RandomMTZig.randUniformArray (len,seed_i))
                                                     [N,Ni]))

                 val Se  = RTensor.*> Je
                                   (RTensor.map (fn (x) => (if Real.> (epsilon,x) then 1.0 else 0.0))
                                                (realRandomTensor 
                                                     (fn (len) => RandomMTZig.randUniformArray (len,seed_e))
                                                     [N,Ne]))

                 val Sp  = RTensor.*>  Jext (RTensor.new ([N,Np],1.0))

                 val S   = foldl1 (fn (x,ax) => RTensor.cat (ax,x,1)) [ Si, Se, Sp ]

                 val out = TextIO.openOut (label ^ "_poisson.dat")
             in
                 List.app (fn (s) => putStrLn out ("# " ^ s))
                          ([
                           label,
                           ("Ce = " ^ (Real.toString Ce)),
                           ("Ci = " ^ (Real.toString Ci)),
                           ("N = " ^ (Int.toString N)),
                           ("Ne = " ^ (Int.toString Ne)),
                           ("Ni = " ^ (Int.toString Ni)),
                           ("Np = " ^ (Int.toString Np)),
                           ("nu_ext = " ^ (Real.toString nu_ext)),
                           ("nu_thresh = " ^ (Real.toString nu_thresh)),
                           ("input_rate = " ^ (Real.toString input_rate)),
                           ("J = " ^ (Real.toString J)),
                           ("Je = " ^ (Real.toString Je)),
                           ("Ji = " ^ (Real.toString Ji)),
                           ("Jext = " ^ (Real.toString Jext)),
                           ("D = " ^ (Int.toString (List.length D)))
                          ]);
                 
                 Net.start (1200.0, N, S, D, Pn, initial, frun, ftime, fspikes, finfo, out);
                 
                 TextIO.flushOut (out);
                 TextIO.closeOut (out); 
                 zeropid
             end)
    end



fun dispatch' x k =
    (let val pid = dispatch x
     in
         if not (pid = zeropid) then k(pid) else ()
     end)


exception Exit of OS.Process.status 

fun main (progname, args) =
    let

        fun recur (paramset) =
            (case paramset of
                 (x1::x2::rst) => 
                 (dispatch' x1 (fn (pid1) => 
                                   dispatch' x2 (fn (pid2) => 
                                                    (Posix.Process.waitpid (Posix.Process.W_CHILD pid1,[]); 
                                                     Posix.Process.waitpid (Posix.Process.W_CHILD pid2,[]); 
                                                     recur rst))))
               | (x::[])       => (let val pid = dispatch x 
                                   in 
                                       if not (pid = zeropid) 
                                       then (Posix.Process.waitpid (Posix.Process.W_ANY_CHILD,[]); ())
                                       else ()
                                   end)

               | ([])          => ())

    in
        case args of
            [x1,x2] =>
            (let
	         val Jgrangemin = Int.fromString x1
	         val Jgrangemax = Int.fromString x2
             in
                 case (Jgrangemin,Jgrangemax) of
                     (SOME Jgrangemin',SOME Jgrangemax') =>
                     (recur (foldl (fn (Ig,ax) => 
                                       let
                                           val lst = List.tabulate
                                                         (Jgrangemax'-Jgrangemin', 
                                                       fn Jg => paramset (Real.fromInt (Jgrangemin'+Jg),Ig))
                                       in
                                           (List.concat lst) @ ax
                                       end)
                                   [] (List.tabulate (10, fn i => Real.fromInt i))))
                     
                   | (_, _) => raise Exit (OS.Process.failure)
             end)
            | [] =>
              (recur ([hd (tl (tl (tl (paramset (10.0, 9.75)))))]))
            | _ => ()
               
    end
    

val _ = let val name = CommandLine.name()
	    val args = CommandLine.arguments()
	    val env  = Posix.ProcEnv.environ()
	in
	    main (name, args)
	end