type variable = string

type p_exp =
    Main of variable * variable * c_exp (* def main with input x output y as c *)
and c_exp =
    Skip
  | Assignment of variable * a_exp      (* x := a *)
  | Sequence of c_exp * c_exp           (* c; c *)
  | If of b_exp * c_exp * c_exp         (* if b then c else c *)
  | While of b_exp * c_exp              (* while b do c *)
and b_exp =
    Boolean of bool                     (* v *)
  | BAnd of b_exp * b_exp               (* b and b *)
  | BNot of b_exp                       (* not b *)
  | BCmpLess of a_exp * a_exp           (* a < a *)
and a_exp =
    Variable of variable                (* x *)
  | Integer of int                      (* n *)
  | Plus of a_exp * a_exp               (* a + a *)
  | Minus of a_exp * a_exp              (* a - a *)
  | Times of a_exp * a_exp              (* a * a *)



module VariableMap = Map.Make(String)

type memory = {
  assignments: int VariableMap.t
}


exception AbsentAssignment of string

let rec evaluate (mem: memory) (command: c_exp) =
  match command with
    Skip -> mem
  | Assignment (v, exp_a) -> {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
    )

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
    )
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
    )
  | BNot (exp_b) -> (
      not (evaluate_b mem exp_b)
    )
  | 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
    )


let reduce (iin : int) (program: p_exp) =
  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
    )