open Ast
open Sast
open Types
open Params

exception Exception of string

type transl_env = {
    tparams: Sast.param_def IntMap.t;
    tactfuns: Sast.actfun_def IntMap.t;
    tkerfuns: Sast.kerfun_def IntMap.t;
    (*moddefs: mod_def StringMap.t; <- only needs to be used when module
    * declarations are allowed *)

    (* actualized module definitions *)
    tmods: Sast.mod_def StringMap.t;
}
let intMapCount m = IntMap.fold (fun k d a -> a+1) m 0
let stringMapCount m = StringMap.fold (fun k d a -> a+1) m 0

let neuron_count { m_inputs=n1; m_outputs=n2; m_neurons=n3 } =
  (List.length n1) + (List.length n2) + (List.length n3)

(* used to calculate the new "global symbol ids" or gsid *)
let symbol_count env =
  (intMapCount env.tparams) + (intMapCount env.tactfuns)
  + (StringMap.fold
       (fun k mdule count -> count + (neuron_count mdule))
       env.tmods
       0)

let create_environ (program:Ast.program) =
    let params = Params.create program in
       assert( IntMap.is_empty params = false );
    { tparams=params; 
      tactfuns=IntMap.empty; tkerfuns=IntMap.empty;
      tmods=StringMap.empty; }

let ast2str = function
   Ast.Float(_) -> "Ast.Float"
 | Ast.Binop(_,_,_) -> "Ast.Binop"
 | Ast.Intgr(_) -> "Ast.Intgr"
 | Ast.SizeOf(_,_)->"Ast.SizeOf"
 | Ast.Id(_)->"Ast.Id"
 | Ast.Negate(_)->"Ast.Negate"
 | Ast.NeuronRef(_,_)->"Ast.NeuronRef"
 | Ast.ExternalRef(_)->"Ast.ExternalRef"
 | Ast.ActRef(_,_,_)->"Ast.ActRef"
 | Ast.WghtRef(_,_)->"Ast.WghtRef"
 | Ast.Conv(_)->"Ast.Conv"
 | Ast.Exp(_)->"Ast.Exp"
 | Ast.Sin(_)->"Ast.Sin"
 | Ast.Cos(_)->"Ast.Cos"
 | Ast.Span(_,_,_)->"Ast.Span"
 | Ast.ForExpr(_,_)->"Ast.ForExpr"

let rec translate expr env =
match expr with 
    Ast.Float(v) -> Sast.Float(v), Types.Matrix([1]) 
  | Ast.Intgr(v) -> Sast.Intgr(v), Types.Int
  | Ast.Id(name) -> Sast.Id(name), Types.MatrixUnknownSize (* local function variable  or neuronref*)
  | Ast.Exp(e1) -> 
      let e1' = translate e1 env in 
        Sast.Unop( Sast.Exp, e1' ), Types.MatrixUnknownSize
  | Ast.Sin(e1) -> 
      let e1' = translate e1 env in 
        Sast.Unop( Sast.Sin, e1' ), Types.MatrixUnknownSize
  | Ast.Cos(e1) -> 
      let e1' = translate e1 env in 
        Sast.Unop( Sast.Cos, e1' ), Types.MatrixUnknownSize
  | Ast.Negate(e1) -> 
      let e1' = translate e1 env in 
        Sast.Unop( Sast.Negate, e1' ), Types.MatrixUnknownSize
  | Ast.Binop(e1,op,e2) -> 
      let e1' = translate e1 env in
      let e2' = translate e2 env in
      let op' = match op with 
           Ast.Add -> Sast.Add | Ast.Sub -> Sast.Sub | Ast.Mul -> Sast.Mul
         | Ast.Div -> Sast.Div | Ast.Pow -> Sast.Pow
      in 
         Sast.Binop(e1',op',e2'), Types.MatrixUnknownSize 
  | Ast.Conv(e1,e2) -> (* Should allow multiple Conv, but type doesn't *)
      let e1' = translate e1 env in
      let e2' = translate e2 env in (* e2 should be WghtRef *)
        Sast.Conv(e1',e2'), Types.MatrixUnknownSize
  | Ast.ExternalRef(extref) ->
      let _,expr,t = translate_extref extref env OutputNeuron in
        expr,t
  | Ast.WghtRef( name, fparams ) ->
      let gsid = IntMap.fold
                   (fun k {kf_name=kf_name} a ->
                      Printf.printf "trying to match %s %s\n" kf_name name;
                      if kf_name=name then
                        k
                      else a)
                   env.tkerfuns
                   (-1)
      in
        if gsid = -1 then 
          raise 
            (Exception 
               (Printf.sprintf "Kernel function %s is undefined." name));
        Sast.WghtRef(gsid, name, fparams), Types.MatrixUnknownSize

  | Ast.ActRef(name,expr,fparams) ->
      let expr' = translate expr env in
      let gsid = IntMap.fold 
                   (fun k {af_name=af_name} a ->
                      if af_name=name then
                        k
                      else a
                   )
                   env.tactfuns 
                   (-1)
      in
        if gsid = -1 then 
          raise 
            ( Exception (Printf.sprintf "Activation function %s is undefined." name));
        Sast.ActRef(gsid, name, expr', fparams), Types.MatrixUnknownSize
               
  | Ast.ForExpr(_,_)-> 
      raise ( Exception ("Error: for expression in invalid environment. "))

  | Ast.NeuronRef(name,indl)->
      let (indl':Sast.expr_detail list) = List.fold_left
                    (fun li i -> (translate_ind i env) :: li)
                    [] indl
      in 
        Sast.NeuronRef(-1,name,indl'), Types.MatrixUnknownSize
  | Ast.Span(_,_,_) ->
      raise ( Exception "Span is not allowed here!")
  | Ast.SizeOf(_,_) ->
      raise ( Exception "SizeOf is not allowed here!")

and translate_ind exprd env =  (* Inside index *)
  match exprd with 
      Ast.Intgr(v) -> Sast.Intgr(v)
    | Ast.Id(name) -> Sast.Id(name) (* must refer to for-loop var *)
    | Ast.Negate(e1) ->
        let e1' = translate_ind e1 env in
          Sast.Unop(Sast.Negate, (e1',Types.Int))
    | Ast.Binop(e1,op,e2) -> 
      let e1' = translate_ind e1 env in
      let e2' = translate_ind e2 env in
      let op' = match op with 
           Ast.Add -> Sast.Add | Ast.Sub -> Sast.Sub | Ast.Mul -> Sast.Mul
        | Ast.Div -> raise( Exception "Error division is not allowed in indices.")
        | Ast.Pow -> raise( Exception "Power operators are not allowed in indices.")
      in 
         Sast.Binop( (e1',Types.Int),op',(e2',Types.Int))
    | Ast.SizeOf(exprl,dim) ->
        let e',t = translate exprl env in 
          Sast.SizeOf(e',dim)
    | Ast.Span(e1,step,e2) -> 
        let e1' = translate_ind e1 env in
        let step' = translate_ind step env in
        let e2' = translate_ind e2 env 
        in 
          Sast.Span(e1',step',e2')
      (* the following errors should have been caught as syntax errors *)
    | Ast.ForExpr(_,_) -> assert false
    | Ast.Cos(_) -> assert false
    | Ast.Sin(_) -> assert false
    | Ast.Exp(_) -> assert false
    | Ast.Conv(_,_) -> assert false
    | Ast.WghtRef(_,_) -> assert false
    | Ast.ActRef(_,_,_) -> assert false
    | Ast.ExternalRef(_) -> assert false
    | Ast.NeuronRef(_,_) -> assert false
    | Ast.Float(_) -> assert false

and translate_extref eref env reqd_nt = match eref with
    Ast.Param(pnum,diml) -> 
      let diml' = List.rev 
                    (List.fold_left 
                       (fun li dim -> (translate_ind dim env)::li ) 
                       [] diml)
      in
          pnum,Sast.Param(pnum,diml'), Types.MatrixUnknownSize
  | Ast.ExtNeuron(mname,mindl,nname,nindl) ->
      Printf.printf "Ast.ExtNeuron(%s,...,%s,...)\n" mname nname;
      if StringMap.mem mname env.tmods = false then
        raise (Exception 
                 ("Could not find "^mname^" module."));
      let gsid =
          let neurons = 
            let tmod = StringMap.find mname env.tmods in 
              List.append tmod.m_inputs (List.append tmod.m_outputs tmod.m_neurons)
          in
            List.fold_left 
              (fun a (name,gsid,_,nt) -> 
                 if name = nname then
                   ( 
                     if nt<>reqd_nt then match nt with 
                         InterNeuron -> raise (Exception 
                            (Printf.sprintf 
                               "Attempted to reference interneuron '%s' outside of module."
                               name 
                            )) 
                       | InputNeuron -> raise (Exception
                            (Printf.sprintf
                               "Attempted to receive action potentials from input neuron '%s' outside of module."
                               name
                            ))
                       | OutputNeuron -> raise (Exception
                            (Printf.sprintf
                               "Attempted to send action potentials to output neuron '%s' outside of module."
                               name
                            ))
                     else gsid
                   )
                 else 
                   a
              ) 
              (-1)
              neurons
      in
        if gsid= -1 then raise (Exception (Printf.sprintf "Could not find the neuron %s in %s." nname mname));
      let translate' indl =
        List.fold_left 
          (fun indl indx -> (translate_ind indx env) :: indl )
          [] indl
      in
        gsid,
       Sast.ExtNeuron( gsid, mname,
                      (translate' mindl),
                      nname,
                      (translate' nindl)),
       Types.MatrixUnknownSize

let translate_for expr env = match expr with
    Ast.ForExpr(varname,e1)->
      let e1' = translate_ind e1 env 
      in 
        Sast.ForExpr(varname,e1'), env
  | Ast.ForExpr(_,e) -> 
      raise ( Exception ("Error: invalid format for for-macro " ^ 
                         (ast2str e)^"."))
  | e -> raise( Exception ("Error: expected for-expression, found "^ (ast2str e)^".")) 

let translate_syn synap env = 
 let expr' = translate synap.s_expr env in
 let (sfor':Sast.expr_detail list),env' = List.fold_left
                    (fun (li,env') sfor -> 
                       let ((fexpr:Sast.expr_detail),env') = 
                         translate_for sfor env in fexpr::li,env' )
                    ([], env)
                    (synap.s_for: Ast.expr list)
 in
   (expr', sfor')

(* Performs semantic analysis AST->SAST for synapses, some sizes may still be
 * undetermined *)
let translate_s synap env = 
 let (gsid,extref',_) = translate_extref synap.s_dest env InputNeuron in
 let (expr',sfor') = translate_syn synap.s_syn env 
 in
   gsid,{ sdest=extref'; sexpr=expr'; sfor=sfor'},env


(* Dumbed-down version of Eval.evali that cannot use SizeOf or Id *)
let rec evalc = function
    Sast.Intgr( v ) -> v
  | Sast.Binop( (e1,_), op, (e2,_) ) ->
      (
      let v1 = evalc e1 and v2 = evalc e2 in
          match op with
            Sast.Add -> v1 + v2
          | Sast.Sub -> v1 - v2
          | Sast.Mul -> v1 * v2
          | Sast.Div -> raise( Exception "Division is not allowed in indices.")
          | Sast.Pow -> raise( Exception "Power operator is not allowed in indices.")
      )
  | e -> raise(Exception ( Printer.string_of_exprd e ^ " can not be used inside the dimension definition of a neuron or module."))

let match_local_neuron idname allneurons =
  List.fold_left 
    (fun a (nname,gsid,t,nt) -> 
       if nname=idname then NeuroDefMatch( (nname,gsid,t,nt) ) 
       else a) 
    NoNeuroDefMatch
    allneurons

let rec convert_ids_to_neurorefs (exprd,t) forl allneurons = 
  (convert_ids_to_neurorefs_d exprd forl allneurons),t 
and convert_ids_to_neurorefs_d exprd forl allneurons = match exprd with
    Sast.Id(idname) ->
      Printf.printf "Trying to match '%s' with a neuron..." idname;
      let mneuron = match_local_neuron idname allneurons
      in 
        (
        match mneuron with
          
          (* TO-DO: should make sure that there are no conflicts between neuron
          * reference and a for-var  *)
          Sast.NeuroDefMatch( (nname,gsid,t,nt) )->
            Printf.printf "success matching to %d.\n" gsid;
                Sast.NeuronRef(gsid,nname,[])
        | NoNeuroDefMatch -> 
            Printf.printf "failed.\n";
            Sast.Id(idname)
        )
  | Sast.NeuronRef( -1, nname, indl ) -> 
      let mneuron = match_local_neuron nname allneurons
      in
        (
          match mneuron with
              Sast.NeuroDefMatch( (nname,gsid,t,nt) )->
                Sast.NeuronRef(gsid,nname,indl)
            | NoNeuroDefMatch ->
                raise (Exception (Printf.sprintf "Unable to resolve local neuron reference `%s'." nname))
        )

  | Sast.NeuronRef( a, b, c ) -> Sast.NeuronRef( a, b, c )
  | Sast.Binop(x1,op,x2) ->
      let x1' = convert_ids_to_neurorefs x1 forl allneurons in
      let x2' = convert_ids_to_neurorefs x2 forl allneurons in
        Sast.Binop(x1',op,x2')
  | Sast.Unop(op,x) ->
      Sast.Unop( op, convert_ids_to_neurorefs x forl allneurons )
  | Sast.Float(v) -> Sast.Float(v)
  | Sast.Intgr(v) -> Sast.Intgr(v)
  | Sast.SizeOf(exprd, d) -> Sast.SizeOf( (convert_ids_to_neurorefs_d exprd forl allneurons), d )
  | Sast.ActRef(a,b,x,c) ->
      let x' = convert_ids_to_neurorefs x forl allneurons in
        Sast.ActRef(a,b,x',c)
  | Sast.WghtRef( gsid, n, fparams ) ->
      Sast.WghtRef( gsid, n, fparams )
  | Sast.Conv( x1, x2 ) ->
      let x1' = convert_ids_to_neurorefs x1 forl allneurons in
      let x2' = convert_ids_to_neurorefs x2 forl allneurons in
        Sast.Conv( x1', x2' )
  | Sast.Span( a, b, c) -> Sast.Span( a, b, c ) (* neurorefs can't be in spans - except through sizeof *)
  | Sast.ForExpr( name, exprd ) -> Sast.ForExpr( name, exprd)(* except through sizeof *)
  | Sast.ExtNeuron( a, b, c, d, e ) ->
      Sast.ExtNeuron( a, b, c, d, e )
  | Sast.Param( pnum, indl ) -> Sast.Param( pnum, indl)



let translate_mdef mdef env synmap = 
  let synmap,mdef' =
    let symc = symbol_count env in
    let conv_neurodefs symc pneurons nt =  (* converts Ast neurons to Sast neurons *)
      List.fold_left 
        (fun (inplist,symc) (name,diml) ->
           let diml' =
             List.fold_left
               (fun diml' expr ->
                 (evalc (translate_ind expr env)) :: diml')
               []
               diml
           in
             (name,symc+1,Types.Matrix(diml'),nt) :: inplist,(symc+1))
        ([],symc)
        pneurons
    in
    let inputs',symc = conv_neurodefs symc mdef.mod_inputs InputNeuron in
    let outputs',symc = conv_neurodefs symc mdef.mod_outputs OutputNeuron in
    let (synaps',env) = List.fold_left  (* external references *)
                    (fun (synmap,env) syn -> 
                       let (gsid,syn,env) = translate_s syn env 
                       in (IntMap.add gsid syn synmap), env )
                    (synmap,env) 
                    mdef.mod_synaps 
    in
    let (neurons',synaps'',env,symc) 
          = List.fold_left   (* local references *)
              (fun (neurli,synmap,env,symc) n -> 
                 let gsid,nt = List.fold_left
                              (fun a ((name,nid,_,nt):neuro_def) ->
                                 if name = n.neuro_name then
                                   nid,nt
                                 else a
                              )
                              (-1,InputNeuron)
                              outputs'
                 in
                 let (symc',gsid) = if gsid = -1 then (symc+1,symc+1) else (symc,gsid) in
                 let (sexpr,sfor) = translate_syn n.neuro_syn env in 
                 let nindl = 
                   List.fold_left 
                     (fun li ind -> (translate_ind ind env) :: li)
                     [] n.neuro_ind
                 in
                 let (n':Sast.neuro_def) = n.neuro_name,gsid,Types.MatrixUnknownSize,nt
                 in
                   n'::neurli,
                   (IntMap.add gsid { sdest=Sast.NeuronRef(gsid,n.neuro_name,nindl); 
                    sexpr=sexpr; sfor=sfor} synmap),env,symc')
              ([],synaps',env,symc)
              mdef.mod_neurons
    in
    let synaps''' = 
      let allneurons = List.append inputs' (List.append outputs' neurons')
        in 
        IntMap.map (fun {sdest=sdest;sexpr=sexpr;sfor=sfor} -> 
                      let sexpr' = (convert_ids_to_neurorefs sexpr sfor allneurons)
                      in
                        print_endline (Printer.string_of_expr sexpr');
                        { sdest=sdest; sexpr=sexpr'; sfor=sfor }) synaps''
    in
      IntMap.iter (fun k d -> print_endline (Printer.string_of_synap k d) ) synaps''';
      synaps''',
      { m_name=mdef.mod_name; m_inputs=inputs'; m_outputs=outputs'; m_neurons=neurons';
        m_synaps=IntMap.fold (fun k d li -> d :: li ) synaps''' [] }
  in let tmods' = 
    StringMap.add mdef.mod_name mdef' env.tmods
  in
    synmap,{ tmods=tmods'; tparams=env.tparams; 
             tactfuns=env.tactfuns; tkerfuns=env.tkerfuns }

(* translates actdef from Ast.program to transl_env  *)
let translate_adef actdef env = match actdef with
    { act_name = name; act_params=fparams; 
      act_local = invar; act_expr=expr }
    ->
      Printf.printf "Translating activation function %s.\n" name;
      let expr' = translate expr env in
      let gsid = (intMapCount env.tparams) + (intMapCount env.tactfuns) + 1 in
      let adef' = { af_gsid = gsid; af_name=name; af_invar=invar; 
          af_params=fparams; af_expr=expr' }
      in
      let adefs' = IntMap.add gsid adef' env.tactfuns in
        Printf.printf "Current activation functions:\n";
        IntMap.iter (fun gsid adef -> 
                       print_endline (Printer.string_of_afundef gsid adef)) 
          adefs';
        { tparams=env.tparams; tactfuns=adefs'; 
          tkerfuns=env.tkerfuns;
          tmods=env.tmods }

(* translates kerneldef from Ast.program to transl_env  *)
let translate_kdef kerdef env = match kerdef with
    { wght_name = name; wght_ind = indices;
      wght_params = fparams; wght_expr = expr }
    ->
      let expr' = translate expr env in
      let gsid = (intMapCount env.tparams) + 
                 (intMapCount env.tactfuns) +
                 (intMapCount env.tkerfuns) + 1
      in
      let kdef' = { kf_gsid=gsid; kf_name=name; kf_indices=indices;
                    kf_params=fparams; kf_expr=expr' }
      in
      let kdefs' = IntMap.add gsid kdef' env.tkerfuns 
      in
        Printf.printf "kern %s\n" name;
          
        { tparams=env.tparams; tactfuns=env.tactfuns;
          tmods=env.tmods;
          tkerfuns=kdefs' }


let translate_p (program:Ast.program) env=
  let env = List.fold_left
               (fun env adef ->
                  let env' = translate_adef adef env in env')
               env
               (List.rev program.p_adef)
  in
  let env = List.fold_left
               (fun env kdef ->
                  let env' = translate_kdef kdef env in env')
               env
               program.p_wdef
  in
    Printf.printf "There are %d kernel definitions.\n" 
      (intMapCount env.tkerfuns);
  let msynaps,env = List.fold_left
              (fun (synm,env) mdef -> translate_mdef mdef env synm )
              (IntMap.empty,env)
              program.p_mdef
  in
  let synapsel,env = List.fold_left 
                   (fun (m,env) syn -> 
                        let (gsid,syn,env) = translate_s syn env 
                        in
                         if IntMap.mem gsid m then 
                           raise( Exception (Printf.sprintf "A synapse is redefining the symbol with ID %d." gsid ));
                          (IntMap.add gsid syn m), env )
                   (msynaps,env)
                   program.p_synap 
  in
  let neurons = 
    StringMap.fold 
      (fun mname mdef a -> match mdef with 
           { m_inputs=inputs; m_outputs=outputs; m_neurons=inters } ->
             let addneurons (nlist:Sast.neuro_def list) nmap =
               List.fold_left
                 (fun nmap (name,gsid,t,nt) -> 
                    Printf.printf "Creating neuron %s.%s with gsid %d\n" mname name gsid;
                    IntMap.add 
                      gsid 
                      { n_mname=mname; n_name=name; 
                        t=t; v=[| [| 0.; 0. |] |] }
                      nmap
                 )
                 nmap
                 nlist
             in
               (addneurons inputs (addneurons outputs (addneurons inters IntMap.empty)))) 
      env.tmods
      IntMap.empty
   in { pparams=env.tparams; actfuns=env.tactfuns; kerfuns=env.tkerfuns;
        syn=synapsel; pneurons=neurons }

let program p =
  let env = create_environ p in
      assert( (IntMap.is_empty env.tparams) = false );
  let p = translate_p p env
  in
      IntMap.iter 
        (fun k pam -> print_endline (Printer.string_of_param pam.pa_num pam)) 
        p.pparams; 
      p
