open Ast

module NameMap = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)

exception ReturnException of int list * int NameMap.t

let rec transform mat00 mat01 mat10 mat11 = function
  | [] -> []
  | [x] -> [x]
  | x :: y :: t -> 
        (x * mat00 + y * mat01) 
        :: (x * mat10 + y * mat11) 
        :: transform mat00 mat01 mat10 mat11 t 

let explode s =
  let rec exp i l =
    if i < 0 then l else exp (i - 1) (s.[i] :: l) in
  exp (String.length s - 1) []

let implode l =
  let res = String.create (List.length l) in
  let rec imp i = function
  | [] -> res
  | c :: l -> res.[i] <- c; imp (i + 1) l in
  imp 0 l

let string_to_int_list s = List.map Char.code (explode s)@[-1]

let int_list_to_string l =
  let ascii_to_char = function
  | -1 -> ','
  | x -> Char.chr x
  in let rec char_to_string res temp = function
  | [] -> res
  | ',' :: tl -> char_to_string ((implode (List.rev temp))::res) [] tl
  | x :: tl -> char_to_string res (x::temp) tl
  in List.rev (char_to_string [] [] (List.map ascii_to_char l))

(* Main entry point: run a program *)

let run (vars, funcs) =
  (* Put function declarations in a symbol table *)
  let func_decls = List.fold_left
      (fun funcs fdecl -> NameMap.add fdecl.fname fdecl funcs)
      NameMap.empty funcs
  in

  (* Invoke a function and return an updated global symbol table *)
  let rec call fdecl actuals globals =

    (* Evaluate an expression and return (value, updated environment) *)
    let rec eval env = function
        Literal(i) -> [i], env
      | Dotop(var, op, [e]) -> 
    let l, (locals, globals) = eval env e in
      let v = List.hd l in
        if op = "nth" then
          [List.nth (NameMap.find var locals).value v], env
        else if op = "getX" then
          let index = v * 2 in
            [List.nth (NameMap.find var locals).value index], env
        else if op = "getY" then
          let index = v * 2 + 1 in
          [List.nth (NameMap.find var locals).value index], env
        else raise (Failure ("undeclared dot operation " ^ op))
      | Curve(e1, e2, e3, e4, e5, e6, e7, e8) ->
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        let l3, env = eval env e3 in
          let l4, env = eval env e4 in
            let l5, env = eval env e5 in
              let l6, env = eval env e6 in
                let l7, env = eval env e7 in
                  let l8, env = eval env e8 in
    let v1 = List.hd l1 in
      let v2 = List.hd l2 in
        let v3 = List.hd l3 in
          let v4 = List.hd l4 in
            let v5 = List.hd l5 in
              let v6 = List.hd l6 in
                let v7 = List.hd l7 in
                  let v8 = List.hd l8 in
    [v1; v2; v3; v4; v5; v6; v7; v8], env

      | Point(e1, e2) -> 
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        let v1 = List.hd l1 in
          let v2 = List.hd l2 in
          [v1; v2], env
      | Layer(cvs) ->
    List.fold_left (fun result s -> result@string_to_int_list s) [] cvs, env
      | AddToLayer(cv, ly) ->
    let locals, globals = env in 
      (NameMap.find ly locals).value@string_to_int_list cv, env
      | Noexpr -> [1], env (* must be non-zero for the for loop predicate *)
      | Id(var) ->
    let locals, globals = env in
    if NameMap.mem var locals then
      (NameMap.find var locals).value, env
    else if NameMap.mem var globals then
      [(NameMap.find var globals)], env
    else raise (Failure ("undeclared identifier " ^ var))
      | Binop(e1, op, e2) ->
    let l1, env = eval env e1 in
      let l2, env = eval env e2 in
        let v1 = List.hd l1 in
          let v2 = List.hd l2 in
    let boolean i = if i then 1 else 0 in
    (match op with
      Add -> [v1 + v2]
    | Sub -> [v1 - v2]
    | Mult -> [v1 * v2]
    | Div -> [v1 / v2]
    | Equal -> [boolean (v1 = v2)]
    | Neq -> [boolean (v1 != v2)]
    | Less -> [boolean (v1 < v2)]
    | Leq -> [boolean (v1 <= v2)]
    | Greater -> [boolean (v1 > v2)]
    | Geq -> [boolean (v1 >= v2)]), env
    | Assign(var, e) ->
    let l, (locals, globals) = eval env e in
    if NameMap.mem var locals then
      l, (NameMap.add var { t = (NameMap.find var locals).t; name = var; value = l } locals, globals)
    else if NameMap.mem var globals then
      l, (locals, NameMap.add var (List.hd l) globals)
    else raise (Failure ("undeclared identifier " ^ var))
      | Call("print", [e]) ->
    let l, env = eval env e in
      let v = List.hd l in
    print_endline (string_of_int v);
    [0], env
      | Call("draw", [e]) ->
    let l, env = eval env e in
      print_endline "draw starts";
      ignore (List.map (fun x -> print_endline (string_of_int x)) l);
      print_endline "draw ends";
    [0], env
      | Call("pause", [e]) ->
    let l, env = eval env e in
      print_endline "pause starts";
      ignore (List.map (fun x -> print_endline (string_of_int x)) l);
      print_endline "pause ends";
    [0], env
      | Call("drawLayer", [e]) ->
    let l, env = eval env e in
      let cvs = int_list_to_string l in
        print_endline "draw layer starts";
        ignore (List.map (fun x -> print_endline x) cvs);
        print_endline "draw layer ends";
    [0], env
      | Call(f, actuals) ->
    let fdecl =
      try NameMap.find f func_decls
      with Not_found -> raise (Failure ("undefined function " ^ f))
    in
    let actuals, env = List.fold_left
        (fun (actuals, env) actual ->
    let v, env = eval env actual in v :: actuals, env)
          ([], env) (List.rev actuals)
    in
    let (locals, globals) = env in
    try
      let globals = call fdecl actuals globals
      in [0], (locals, globals)
    with ReturnException(v, globals) -> v, (locals, globals)
    in

    (* Execute a statement and return an updated environment *)
    let rec exec env = function
  Block(stmts) -> List.fold_left exec env stmts
      | Expr(e) -> let _, env = eval env e in env
      | If(e, s1, s2) ->
    let l, env = eval env e in
      let v = List.hd l in
    exec env (if v != 0 then s1 else s2)
      | While(e, s) ->
    let rec loop env =
      let l, env = eval env e in
        let v = List.hd l in
      if v != 0 then loop (exec env s) else env
    in loop env
      | For(e1, e2, e3, s) ->
    let _, env = eval env e1 in
    let rec loop env =
      let l, env = eval env e2 in
        let v = List.hd l in
      if v != 0 then
        let _, env = eval (exec env s) e3 in
        loop env
      else
        env
    in loop env
      | Return(e) ->
    let v, (locals, globals) = eval env e in
    raise (ReturnException(v, globals))
    in

    (* Enter the function: bind actual values to formal arguments *)
    let locals =
      try List.fold_left2
    (fun locals formal actual -> NameMap.add formal.name { t = formal.t; name = formal.name; value = actual } locals)
    NameMap.empty fdecl.formals actuals
      with Invalid_argument(_) ->
  raise (Failure ("wrong number of arguments passed to " ^ fdecl.fname))
    in
    (* Initialize local variables to 0 *)
    let locals = List.fold_left
  (fun locals local -> NameMap.add local.name { t = local.t; name = local.name; value = local.value } locals) locals fdecl.locals
    in
    (* Execute each statement in sequence, return updated global symbol table *)
    snd (List.fold_left exec (locals, globals) fdecl.body)

  (* Run a program: initialize global variables to 0, find and run "main" *)
  in let globals = List.fold_left
      (fun globals vdecl -> NameMap.add vdecl 0 globals) NameMap.empty (List.map (fun x -> x.name) vars)
  in try
    call (NameMap.find "main" func_decls) [] globals
  with Not_found -> raise (Failure ("did not find the main() function"))
