(* If you would like, write how many hours you spent on this homework:

  XXX

*)

module StringMap = Map.Make(String)

type value = int
type op = | Add | Mul | EQ | Sub | Gt

let eval_binop (o : op) (v1 : value) (v2 : value) : value = 
    match o with 
    | Add -> v1 + v2
    | Mul -> v1 * v2
    | Sub -> v1 - v2
    | EQ -> if v1 = v2 then 1 else 0
    | Gt -> if v1 > v2 then 1 else 0

module StackLang = struct

  type instr = 
    | Push of value
    | PrimOp of op
    | Get of string
    | Set of string
    | JumpIfNotOne of int
    | Jump of int

  type config = {
    pc : int;
    ins : instr list;
    stack : value list;
    env : value StringMap.t
  }

  let step (cfg : config) : config =
    match List.nth cfg.ins cfg.pc with 
    | Push v -> { cfg with pc = cfg.pc + 1; stack = v :: cfg.stack }
    | Set s -> 
      (match cfg.stack with 
        | [] -> failwith "Stack empty"
        | v::stack' -> { cfg with pc = cfg.pc + 1; stack = stack'; env = StringMap.add s v cfg.env }
      )
    | Get s -> 
      (match StringMap.find_opt s cfg.env with
      | None -> failwith "Unknown variable"
      | Some v -> { cfg with pc = cfg.pc + 1; stack = v :: cfg.stack }
      )
    | PrimOp op -> 
      (match cfg.stack with 
        | v1 :: v2 :: stack' -> 
          (* Important: ordering of v1 and v2 is flipped. This is on purpose. *)
          let v = eval_binop op v2 v1 in { cfg with pc = cfg.pc + 1; stack = v :: stack' }
        | _ -> failwith "Not enough elements on stack."
      )
    | Jump n -> { cfg with pc = cfg.pc + n }
    | JumpIfNotOne n -> 
        (match cfg.stack with
         | v :: stack' -> 
             if v <> 1 
             then { cfg with pc = cfg.pc + n; stack = stack' } 
             else { cfg with pc = cfg.pc + 1; stack = stack' }
         | _ -> failwith "Not enough elements on stack"
        )

    type program = instr list

    let rec loop (cfg : config) : config =
        if 0 <= cfg.pc && cfg.pc < List.length cfg.ins then 
          loop (step cfg)
        else cfg

    let init_config p = { pc = 0; ins = p; stack = []; env = StringMap.empty }

    let execute (p : program) : value StringMap.t =
      (loop (init_config p)).env

    let get_pc_list (p : program) : int list * value StringMap.t =
        let c = init_config p in 
        let history = ref [c.pc] in 
        let rec go cfg = 
          if 0 <= cfg.pc && cfg.pc < List.length cfg.ins then
            let cfg' = step cfg in 
            history := !history @ [cfg'.pc];
            go cfg' 
          else cfg
        in 
        let cfg' = go c in 
        (!history, cfg'.env)
end


module MiniPython = struct
  type expr = | Value of value | Var of string | App of op * expr list

  type cmd = | Assign of string * expr | Ite of expr * cmd * cmd | Seq of cmd * cmd | While of expr * cmd | Skip


  let rec eval_expr (m : value StringMap.t) (e: expr) : value = 
    match e with 
    | Value v -> v
    | Var x -> StringMap.find x m
    | App (o, es) -> 
        (match es with 
          | [e1; e2] -> eval_binop o (eval_expr m e1) (eval_expr m e2)
          | _ -> failwith "Got wrong number of arguments"
        )
      
  let rec eval_cmd (m : value StringMap.t) (c : cmd) : value StringMap.t = 
    match c with 
    | Assign (x, e) -> StringMap.add x (eval_expr m e) m
    | Ite (e, c1, c2) -> if eval_expr m e = 1 then eval_cmd m c1 else eval_cmd m c2
    | Seq (c1, c2) -> let m' = eval_cmd m c1 in eval_cmd m' c2
    | While (e, c) -> 
        if eval_expr m e = 1 then 
            let m' = eval_cmd m c in 
            eval_cmd m' (While (e, c))
        else 
          m
    | Skip -> m
end

let rec compile_expr (e : MiniPython.expr) : StackLang.program = 
  failwith "TODO"

let rec compile_cmd (c : MiniPython.cmd) : StackLang.program = 
  failwith "TODO"

(*
    Below, we have three example programs, along with a way of testing the return values of each when run under both MiniPythong and StackLang.
*)

let if_cmd : MiniPython.cmd = 
    (*
      x := 0;
      y := 1;
      result := 0;
      if x == y {
        result := 1;
      }
      else {
        result := 2;
      }
  *)
  Seq (
    Assign ("x", Value 0),
    Seq (
      Assign ("y", Value 1),
      Seq (
        Assign ("result", Value 0),
        Ite (
          App (EQ, [Var "x"; Var "y"]),
          Assign ("result", Value 1),
          Assign ("result", Value 2)
        )
      )
    )
  )

let while_cmd : MiniPython.cmd = 
    (* 
      n := 2;
      while (n > 0) {
        n := n - 1;
      }
    *)
  Seq (
    Assign ("n", Value 2),
    While (
      App (Gt, [Var "n"; Value 0]),
      Seq (
        Assign ("n", App (Sub, [Var "n"; Value 1])),
        Skip
      )
    )
  )


let fib_cmd n : MiniPython.cmd =
    (*
      n := (the number we pass in);
      a := 0;
      b := 1;
      sum := 0; 
      if n == 0 then
          skip
      else  if n == 1 then 
          sum := 0
      else 
      while (n > 1) {
        sum := a + b;
        n := n - 1; 
        a := b; 
        b := sum;
      }
    *)
  Seq (
    Assign ("n", Value n),
    Seq (
      Assign ("a", Value 0),
      Seq (
        Assign ("b", Value 1),
        Seq (
          Assign ("sum", Value 0),
          Ite (
            App (EQ, [Var "n"; Value 0]),
            Skip,
            Ite (
              App (EQ, [Var "n"; Value 1]),
              Assign ("sum", Value 0),
              While (
                App (Gt, [Var "n"; Value 1]),
                Seq (
                  Assign ("sum", App (Add, [Var "a"; Var "b"])),
                  Seq (
                    Assign ("n", App (Sub, [Var "n"; Value 1])),
                    Seq (
                      Assign ("a", Var "b"),
                      Assign ("b", Var "sum")
                    )
                  )
                )
              )
            )
          )
        )
      )
    )
  )


(* Run the program `c` under both MiniPython and StackLang, and return their
final values on the variable `out`. You can run this in utop to see if your
compiler is correct. *)
let test_compile c out =
  let compiled = compile_cmd c in 
  let result_store_1 = MiniPython.eval_cmd StringMap.empty c in 
  let result_store_2 = StackLang.execute compiled in
  (StringMap.find_opt out result_store_1, StringMap.find_opt out result_store_2)

let test_if_cmd () = test_compile if_cmd "result"

let test_while_cmd () = test_compile while_cmd "n"

let test_fib_cmd n = 
  (* Run all test cases for fibonacci from 0 to (n - 1). Should give you an output like 

  [(Some 0, Some 0); (Some 0, Some 0); (Some 1, Some 1); (Some 2, Some 2);
  (Some 3, Some 3); (Some 5, Some 5); (Some 8, Some 8); (Some 13, Some 13)]

  for n = 8.
  *)
  List.map (fun i -> test_compile (fib_cmd i) "sum") (List.init n (fun i -> i))