open Absyn

exception TODO
exception RuntimeError of string

(**** Easy control for dynamic vs. static scoping. ******)
type scoping = Dynamic | Static 
let scoping = Dynamic

type value = 
    Num of int 
  | ProcVal of env * var list * exp 
  | Bool of bool 
and  env = Extend of var list * value array * env 
           | ExtendRec of var list * var list list * exp list * env
           | Empty
               
let rec print_value (v:value) : unit = 
  match v with 
      Num i -> Format.print_int i;
    | Bool b -> if b then Format.print_string "true" else 
        Format.print_string "false"
    | ProcVal (env,v,str) -> Format.print_string "<closure>"
        
let closure (ids:var list) (body:exp) (env:env) : value = 
  ProcVal (env, ids, body)
    
(* ------- Environments --------------- *)
    
let list_find_position (x:'a) (l:'a list) :int option = 
  let rec helper l i = 
    match l with 
        [] -> None
      | (hd::tl) -> if x = hd then Some i else 
          helper tl (i+1)
  in
    helper l 0

let extend_env (vars:var list) (vals:value list) (old_env:env) : env = 
  Extend(vars, Array.of_list vals, old_env)

let extend_env_rec (proc_names:var list) (idss:var list list) (bodies:exp list) (old_env:env) : env = 
  (*** Change from problem 5, don't used ExtendRec *****)
  let num_procs = List.length proc_names             in 
  let array     = Array.make num_procs (Num 0)       in
  let new_env   = Extend (proc_names, array,old_env) in 
  let rec loop idss bodies n : unit = 
    match idss, bodies with 
        (ids::is), (body::bs)  -> 
          (array.(n) <- closure ids body new_env;
           loop is bs (n+1))
      | [], [] -> () 
      | _,_ -> failwith "BUG: Cannot happen" in
    begin
      loop idss bodies 0;
      new_env
    end 
      

let rec apply_env (env:env) (sym:var) :value = 
  match env with 
      Empty -> raise (RuntimeError ("Unbound variable " ^ sym))
    | Extend (vars, vals, old_env) -> 
        (match list_find_position sym vars with 
             Some pos -> vals.(pos)
           | None -> apply_env old_env sym)
    | ExtendRec (proc_names, idss, bodies, old_env) ->
        match (list_find_position sym proc_names) with 
            Some pos -> 
              closure (List.nth idss pos) (List.nth bodies pos) env
          | None -> apply_env old_env sym


(* ----- Booleans ------ *)

(**** Changes for part 2 below *******)
let is_true (v:value) :bool = 
  match v with 
      Bool b -> b
    | _ -> raise (RuntimeError "Invalid value used as bool")
        
let true_value  :value = Bool true
let false_value :value = Bool false
                           
(* ----------- Primitives ----------- *)

let rec apply_prim (p:prim) (args:value list) = 
  match p,args with 
      Plus,  [Num v1;Num v2] -> Num (v1 + v2)
    | Times, [Num v1;Num v2] -> Num (v1 * v2)
    | Minus, [Num v1;Num v2] -> Num (v1 - v2)
    | And,   [v1;v2] ->
        if (is_true v1 && is_true v2)
        then true_value
        else false_value
    | Or,    [v1;v2] ->
        if (is_true v1 || is_true v2)
        then true_value
        else false_value
    | Not,   [v1] ->
        if (is_true v1) 
        then false_value 
        else true_value
    | Add1, [Num v1] -> Num (v1 + 1)
    | Sub1, [Num v1] -> Num (v1 - 1)
        (****** New branches added for problem 1 **********************)
    | Zero, [Num v1] -> if v1 = 0 then true_value else false_value
    | LT, [Num v1; Num v2] -> if v1 < v2 then true_value else false_value
    | GT, [Num v1; Num v2] -> if v1 > v2 then true_value else false_value
    | Equal, [Num v1;Num v2] -> 
        if (v1 = v2) then true_value  else false_value
    | _, _ -> raise (RuntimeError "Unimplemented  primitive or invalid arguments")

(* ------------ Evaluator ---------------------- *)

(******* Addition for problem 3 & 4 ********
  - new argument, current_env
  - check length of ids and args
  - extend current_env instead of env stored in closure
*)
let rec apply_procval (v:value) (args:value list) (current_env:env) : value = 
  match v with 
      ProcVal (env, ids, body) -> 
        if (List.length ids = List.length args) then 
          eval body (extend_env ids args (if scoping = Dynamic then  current_env else env))
        else 
          raise (RuntimeError "Wrong number of arguments to procedure")
    | _ -> raise (RuntimeError "Attempt to apply non-procedure")

and  eval (exp:exp) (env:env) :value = 
  match exp with 
      LitInt i -> Num i
    | LitBool b -> 
        if b then true_value 
        else false_value
    | Var v -> apply_env env v
    | Prim (p, rands) -> 
        let args = eval_rands rands env in
          apply_prim p args
    | If (test_exp, true_exp, false_exp) -> 
        if (is_true (eval test_exp env))
        then (eval true_exp env)
        else (eval false_exp env)
    | Let (vars,rands,body) -> 
        let args = eval_rands rands env in 
          eval body (extend_env vars args env)
    | Lambda (ids,body) ->              
        closure ids body env
    | App (rator,rands) -> 
        let proc = eval rator env in 
        let args = eval_rands rands env in 
          (**** change for prob 4 -- new argument ***)
          apply_procval proc args env

    | Letrec (proc_names, idss, bodies, letrec_body) -> 
        eval letrec_body (extend_env_rec proc_names idss bodies env)
          
and eval_rands (rands:exp list) (env:env): value list = 
  List.map (fun x -> eval x env) rands