open Ast
open Sast
open Builtin
open Printf
open Printast
open Exception

(* The map records bits length which decides during compilation. *)
module BitlenMap = Map.Make(struct
    type t = int
    let compare x y = Pervasives.compare x y
end);;

(* The assign the index to special function to avoid names conflict. *)
let f_counter = ref 0;;

module Semantic = struct

(* Given a function in type of xexpr, return its name as string. *)
let find_fun_name = function
| Xid(s, _) -> s
| Xlambda(_, _, i) -> sprintf "lambda_%d" i
| _ -> raise(Dev_Error("semantic.find_fun_name: it's not a function."))
;;

(* Return a list filled with a, whose length is l. *)
let rec build_list a l =
    if l = 1 then [a]
    else a::(build_list a (l-1));;

(* Given the basic type of vector, return a expr with that type. *)
let rec vector_to_expr t i_l =
    if List.length i_l = 1 then
        Vector_Lit(build_list (type_to_expr t) (List.hd i_l))
    else
        Vector_Lit(build_list (vector_to_expr t (List.tl i_l)) (List.hd i_l))

and string_len_n n =
    if n = 1 then
        "1"
    else
        "1" ^ string_len_n (n-1)

(* Given a type, return a expr with that type. *)
and type_to_expr = function
| Int -> Int_Lit("0")
| Bits(i) ->
    Bit_Binary_Lit(string_len_n i)
| String -> String_Lit("a")
| Vector(t, i_l) -> vector_to_expr t i_l
| Fun(_, _, _) -> raise(Dev_Error("type_to_expr.Special or Wild_Card"))
| Special(s) -> raise(Dev_Error("type_to_expr.Special or Wild_Card"))
| Wild_Card(i) -> raise(Dev_Error("type_to_expr.Special or Wild_Card"));;

(* Given the basic type of vector, return a xexpr with that type. *)
let rec vector_to_xexpr t i_l =
    if List.length i_l = 1 then
        Xvector_Lit(build_list (type_to_xexpr t) (List.hd i_l))
    else
        Xvector_Lit(build_list (vector_to_xexpr t (List.tl i_l)) (List.hd i_l))

(* Given a type, return a xexpr with that type. *)
and type_to_xexpr = function
| Int -> Xint_Lit("0")
| Bits(i) ->
    Xbit_Binary_Lit(string_of_int(int_of_float(10.0**(float_of_int (i-1)))), i)
| String -> Xstring_Lit("a")
| Vector(t, i_l) -> vector_to_xexpr t i_l
| Fun(_, _, _) -> raise(Dev_Error("type_to_xexpr.Special or Wild_Card"))
| Special(s) -> raise(Dev_Error("type_to_xexpr.Special or Wild_Card"))
| Wild_Card(i) -> raise(Dev_Error("type_to_xexpr.Special or Wild_Card"));;

(* Return the type of xexpr. *)
let rec get_type = function
| Xint_Lit(_)
| Xbin_Lit(_)
| Xhex_Lit(_) -> Int
| Xbit_Binary_Lit(_, i) -> Bits(i)
| Xbit_Hex_Lit(_, i) -> Bits(i)
| Xstring_Lit(_) -> String
| Xvector_Lit(hd::tl) ->
    let t = get_type hd in
    let len = (List.length tl) + 1 in
    begin match t with
    | Int -> Vector(Int, [len])
    | Bits(i) -> Vector(Bits(i), [len])
    | String -> Vector(String, [len])
    | Vector(t, l) -> Vector(t, len::l)
    | _ -> raise(Dev_Error("vector_lit must have concrete type")) end
| Xid(_, t) -> t
| Xidd(_, _, t) -> t
| Xvec_Dimension(_) -> Int
| Xlet(_, xe) -> get_type xe
| Xlambda(idt_l, xe, _) ->
    let t_list = List.map (fun idt -> idt.t) idt_l in
    Fun(Fix, get_type xe, t_list)
| Xmake_Vector(t, _) -> t
| Xfuncall(_, _, t, _) -> t
| Xvector_Lit([]) -> raise(Dev_Error("vector_lit cannot be zero size"))
;;

(* Confirm the uncertain type according to b_map. *)
let rec ass_type b_map = function
| Wild_Card(i) -> BitlenMap.find i b_map
| Bits(i) ->
    if i <= 0 then BitlenMap.find i b_map
    else Bits(i)
| Vector(t, l) -> Vector(ass_type b_map t, l)
| Fun(fc, t, t_l) ->
    let t' = ass_type b_map t in
    let t_l' = List.map (fun t -> ass_type b_map t) t_l in
    Fun(fc, t', t_l')
| x -> x
;;

(* Determine whether two lists of index are the same
    negative number can match with any number. *)
let rec vector_list_eq l1 l2 =
    match l1, l2 with
    | [], [] -> true
    | h1::t1, h2::t2 ->
        if h1 = 0 || h2 = 0 then
            true
        else if h1 = h2 || h1 < 0 || h2 < 0 then
            vector_list_eq t1 t2
        else
            false
    | _, _ -> false;;

(* Determine whether t1 t2 can be the same types
   b_map is used for handling the Bits(0), Bits(-1), Wild_Card(i)... type. *)
let rec compatible_type t1 t2 b_map =
    (*fprintf stderr "t1=%s, t2=%s\n" (string_type t1) (string_type t2);*)
    match (t1, t2) with
    | Wild_Card(i), Wild_Card(j) ->
        raise(Dev_Error("compare wild_card type to wild_card ?"))
    | t, Wild_Card(i) | Wild_Card(i), t ->
        if BitlenMap.mem i b_map then
            ((BitlenMap.find i b_map) = t), b_map
        else
            true, BitlenMap.add i t b_map
    | (Bits(n), Bits(m)) ->
        if (((n > 0 && m > 0) || (n <= 0 && m <= 0)) && n <> m) then
            false, b_map
        else if n == m then
            true, b_map
        else if n <= 0 then
            if BitlenMap.mem n b_map then
                ((BitlenMap.find n b_map) = Bits(m)), b_map
            else
                true, BitlenMap.add n (Bits(m)) b_map
        else
            if BitlenMap.mem m b_map then
                ((BitlenMap.find m b_map) = Bits(n)), b_map
            else
                true, BitlenMap.add m (Bits(n)) b_map
    | Fun(Indef, _, _), Fun(_, _, _) -> true, b_map
    | Fun(_, _, _), Fun(_, _, _) -> true, b_map
    | Vector(t1', l1), Vector(t2', l2) -> 
        ((fst (compatible_type t1' t2' b_map)) && (vector_list_eq l1 l2)), b_map
    | x, y ->
        x = y, b_map

(* try to evaluate some expressions at compile time *)
let easy_eval = function
| Xint_Lit(s) -> int_of_string s
| _ -> 0

(* Return the last n elements in list l *)
let rec cut_list l n =
    if n = 0 then
        ([], l)
    else
        let (l1, l2) = cut_list (List.tl l) (n-1) in
        ((List.hd l)::l1, l2)

(* Add dimension identifiers @1 ... @n to varmap *)
let rec add_vec_dim_ids n varmap =
    if n <= 0 then
        varmap
    else
        let new_varmap = VarMap.add ("@" ^ (string_of_int n)) Int varmap in
        add_vec_dim_ids (n-1) new_varmap;;

(* Check the semantic of an expr, return the xexpr and variable map.
   If there is an error, throw the exceptions. *)
let rec check_expr varmap = function
| Int_Lit(s) -> Xint_Lit(s), varmap
| Bin_Lit(s) -> Xbin_Lit(s), varmap
| Hex_Lit(s) -> Xhex_Lit(s), varmap
| Bit_Binary_Lit(s) -> Xbit_Binary_Lit(s, String.length s), varmap
| Bit_Hex_Lit(s) -> Xbit_Hex_Lit(s, 4 * (String.length s)), varmap
| String_Lit(s) -> Xstring_Lit(s), varmap
| Vector_Lit(expr_l) ->
    (* Check whether all types in expr_l are same *)
    let xexpr_l = List.map (fun e -> (fst (check_expr varmap e))) expr_l in
    let expr_t_l = List.map (fun xe -> get_type xe) xexpr_l in
    ignore (List.fold_left
        (fun t1 t2 ->
            if t1 = t2 then
                t1
            else
                raise(Invalid_Vector((string_expr 0 (Vector_Lit(expr_l))))))
        (List.hd expr_t_l) (List.tl expr_t_l));
    Xvector_Lit(xexpr_l), varmap
| Id(s) ->
    if VarMap.mem s varmap then
        Xid(s, VarMap.find s varmap), varmap
    else 
        raise(Indefined_Id(s))
| Idd(s, expr_l) ->
    let xexpr_l = List.map (fun e -> (fst (check_expr varmap e))) expr_l in
    let expr_t_l = List.map (fun xe -> get_type xe) xexpr_l in
    List.iter (fun t -> if t = Int then () else
        raise(Invalid_Ind(string_expr 0 (Idd(s, expr_l))))) expr_t_l;
    let s_type =
        if VarMap.mem s varmap then
            VarMap.find s varmap
        else
            raise(Indefined_Id(s)) in
    begin
        match s_type with
        | Vector(ctype, int_l) ->
            (* If the query dimension is higher then declaration, it fails. *)
            if List.length expr_l > List.length int_l then
                raise(Vector_Dim(s, List.length int_l))
            else
                let (l1, l2) = cut_list int_l (List.length expr_l) in
                (* Check if the index is out of bound. *)
                let valid = List.fold_left2
                    (fun b max i -> b && max > easy_eval i && easy_eval i >= 0)
                    true l1 xexpr_l in
                if valid && List.length l2 > 0 then
                    Xidd(s, xexpr_l, Vector(ctype, l2)), varmap
                else if valid then
                    Xidd(s, xexpr_l, ctype), varmap
                else
                    raise(Invalid_Ind(string_expr 0 (Idd(s, expr_l)))), varmap
        | _ -> raise(Not_Vector(s))
    end
| Vec_Dimension(i) ->
    if VarMap.mem ("@" ^ (string_of_int i)) varmap && (i >= 0) then
        Xvec_Dimension(i), varmap
    else
        raise(Make_Vec_Bound(i))
| Funcall(expr, expr_l) ->
    let xfun, varmap' = check_expr varmap expr in
    let xexpr_l, varmap = List.fold_left (fun (xl, m) e ->
        let (xe, map) = (check_expr m e) in
        xe::xl, map)
        ([], varmap') expr_l in
    let xexpr_l = List.rev xexpr_l in
    let xexpr_t_l = List.map (fun xe -> get_type xe) xexpr_l in
    let xexpr, varmap = check_expr varmap expr in
    let xexpr_t = get_type xexpr in
    begin match xexpr_t with
    | Fun(Indef, _, _) ->
        raise(Dev_Error("Semantic.Funcall unexpected case - indef"))
    | Fun(Fix, out_t, int_t_l) ->
        if List.length xexpr_t_l <> List.length int_t_l then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            let bit_map = List.fold_left2
                (fun b_map t1 t2 -> 
                    let b, new_b_map = compatible_type t1 t2 b_map in
                    if b then
                        new_b_map
                    else
                        raise(Wrong_Argu_Type(string_expr 0 (expr))))
                BitlenMap.empty xexpr_t_l int_t_l in
            let fun_name = find_fun_name xexpr in
            let exact_out_t = begin match fun_name with
            | "bits-of-int" -> Bits((easy_eval (List.hd xexpr_l)))
            | "bits-of-string" -> Bits((easy_eval (List.hd xexpr_l)))
            | "pad" -> Bits((easy_eval (List.nth xexpr_l 1)))
            | _ -> ass_type bit_map out_t end in
            Xfuncall(fun_name, xexpr_l, exact_out_t, xexpr), varmap
    | Fun(Len_Flex, out_t, [in_t]) ->
        if List.length xexpr_t_l <= 0 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            let bit_map = List.fold_left
                (fun b_map t -> 
                    let b, new_b_map = compatible_type in_t t b_map in
                    if b then
                        new_b_map
                    else
                        raise(Wrong_Argu_Type(string_expr 0 (expr))))
                BitlenMap.empty xexpr_t_l in
            let exact_out_t = ass_type bit_map out_t in
            Xfuncall(find_fun_name xexpr, xexpr_l, exact_out_t, xexpr), varmap
    | Special("group") ->
        if List.length xexpr_l <> 2 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            begin match (List.hd xexpr_t_l)
                , (List.nth xexpr_t_l 1) with
            | Bits(i), Int -> 
                let l = easy_eval (List.nth xexpr_l 1) in
                    if l > 0 then
                        let out_t = Vector(Bits(l), [(i+l-1) / l]) in
                        let fun_name = sprintf "group__%d" !f_counter in
                        f_counter := 1 + !f_counter;
                        Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
                    else
                        raise(Eval_Fail(string_expr 0 (List.nth expr_l 1)))
            | _ ->
                raise(Wrong_Argu_Type(
                    string_expr 0 (Funcall(expr, expr_l)))) end
    | Special("merge") ->
        if List.length xexpr_l <> 1 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            begin match (List.hd xexpr_t_l) with
            | Vector(Bits(i), [vec_l]) -> 
                let out_t = Bits(i * vec_l) in
                let fun_name = sprintf "merge__%d" !f_counter in
                f_counter := 1 + !f_counter;
                Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
            | _ -> raise(Wrong_Argu_Type(string_expr 0
                    (Funcall(expr, expr_l)))) end
    | Special("transpose") ->
        if List.length xexpr_l <> 1 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            begin match (List.hd xexpr_t_l) with
            | Vector(t, [d1;d2]) ->
                let out_t = Vector(t, [d2; d1]) in
                let fun_name = sprintf "transpose__%d" !f_counter in
                f_counter := 1 + !f_counter;
                Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
            | _ -> raise(Wrong_Argu_Type(string_expr 0
                    (Funcall(expr, expr_l)))) end
    | Special("map") ->
        if List.length xexpr_l <> 2 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            let xexpr_t_1 = List.hd xexpr_t_l in
            let xexpr_t_2 = List.nth xexpr_t_l 1 in
            begin match xexpr_t_1, xexpr_t_2 with
            | Fun(Fix, out_t, [in_t]), Vector(t, l) ->
                let in_fun_name = begin match List.hd xexpr_l with
                | Xid(name, _) -> name
                | Xlambda(_, _, i) -> "lambda__"
                | _ -> raise(Dev_Error("Semantic.Special.map")) end in
                let tmp_out, b_map =
                    if List.length l = 1 then
                        let valid, m =
                            compatible_type in_t t BitlenMap.empty in
                        if valid then Vector(out_t, l), m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l)))))
                    else
                        let valid, m = compatible_type
                            in_t (Vector(t, List.tl l)) BitlenMap.empty in
                        if valid then Vector(out_t, [List.hd l]), m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l))))) in
                let real_out = begin match tmp_out with
                | Vector(Vector(t, inl), outl) -> Vector(t, inl@outl)
                | t -> t end in
                let fun_name = sprintf "map__%d" !f_counter in
                f_counter := 1 + !f_counter;
                if in_fun_name = "lambda__" then
                    Xfuncall(fun_name, xexpr_l, real_out, xexpr), varmap
                else Xfuncall(fun_name,
                        [Xid(in_fun_name, Fun(Fix, out_t,
                            [ass_type b_map in_t])); (List.nth xexpr_l 1)]
                        , real_out, xexpr), varmap
            | Special(s), Vector(t, l) ->
                let in_t = if List.length l = 1
                    then t else Vector(t, List.tl l) in
                let tmp_expr = type_to_expr in_t in
                begin match check_expr varmap (Funcall(Id(s), [tmp_expr])) with
                | Xfuncall(fun_name, [in_xe_l], out_t, xexpr), _ ->
                    let in_t = get_type in_xe_l in
                    let tmp_out, b_map =
                        if List.length l = 1 then
                            let valid, m =
                                compatible_type in_t t BitlenMap.empty in
                            if valid then Vector(out_t, l), m
                            else raise(Wrong_Argu_Type(
                                (string_expr 0 (Funcall(expr, expr_l)))))
                        else
                            let valid, m = compatible_type
                                in_t (Vector(t, List.tl l)) BitlenMap.empty in
                            if valid then Vector(out_t, [List.hd l]), m
                            else raise(Wrong_Argu_Type((string_expr 0 
                                (Funcall(expr, expr_l))))) in
                    let real_out = begin match tmp_out with
                    | Vector(Vector(t, inl), outl) -> Vector(t, inl@outl)
                    | t -> t end in
                    let fun_name' = sprintf "map__%d" !f_counter in
                    f_counter := 1 + !f_counter;
                    Xfuncall(fun_name',
                        [Xid(fun_name, 
                            Fun(Fix, out_t, [in_t])); (List.nth xexpr_l 1)]
                        , real_out, xexpr), varmap
                | _ -> raise(Dev_Error("")) end
            | _ -> raise(Wrong_Argu_Type("1" ^ 
                (string_expr 0 (Funcall(expr, expr_l))))) end
    | Special("reduce") ->
        if List.length xexpr_l <> 2 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            let xexpr_t_1 = List.hd xexpr_t_l in
            let xexpr_t_2 = List.nth xexpr_t_l 1 in
            begin match xexpr_t_1, xexpr_t_2 with
            | Fun(Len_Flex, out_t, [in_t]), Vector(t, l) ->
                let b_map =
                    if List.length l = 1 then
                        let valid, m = compatible_type in_t t BitlenMap.empty in
                        if valid then m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l)))))
                    else
                        let valid, m = compatible_type
                            in_t (Vector(t, List.tl l)) BitlenMap.empty in
                        if valid then m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l))))) in
                let out_t = ass_type b_map out_t in
                let in_t = ass_type b_map in_t in
                let fun_name = sprintf "reduce__%d" !f_counter in
                f_counter := 1 + !f_counter;
                let fst_xe' =
                    begin match List.hd xexpr_l with
                    | Xid(s, Fun(Len_Flex, o, [i])) ->
                        Xid(s, Fun(Fix, out_t, [in_t;in_t]))
                    | _ ->
                        raise(Dev_Error("Semantic: special(reduce).len_flex"))
                    end in
                Xfuncall(fun_name, fst_xe'::(List.tl xexpr_l), out_t, xexpr),
                varmap
            | Fun(Fix, out_t, [in_t; in_t']), Vector(t, l) ->
                let _ =
                    if List.length l = 1 && in_t' = in_t then
                        let valid, m = compatible_type in_t t BitlenMap.empty in
                        if valid then m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l)))))
                    else
                        let valid, m = compatible_type
                            in_t (Vector(t, List.tl l)) BitlenMap.empty in
                        if valid then m
                        else raise(Wrong_Argu_Type(
                            (string_expr 0 (Funcall(expr, expr_l))))) in
                let fun_name = sprintf "reduce__%d" !f_counter in
                f_counter := 1 + !f_counter;
                Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
            | _ -> raise(Wrong_Argu_Type(
                (string_expr 0 (Funcall(expr, expr_l))))) end
    | Special("zero") ->
        if List.length xexpr_l <> 1 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            begin match (List.hd xexpr_t_l) with
            | Int -> 
                let l = easy_eval(List.hd xexpr_l) in
                let out_t = Bits(l) in
                let fun_name = sprintf "zero__%d" !f_counter in
                f_counter := 1 + !f_counter;
                Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
            | _ -> raise(Wrong_Argu_Type(
                string_expr 0 (Funcall(expr, expr_l)))) end
    | Special("rand") ->
        if List.length xexpr_l <> 1 then
            raise(Wrong_Argu_Len(string_expr 0 expr))
        else
            begin match (List.hd xexpr_t_l) with
            | Int -> 
                let l = easy_eval(List.hd xexpr_l) in
                let out_t = Bits(l) in
                let fun_name = sprintf "rand__%d" !f_counter in
                f_counter := 1 + !f_counter;
                Xfuncall(fun_name, xexpr_l, out_t, xexpr), varmap
            | _ -> raise(Wrong_Argu_Type(
                string_expr 0 (Funcall(expr, expr_l)))) end
    | x -> raise(Not_Function(string_expr 0 expr))
    end
| Lambda(idt_list, exp) ->
    let new_varmap = List.fold_left
        (fun m idt -> VarMap.add idt.id idt.t m) varmap idt_list in
    let xexpr, _ = check_expr new_varmap exp in
    let n = !f_counter in
    f_counter := 1 + !f_counter;
    Xlambda(idt_list, xexpr, n), varmap
    
| Let(let_arg_l, expr) ->
    let xlet_arg_l, varmap' = List.fold_left
        (fun (l', m) (idt, e) ->
            let xe, m = check_expr m e in
            if (get_type xe) = idt.t then
                (idt, xe)::l', VarMap.add idt.id idt.t m
            else
                raise(Bind_Wrong_Type(idt.id)))
        ([], varmap)
        let_arg_l in
    let xlet_arg_l = List.rev xlet_arg_l in
    let xexpr, t = check_expr varmap' expr in
    Xlet(xlet_arg_l, xexpr), varmap
| Make_Vector(c_type, exp) ->
    begin
    match c_type with
    | Vector(_, l) ->
        let new_varmap = add_vec_dim_ids (List.length l) varmap in
        let xexpr, m = check_expr new_varmap exp in
        Xmake_Vector(c_type, xexpr), varmap
    | _ -> raise(Wrong_Argu_Type("make-vector")) end
;;

(* Check the semantic of defvar, return xdefvar and variable map. *)
let check_defvar varmap defv =
    if VarMap.mem defv.vname.id varmap then
        raise(Bind_Twice(defv.vname.id))
    else
        let xexpr, newmap = check_expr varmap defv.vbody in
        if fst (compatible_type defv.vname.t
            (get_type xexpr) BitlenMap.empty) then
            {xvname = defv.vname; xvbody = xexpr}
                , (VarMap.add defv.vname.id defv.vname.t newmap)
        else
            raise(Bind_Wrong_Type(defv.vname.id));;

(* Check the semantic of defun, return xdefun and variable map. *)
let check_defun varmap defun =
    if VarMap.mem defun.fname.id varmap then
        raise(Bind_Twice(defun.fname.id))
    else
        let varmap' = VarMap.add defun.fname.id
            (Fun(Fix, defun.fname.t,
                (List.map (fun idt -> idt.t) defun.fargu) )) varmap in
        let infunmap = List.fold_left
            (fun m idt -> VarMap.add idt.id idt.t m)
            varmap'
            defun.fargu in
        let xexpr, infunmap = check_expr infunmap defun.fbody in
        if fst (compatible_type defun.fname.t
            (get_type xexpr) BitlenMap.empty) then
            {xfname = defun.fname; xfargu = defun.fargu; xfbody = xexpr},
            VarMap.add
                defun.fname.id
                (Fun(Fix, defun.fname.t,
                    (List.map (fun idt -> idt.t) defun.fargu)))
                varmap
        else
            raise(Defun_Wrong_Type(defun.fname.id))

let check_clip varmap = function
| Expr(expr, b) ->
    let (xe, m) = check_expr varmap expr in
    Xexpr(xe, b), m
| Defvar(defvar) ->
    let (xdefv, m) = check_defvar varmap defvar in
    Xdefvar(xdefv), m
| Defun(defun) ->
    let (xdefun, m) = check_defun varmap defun in
    Xdefun(xdefun), m
;;

let check_ast program =
    List.rev (snd (List.fold_left
        (fun (m ,p_list) c -> 
            let (xc, map) = (check_clip m c) in 
            (map, xc::p_list))
        ((add_builtin_fun VarMap.empty), [])
        program));;

end