open Set
open Ast
open Sast
open Printf
open Printer
open Ppm
open Str
open Bigarray

exception InputDeclException of string
exception Exception of string

module IntSet = Set.Make(struct type t = int let compare x y = Pervasives.compare x y end)

let openStream ( param : Sast.param_def ) = 
  Printf.printf "Attempting to open %s.\n" (string_of_param param.pa_num param);
  match param with
   { pa_ptype=InParam; pa_io=SpaceDelimFile(fname,chnl) } -> 
      {  pa_io = SpaceDelimFile( fname, OpenIn(open_in fname) );
         pa_num=param.pa_num; pa_type=param.pa_type; 
         pa_ptype = InParam; pa_data = param.pa_data } 
 | { pa_ptype=OutParam; pa_io=SpaceDelimFile(fname,chnl) } -> 
      {  pa_io = SpaceDelimFile( fname, OpenOut(open_out fname) );
         pa_num=param.pa_num; pa_type=param.pa_type; 
         pa_ptype = OutParam; pa_data = param.pa_data; }
 | { pa_ptype=InParam; pa_io=InPpm(fname, rtime) } -> 
     let data' = Ppm.read_ppm fname in
      {  pa_io = InPpm( fname, rtime );
         pa_num=param.pa_num; pa_type=param.pa_type; 
         pa_ptype = InParam; pa_data = ImageC3(data') } 
 | { pa_ptype=OutParam; pa_io=OutPpmSeq(fname) } ->  (* do nothing*)
     param
 | _ -> raise( Exception "Attempted to open parameter with invalid combination." )



let rec eval_dim = function
    [] -> []
 |  Ast.Intgr(x) :: tl -> x :: (eval_dim tl)
 |  _ -> raise ( InputDeclException "error")

(* exprl should be a list of const_int_expr *)
let rec set_input_dim m (pnum, exprl) =
    let {pa_num=pa_num; pa_type=pa_type;pa_ptype=pa_ptype;
                pa_io=pa_io;pa_data=pa_data } =
      try
        Sast.IntMap.find pnum m
      with Not_found -> 
        raise( Exception ("Param $"^(string_of_int pnum)^" not found in set_input_dim."))
    in let diml = eval_dim exprl in
    let totalsize = List.fold_left (fun siz i -> siz * i) 1 diml
    in let m = IntMap.add 
                 pa_num 
                 { pa_num=pa_num; pa_ptype=pa_ptype;
                   pa_type=Types.Matrix(diml); pa_io=pa_io;
                   pa_data=Unallocated } m
       in (m : param_def IntMap.t)

let rec find_inputs = function
    Ast.ExternalRef(param), set -> add_param( param, set)
  | Ast.Binop(expr1,_,expr2), set -> 
      let set' = find_inputs (expr1, set) in
        find_inputs(expr2, set')
  | Ast.SizeOf(expr,_), set -> find_inputs (expr, set)
  | Ast.Negate(expr), set -> find_inputs (expr, set)
  | Ast.NeuronRef(_,_), set -> set (* Need to make sure indices don't include param! *)
  | Ast.ActRef(_,expr,_), set -> find_inputs( expr, set )
  | Ast.WghtRef(_,_), set -> set
  | Ast.Conv(expr,_), set -> find_inputs( expr, set )
  | Ast.Exp(expr), set -> find_inputs( expr, set )
  | Ast.Sin(expr), set -> find_inputs( expr, set )
  | Ast.Cos(expr), set -> find_inputs( expr, set )
  | Ast.Span(_,_,_), set -> set (* Param must not be in Span *)
  | _,set -> set
and add_param = function
    Ast.Param( num, _ ),set -> IntSet.add num set
  | Ast.ExtNeuron(_,_,_,_),set-> set

let rec find_outputs = function
  Ast.ExtNeuron(_,_,_,_), set -> set
| Ast.Param(num, _), set -> IntSet.add num set

(* creates an IntMap of params *)
let create (program:Ast.program) = 
    let params_in = List.fold_left 
                      (fun s x -> find_inputs(x.s_syn.s_expr,s))  
                IntSet.empty program.p_synap
    and params_out = List.fold_left (fun s x -> find_outputs(x.s_dest,s))
                IntSet.empty program.p_synap
    in
    let push_param = fun pt m x -> 
      let pa_io' = match pt with
          Sast.InParam ->
            if Str.string_match (Str.regexp "\\(.*\\.ppm\\):\\[\\([0-9]+\\)\\]") Sys.argv.(x) 0 then
              ( 
                let fname = Str.matched_group 1 Sys.argv.(x) in
                  Printf.printf "g1='%s'  g2='%s'\n" fname (Str.matched_group 2 Sys.argv.(x)) ;
                let rtime = int_of_string (Str.matched_group 2 Sys.argv.(x)) in
                  InPpm(fname,rtime)
              )
            else SpaceDelimFile(Sys.argv.(x),Closed);
        | Sast.OutParam ->
            if Str.string_match (Str.regexp "\\(.*\\.ppm\\)") Sys.argv.(x) 0 then
              ( 
                let fname = Str.matched_group 1 Sys.argv.(x) in
                  OutPpmSeq(fname)
              )
            else SpaceDelimFile(Sys.argv.(x),Closed);
      in 
        IntMap.add x 
           { pa_num=x;
             pa_type=Types.MatrixUnknownSize; 
             pa_ptype=pt; pa_io= pa_io';
             pa_data = Unallocated;
           } m
    in
    let params = List.fold_left (push_param InParam) IntMap.empty (IntSet.elements params_in)
    in 
      print_endline "Printing input parameters found...";
      IntMap.iter 
        (fun k x -> print_endline (Printer.string_of_param k x) )
        params;
    let params = List.fold_left set_input_dim params program.p_indim
    in let params = List.fold_left 
                (push_param OutParam) params (IntSet.elements params_out)
    in
       assert( IntMap.is_empty params = false );
       (*IntMap.iter (fun k x ->print x) params;*)
       params

