open Types

module Utility = Utility;;

Random.self_init ()

let rec evaluate (mem: memory) (command: c_exp) =
  match command with
    Skip -> mem
  | Assignment (v, exp_a) -> {
      (* Map.add replaces the previeus value *)
      assignments = VariableMap.add v (evaluate_a mem exp_a) mem.assignments
    }
  | Sequence (exp_c1, exp_c2) -> (
      let mem2 = evaluate mem exp_c1 in
      evaluate mem2 exp_c2
    )
  | If (exp_b, exp_c1, exp_c2) -> (
      if evaluate_b mem exp_b then
        evaluate mem exp_c1
      else
        evaluate mem exp_c2
    )
  | While (exp_b, exp_c) -> (
      if evaluate_b mem exp_b then
        let mem2 = evaluate mem exp_c in
        evaluate mem2 command
      else
        mem
    )
  | For (exp_c1, exp_b, exp_c2, body_c) -> (
      let mem2 = evaluate mem exp_c1 in
      let rec f localmem =
        if (evaluate_b localmem exp_b)
        then f (
            let tmpmem = (evaluate localmem body_c) in
            (evaluate tmpmem exp_c2))
        else localmem
      in
      f mem2
    )

and evaluate_a (mem: memory) (exp_a: a_exp) =
  match exp_a with
    Variable v -> (
      match VariableMap.find_opt v mem.assignments with
        None -> raise (AbsentAssignment ("The variable " ^ v ^ " is not defined."))
      | Some a -> a
    )
  | Integer n -> n
  | Plus (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val + exp_a2val
    )
  | Minus (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val - exp_a2val
    )
  | Times (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val * exp_a2val
    )
  | Division (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      try
        exp_a1val / exp_a2val
      with Division_by_zero -> raise (DivisionByZero "Dividing by zero")
    )
  | Modulo (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val mod exp_a2val
    )
  | Power (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      Utility.pow exp_a1val exp_a2val
    )
  | PowerMod (exp_a1, exp_a2, exp_a3) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      let exp_a3val = evaluate_a mem exp_a3 in
      Utility.powmod exp_a1val exp_a3val exp_a2val
    )
  | Rand (exp_a) -> (
      Random.int (evaluate_a mem exp_a)
    )
and evaluate_b (mem: memory) (exp_b: b_exp) =
  match exp_b with
    Boolean b -> b
  | BAnd (exp_b1, exp_b2) -> (
      let exp_b1val = evaluate_b mem exp_b1 in
      let exp_b2val = evaluate_b mem exp_b2 in
      exp_b1val && exp_b2val
    )
  | BOr (exp_b1, exp_b2) -> (
      let exp_b1val = evaluate_b mem exp_b1 in
      let exp_b2val = evaluate_b mem exp_b2 in
      exp_b1val || exp_b2val
    )
  | BNot (exp_b) -> (
      not (evaluate_b mem exp_b)
    )
  | BCmp (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val = exp_a2val
    )
  | BCmpLess (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val < exp_a2val
    )
  | BCmpLessEq (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val <= exp_a2val
    )
  | BCmpGreater (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val > exp_a2val
    )
  | BCmpGreaterEq (exp_a1, exp_a2) -> (
      let exp_a1val = evaluate_a mem exp_a1 in
      let exp_a2val = evaluate_a mem exp_a2 in
      exp_a1val >= exp_a2val
    )


let reduce (program: p_exp) (iin : int) =
  match program with
    Main (vin, vout, expression) -> (
      let mem : memory = {assignments = (VariableMap.empty |> VariableMap.add vin iin)} in
      match VariableMap.find_opt vout (evaluate mem expression).assignments with
        None -> raise (AbsentAssignment ("The output variable is not defined (" ^ vout ^ ")"))
      | Some a -> a
    )