open Sast
open Params
open Printer

exception Exception of string

type for_binding = string * int

type environ = {
  forbindings: for_binding list;
  (* varbindings are not passed into the scopes of nested functions *)
  varbindings: float StringMap.t;
}

(* The numbers should be resolved to Intgr types before this is called. *)
let rec nums_of_span span = match span with
    Sast.Span(Intgr(x1),Intgr(x2),Intgr(x3))-> nums_in_range x1 x2 x3
  | exprd -> raise( Exception ("Expected Span, found " ^ Printer.string_of_exprd exprd ))
and
    nums_in_range i1 incr i2 = 
  if ((i1 < i2 && incr > 0) || ( i1 > i2 && incr < 0))  then
    i1 :: (nums_in_range (i1+incr) incr i2 )
  else if i1 == i2 then [i1]
  else
        raise( Exception (Printf.sprintf "Error: [%d:%d:%d] is an invalid span!" i1 incr i2))

(* Generalize fold_left over three lists. 
 * http://svn.metaprl.org/svnroot/mojave/metaprl-tags/before_ocaml_3_02/mllib/list_util.ml
 *) 
let rec fold_left3 f arg l1 l2 l3 = 
  match l1, l2, l3 with 
      h1 :: t1, h2 :: t2, h3 :: t3 -> 
        fold_left3 f (f arg h1 h2 h3) t1 t2 t3 
    | [], [], [] -> arg 
    | _ -> raise (Failure "List_util.fold_left3") 
           

let rec retrieve_binding varname fbindings = match fbindings with
    [] -> raise 
            (Exception 
               (Printf.sprintf "Unable to find for-var '%s'." varname))
  | (bvar,bval)::fbindings' ->
      if bvar = varname then bval
      else
        retrieve_binding varname fbindings'

(* calculate the "step" of a dimension *)
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 (* the size of dims decreases as recurses *)
            d * (get_step (ndimsleft-1) t)
          else if ndimsleft < (List.length dims)+1 then
              get_step ndimsleft (Types.Matrix(dims):Types.t)
          else
            raise (Exception "Opps hola.")
      | _ -> raise (Exception "Type of param must be matrix.")

(* For an expression that is multidimensional, iterate over all
* of the indices and evaluate each index. *)
let rec iter_matrix_expr f a expr_dim =
  iter_matrix_expr_loop f a expr_dim expr_dim [] 
and 
    iter_matrix_expr_loop f a uneval_dim expr_dim expr_indl =
  match uneval_dim with
      [] -> (f a expr_indl)
    | dim_expr :: uneval_dim' ->
        let forloop ind a f uneval_dim expr_dim expr_indl =
          (iter_matrix_expr_loop f a uneval_dim' expr_dim (ind::expr_indl))
        in
          List.fold_left 
            (fun a ind -> forloop ind a
                            f uneval_dim expr_dim expr_indl)
            a
            (nums_in_range 1 1 dim_expr)

(* Main expression evaluation code.
 * If t is a multi-element matrix than eindl should indicate which
 * element of the matrix is currently being calculated.
 * *)
let rec eval (expr,t) eindl (p:Sast.prog) (env:environ) = match expr with
    Sast.Float( v ) -> v
  | Sast.Intgr( v ) -> float_of_int v
  | Sast.Unop( op, e1 ) ->
      (
      let v1 = eval e1 eindl p env in
          match op with
            Sast.Negate -> (-. v1)
          | Sast.Sin -> (sin v1 )
          | Sast.Cos -> (cos v1 )
          | Sast.Exp -> (exp v1 )
      )
  | Sast.Binop( e1, op, e2 ) ->
      (
      let v1 = eval e1 eindl p env and v2 = eval e2 eindl p env in
          match op with
            Sast.Add -> v1 +. v2
          | Sast.Sub -> v1 -. v2
          | Sast.Mul -> v1 *. v2
          | Sast.Div -> v1 /. v2
          | Sast.Pow -> v1 ** v2
      )
  | Sast.Conv( (ed1,Types.MatrixUnknownSize), e2 ) -> assert false
  | Sast.Conv( (ed1,Types.Int), e2 ) -> assert false
  | Sast.Conv( (ed1,Types.Matrix(diml)), e2 ) ->
      let e1 = (ed1,Types.Matrix(diml)) in
      let valu = iter_matrix_expr 
                   (fun a eindl ->
                      let v1 = eval e1 eindl p env in
                      let v2 = eval e2 eindl p env in
                      a +. (v1 *. v2))
                   0.0 diml
      in 
        valu
  | Sast.Param( pnum, indl ) -> 
      (*Printf.printf "eval'ing pam %d\n" pnum;*)
      let pam = (IntMap.find pnum p.pparams) in
        (
        match pam with
            { pa_data=FloatArray(data) } -> 
              let dloc = (dataloc indl p env pam.pa_type eindl) in 
              let v = data.(dloc) in
                (*Printf.printf "Value grabbed from %d at %d is %.2f\n" 
                  pnum dloc v;*)
                v
          | { pa_data=ImageC3(data); pa_type=Types.Matrix(diml) } ->
              let i1,i2,i3 = (dataloc3 indl p env diml eindl) in 
              (float_of_int data.{i1,i2,i3}) /. 255.
          | { pa_data=ImageC3(data) } -> assert false;
          | { pa_data=Unallocated } ->
              raise( Exception ("Data must be allocated."))
        )
  | Sast.WghtRef( gsid, name, fparams ) ->
      if IntMap.mem gsid p.kerfuns = false then 
        raise(Exception (Printf.sprintf 
                           "Could not find %s() with GSID %d in p.kerfuns" name gsid ));
      let diml = (match t with Types.Matrix(diml) -> diml | _ -> assert false) in
      let kdef = (IntMap.find gsid p.kerfuns) in
      let varbindings' = fold_left3
                          (fun vbinds str num dim ->
                             StringMap.add str 
                               ((float_of_int num) -. ((float_of_int dim) +. 1.) /. 2.0 )
                               vbinds
                          )
                          StringMap.empty
                          kdef.kf_indices 
                          eindl
                          diml 
      in
      let varbindings' = List.fold_left 
                        (fun vbinds (str,num) ->
                           StringMap.add str num vbinds
                        )
                        varbindings'
                        (List.append kdef.kf_params fparams)  
      in let env' = {varbindings=varbindings'; forbindings=env.forbindings } in
        eval kdef.kf_expr eindl p env'
  | Sast.ActRef( gsid, name, exprt, pbindings ) ->
      if IntMap.mem gsid p.actfuns = false then 
        raise(Exception (Printf.sprintf "Could not find %s() with GSID %d in p.actfuns" name gsid ));
      let adef = (IntMap.find gsid p.actfuns) in
      let inval = eval exprt eindl p env in
      let varbindings' = List.fold_left 
                        (fun vbinds (str,num) ->
                           StringMap.add str num vbinds
                        )
                        (StringMap.add adef.af_invar inval StringMap.empty)
                        (List.append adef.af_params pbindings)
      in let env' = {varbindings=varbindings'; forbindings=env.forbindings } in
        eval adef.af_expr eindl p env'
  | Sast.Id(name) ->
      if StringMap.mem name env.varbindings = false then
        raise(Exception (Printf.sprintf "Could not find local variable %s" name ));
      StringMap.find name env.varbindings

  | Sast.ExtNeuron(gsid,mname,mindl,nname,nindl) ->
      if IntMap.mem gsid p.pneurons = false then
        raise(Exception "Neuron reference is unresolved - should be caught in validation.");
      let neuron = IntMap.find gsid p.pneurons in
      let nloc = dataloc nindl p env t eindl in
        neuron.v.(nloc).(1)
  | Sast.NeuronRef(gsid,nname,nindl) ->
      let neuron = IntMap.find gsid p.pneurons in
        Printf.printf 
          "Eval'ing [%s] with bindings %s\n" 
          (Printer.string_of_indices nindl)
          (Printer.string_of_forbindings env.forbindings);
      let nloc = dataloc nindl p env t eindl in
        Printf.printf "Retrieving %s[(%d)]==%f\n" nname nloc neuron.v.(nloc).(1);
        neuron.v.(nloc).(1)
  | Sast.Span(_,_,_) -> assert false (* incorrect contexts *)
  | Sast.ForExpr(_,_) -> assert false
  | Sast.SizeOf(_,_) -> assert false

and evali (expr,t) (p:Sast.prog) forbindings expr_indl = match expr with
    Sast.Intgr( v ) -> v, expr_indl
  | Sast.Unop( Negate, e1 ) ->
      let v1,eindl1 = evali e1 p forbindings expr_indl  in
        -(v1),eindl1
  | Sast.Binop( e1, op, e2 ) ->
      (
      let v1,eindl1 = evali e1 p forbindings expr_indl 
      and v2,eindl2 = evali e2 p forbindings expr_indl in
        assert (eindl1 = eindl2);
          match op with
            Sast.Add -> (v1 + v2), eindl1
          | Sast.Sub -> (v1 - v2), eindl1
          | Sast.Mul -> (v1 * v2), eindl1
          | Sast.Div -> raise( Exception "Division is not allowed in indices.")
          | Sast.Pow -> raise( Exception "Power operator is not allowed in indices.")
      )
  | Sast.Id(v) -> 
      let num = retrieve_binding v forbindings in
        num, expr_indl
  | Sast.SizeOf(v,i) ->
      ( match v with 
          Sast.Param(pnum,indl) -> 
            let pam = (IntMap.find pnum p.pparams) in 
            let diml =
              (
              match pam.pa_type with
                  Types.Matrix(diml) -> diml
                | Types.MatrixUnknownSize ->
                    raise( Exception 
                             ("size-macro cannot be used on " ^ 
                              "parameters that don't have their dimensions defined.") )
                | _ -> raise( Exception "You don't see this message, because this error can NEVER occur.")
              )
            in
              (
              try (List.nth diml (i-1)), expr_indl
              with _ -> 
                raise( Exception (Printf.sprintf 
                        "Attempted to retrieve %d dimension size of parameter that has %d dimensions"
                                    i (List.length diml)))
              )
        | Sast.Id(vref) -> raise( Exception "Local neuron references are not yet allowed in the size-macros.")
        | Sast.ExtNeuron(_,_,_,_,_) -> raise( Exception "External neuron references are not yet allowed in size-macros.")
        | _ -> raise( Exception ("The size-macro does not and likely will not support '"^(Printer.string_of_exprd v)^"'."))
      )
  | Sast.Span(e1,e2,e3) ->
      let nums = nums_of_span (Sast.Span(e1,e2,e3)) in
        Printf.printf "Span num=%d\n" (List.nth nums (List.hd expr_indl));
        (List.nth nums (List.hd expr_indl)), (List.tl expr_indl)
  | _ -> assert false (* compiler error if this is not caught before it reaches here *)

(* calculates the index of the data array from the corresponding matrix index *)
and dataloc (indl:expr_detail list) p evalenv t expr_indl = 
  match indl with
      [] ->
        ( match expr_indl with
              [] -> 0
            | dval :: eindl' -> 
                let ndims' = List.length eindl' in 
                let step = get_step ndims' t in 
                  (dval-1) * step + (dataloc indl p evalenv t eindl')
        )
    | dexpr :: indl' ->

        let (dval:int),expr_indl' = 
          evali (dexpr,Types.Int) p evalenv.forbindings expr_indl in
        let ndims' = List.length indl' in
        let step = get_step ndims' t in
          (dval-1) * step + (dataloc indl' p evalenv t expr_indl')

and dataloc3 (indl:expr_detail list) p evalenv (diml:int list) expr_indl = 
  assert (List.length indl = 3);
  assert (List.length diml = 3);
  let evalit ind expr_indl = evali (ind,Types.Int) p evalenv.forbindings expr_indl in
  let i1,expr_indl = evalit (List.hd indl) expr_indl in
  let i2,expr_indl = evalit (List.hd (List.tl indl)) expr_indl in
  let i3,expr_indl = evalit (List.hd (List.tl (List.tl indl))) expr_indl in
    (i1-1,i2-1,i3-1)
