open Str
open Ast
open Sast
open Printf
open Params
open Translate1
open Translate2
open Validate
open Printer
open Eval
open Ppm
open Bigarray

exception Exception of string

let write_param pam tstep = match pam with 
    { pa_io=SpaceDelimFile( fname, OpenOut(cout) ); pa_data=FloatArray(data) }->
         (*Printf.printf "$%d " pam.pa_num;
         List.iter (fun v -> print_float v; print_string "," ) (Array.to_list data);
         print_endline ""; *)
         Array.iter (fun x-> output_string cout ((string_of_float x)^" ")) data;
         output_string cout "\n";
  | { pa_io=OutPpmSeq( fpattern ); pa_data=ImageC3(data) } ->
          let (fname:string) = Printf.sprintf "../tests/images/out/temp%03d.ppm" tstep 
          in
             Printf.printf "Trying to write PPM file '%s'..." fname;
             Ppm.output_ppm fname data;
             Printf.printf "success.\n"
  | _ -> ()
            
let update_param pam tstep = 
  Printf.printf "Calling update_param %d\n" tstep;
  try
  match pam with
    { pa_io = SpaceDelimFile( fname, OpenIn(cin) ) } ->
          let line = input_line cin in
          let vals = List.rev
                       (
                       List.fold_left (fun li word -> 
                                         (float_of_string word)::li ) 
                         [] (Str.split(Str.regexp " ") line)
                       )
          in
            Printf.printf "$%d " pam.pa_num; 
            List.iter (fun v -> print_float v; print_string "," ) vals;
            print_endline "";
             ({pa_num=pam.pa_num; pa_type=pam.pa_type; pa_ptype=pam.pa_ptype; pa_io=pam.pa_io; 
               pa_data=FloatArray(Array.of_list vals) }, true)

  | { pa_io = InPpm( fname, maxreads ); pa_data=ImageC3(data) } -> 
          if tstep > maxreads then raise End_of_file;
          Printf.printf "Trying to read PPM file '%s'..." fname;
          Array3.blit ( Ppm.read_ppm fname ) data; 
          Printf.printf "success.\n";
          pam,true
        (* The output params are not updated here, they are updated when their
        * corresponding synapses are updated *)
  | { pa_io = OutPpmSeq( fpat ) } -> pam,true
  | { pa_io = InPpm( fname, maxreads ) } -> raise (Exception "InPpm must be of type ImageC3.")
  | { pa_io = SpaceDelimFile( fname, OpenOut(cin) ) } -> pam,true
  | { pa_io = SpaceDelimFile( fname, Closed ) } -> raise (Exception "Parameter file is closed!")
               
  with End_of_file -> 
    pam,false
       
let rec get_step ndimsleft t = 
  if ndimsleft = 0 then 1
  else 
    match t with 
        Types.Matrix( d :: dims ) -> 
          if ndimsleft = List.length dims+1 then
            d * (get_step (ndimsleft-1) t)
          else if ndimsleft < List.length dims+1 then
            (get_step (ndimsleft-1) t)
          else
            raise (Exception "Opps hola.")
      | _ -> raise (Exception "Type of param must be matrix.")

let rec find_for_expr forvar forlist = match forlist with
    [] -> raise (Exception ("Didn't find "^ forvar ^" in for-list."))
  | ForExpr(forvar', exprd) :: forlist' ->
      if forvar=forvar' then
        ( forvar',exprd )
      else 
        find_for_expr forvar forlist'
  | _ -> raise (Exception "Invalid expression found in for-list.")

let rec iter_matrix_expr f expr_dim =
  iter_matrix_expr_loop f expr_dim expr_dim [] 
and 
    iter_matrix_expr_loop f uneval_dim expr_dim expr_indl =
  match uneval_dim with
      [] -> (f expr_indl)
    | dim_expr :: uneval_dim' ->
        for ind = 1 to dim_expr do
          (iter_matrix_expr_loop f uneval_dim' expr_dim (ind::expr_indl))
        done

(* executes f(evalenv) when evalenv has all of the forbindings *)
let rec iter_forloop f evalenv dimUneval sfor =
  match dimUneval with
      [] -> (f evalenv)
    | dim_expr :: dimUneval' ->
        let forvarl = Translate2.retrieve_forvar dim_expr in
        let forvar = match forvarl with
            [] -> raise( Exception
                           ("Blaahhk - no allowy?"))
          | [varn] -> varn
          | varn :: li ->
              raise( Exception
                       "Only a single for-variable can be in a single dimension index.")
        in
        let (forvar,span) = (find_for_expr forvar sfor) in
        let nums = Eval.nums_of_span span in 
          List.iter (fun n -> 
                       let evalenv' = 
                         { forbindings = ( (forvar,n) :: evalenv.forbindings ); 
                           varbindings = evalenv.varbindings } 
                       in 
                         iter_forloop f evalenv' dimUneval' sfor ) nums

let iter_forloop_and_matrix_expr f evalenv indl sfor expr_diml = 
  (iter_forloop
     (fun evalenv -> 
        (iter_matrix_expr 
           (f evalenv) 
           expr_diml))
     evalenv indl sfor)

let update_synap syn p tstep = 
  Printf.printf "Calling update_synap t=%d for %s\n" 
    tstep (Printer.string_of_synap 0 syn);
  match syn with
    { sdest=sdest; sexpr=(sexprd,Types.Int); sfor=sfor} -> assert false
  | { sdest=sdest; sexpr=(sexprd,Types.MatrixUnknownSize); sfor=sfor} -> assert false
  | { sdest=sdest; sexpr=(sexprd,Types.Matrix(ediml)); sfor=sfor} -> 
        ( 
          match sdest with 
              Param(pnum, indl) ->
                Printf.printf("Param\n");
                let pam = (IntMap.find pnum p.pparams) in
                  (
                  match pam with
                      { pa_data=FloatArray(data) } -> 
                        iter_forloop_and_matrix_expr
                          (fun evalenv eindl -> 
                             Printf.printf "expr index=";
                             List.iter (Printf.printf "%d,") eindl;
                             Printf.printf "\n";
                             let dloc = dataloc indl p evalenv pam.pa_type eindl in
                             let valu = Eval.eval syn.sexpr eindl p evalenv in 
                             Printf.printf "dloc=%d value=%.2f\n" dloc valu;
                               data.(dloc)  <-  valu 
                          )
                          { forbindings=[]; varbindings=StringMap.empty } 
                          indl 
                          sfor
                          ediml
                    | { pa_data=ImageC3(data); pa_type=Types.Matrix(diml) } -> 
                        iter_forloop_and_matrix_expr 
                          (fun evalenv eindl -> 
                             let (i1,i2,i3) = dataloc3 indl p evalenv diml eindl in
                               let pixval = int_of_float 
                                              (255.*.(Eval.eval syn.sexpr eindl p evalenv))
                               in
                                 data.{i1,i2,i3} <- pixval;
                          )
                          { forbindings=[]; varbindings=StringMap.empty } 
                          indl 
                          sfor
                          ediml;
                        Printf.printf "success!\n"
                    | { pa_data=Unallocated }  -> raise (Exception "Parameter data unallocated!")
                    | { pa_data=ImageC3(data)} -> raise (Exception "Image must have type of Matrix!")
                  )
            | ExtNeuron( gsid, _,[],nname,nindl) ->
                Printf.printf("ExtNeuron\n");
                iter_forloop_and_matrix_expr 
                  (fun evalenv eindl -> 
                     let newval = Eval.eval syn.sexpr eindl p evalenv in 
                     let neuron = (IntMap.find gsid p.pneurons) in 
                     let nloc = dataloc nindl p evalenv neuron.t eindl in
                       Printf.printf "neuron %s[(%d)] = %.3f\n" nname nloc newval;
                       print_endline (Printer.string_of_neuronvals neuron.v);
                       neuron.v.(nloc).(0) <- newval
                  )
                  { forbindings=[]; varbindings=StringMap.empty }
                  nindl
                  sfor
                  ediml

            | NeuronRef( gsid, name, nindl) ->
                Printf.printf "NeuronRef of %d/%s\n" gsid name;
                iter_forloop_and_matrix_expr 
                  (fun evalenv eindl ->
                     Printf.printf 
                       "with bindings %s\n" 
                       (Printer.string_of_forbindings evalenv.forbindings);
                     let newval = Eval.eval syn.sexpr eindl p evalenv in 
                     let neuron = (IntMap.find gsid p.pneurons) in 
                     let nloc = dataloc nindl p evalenv neuron.t eindl in 
                       Printf.printf "local neuron %s[(%d)] == %.3f\n" name nloc newval;
                       neuron.v.(nloc).(0) <- newval
                  )
                  { forbindings=[]; varbindings=StringMap.empty }
                  nindl
                  sfor
                  ediml;
                Printf.printf("end NeuronRef\n")

            | expr -> raise(Exception 
                              (Printf.sprintf
                                 "Unhandled code %s N2n9." (Printer.string_of_exprd expr )) )
        ) 


let rec run_step (p:Sast.prog) tstep =
  Printf.printf "#### Time step %d ####\n" tstep;
  let params',success = IntMap.fold 
                          (fun k pam (m,a) -> 
                               let (pam',b) = update_param pam tstep in (IntMap.add pam'.pa_num pam' m),(b && a) ) 
                          p.pparams (IntMap.empty, true) 
  in 
  let p = { syn=p.syn; pparams=params'; 
            actfuns=p.actfuns; kerfuns=p.kerfuns; pneurons=p.pneurons }
  in
    if success=false then () 
    else 
      ( 
        IntMap.iter (fun k syn -> update_synap syn p tstep) p.syn;
        IntMap.iter (fun k n -> (*Printf.printf 
                                  "Neuron %d [|%f,%f|]\n" k n.v.(0).(0) n.v.(1);  *)
                                for i = 0 to Array.length n.v-1 do 
                                  n.v.(i).(1) <- n.v.(i).(0) 
                                done;
                    ) p.pneurons;
        IntMap.iter (fun k pam -> write_param pam tstep) params';
        run_step p (tstep+1)
      )


let run (p:Sast.prog)  =
   let pparams' = IntMap.map Params.openStream p.pparams in
     run_step {syn=p.syn; pparams=pparams'; 
               actfuns=p.actfuns; kerfuns=p.kerfuns; pneurons=p.pneurons } 1

let _ =
  let lexbuf = Lexing.from_channel stdin in
  let p = Parser.program Scanner.token lexbuf in
  let p = Translate1.program p in
    print_endline "\n\nProgram after Translate1:";
    print_string (Printer.string_of_program p);
  let p = Translate2.program p in
    print_endline "\n\nProgram after Translate2:";
    print_string (Printer.string_of_program p);
    Validate.sast2 p;
    IntMap.iter (fun pnum pam -> 
                   print_endline 
                     (Printer.string_of_param pnum pam)) 
      p.pparams;
    ignore (run p)

