open Types

module Utility = Utility;;

Random.self_init ()

let (let*) = Result.bind

let rec evaluate (mem: memory) (command: c_exp) : (memory, [> error]) result =
  match command with
    Skip -> Ok mem
  | Assignment (v, exp_a) ->
    let* vval = evaluate_a mem exp_a in
    Ok {
      (* Map.add replaces the previus value *)
      assignments = VariableMap.add v vval 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) -> (
      let* guard = evaluate_b mem exp_b in
      if guard then
        evaluate mem exp_c1
      else
        evaluate mem exp_c2
    )
  | While (exp_b, exp_c) -> (
      let* guard = evaluate_b mem exp_b in
      if guard then
        let* mem2 = evaluate mem exp_c in
        evaluate mem2 command
      else
        Ok mem
    )
  | For (exp_c1, exp_b, exp_c2, body_c) -> (
      let* mem2 = evaluate mem exp_c1 in
      let rec f (localmem: memory) : (memory, [> error]) result =
        let* guard = (evaluate_b localmem exp_b) in
        if guard
        then
            let* stepmem = evaluate localmem body_c in
            let* incrementmem = evaluate stepmem exp_c2 in
            f incrementmem
        else Ok localmem
      in
      f mem2
    )


and evaluate_a (mem: memory) (exp_a: a_exp) : (int, [> error]) result =
  match exp_a with
    Variable v -> (
      match VariableMap.find_opt v mem.assignments with
        None -> Error (`AbsentAssignment ("The variable " ^ v ^ " is not defined."))
      | Some a -> Ok a
    )
  | Integer n -> Ok n
  | Plus (exp_a1, exp_a2) -> (
      let* exp_a1val = evaluate_a mem exp_a1 in
      let* exp_a2val = evaluate_a mem exp_a2 in
      Ok (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
      Ok (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
      Ok (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
        Ok (exp_a1val / exp_a2val)
      with Division_by_zero -> Error (`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
      Ok (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
      Ok (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
      Ok (Utility.powmod exp_a1val exp_a3val exp_a2val)
    )
  | Rand (exp_a) -> (
      let* exp_aval = evaluate_a mem exp_a in
      Ok (Random.int exp_aval)
    )


and evaluate_b (mem: memory) (exp_b: b_exp) : (bool, [> error]) result =
  match exp_b with
    Boolean b -> Ok b
  | BAnd (exp_b1, exp_b2) -> (
      let* exp_b1val = evaluate_b mem exp_b1 in
      let* exp_b2val = evaluate_b mem exp_b2 in
      Ok (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
      Ok (exp_b1val || exp_b2val)
    )
  | BNot (exp_b) -> (
      let* exp_bval = evaluate_b mem exp_b in
      Ok (not exp_bval)
    )
  | BCmp (exp_a1, exp_a2) -> (
      let* exp_a1val = evaluate_a mem exp_a1 in
      let* exp_a2val = evaluate_a mem exp_a2 in
      Ok (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
      Ok (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
      Ok (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
      Ok (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
      Ok (exp_a1val >= exp_a2val)
    )


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