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 *)
  | For of c_exp * b_exp * c_exp * c_exp (* for (c; b; c) do c *)
and b_exp =
    Boolean of bool                      (* v *)
  | BAnd of b_exp * b_exp                (* b && b *)
  | BOr of b_exp * b_exp                 (* b || b *)
  | BNot of b_exp                        (* not b *)
  | BCmp of a_exp * a_exp                (* a == a *)
  | BCmpLess of a_exp * a_exp            (* a < a *)
  | BCmpLessEq of a_exp * a_exp          (* a <= a *)
  | BCmpGreater of a_exp * a_exp         (* a > a *)
  | BCmpGreaterEq 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 *)
  | Division of a_exp * a_exp            (* a / a *)
  | Modulo of a_exp * a_exp              (* a % a *)
  | Power of a_exp * a_exp               (* a ^ a *)
  | PowerMod of a_exp * a_exp * a_exp    (* a ^ a % a *)
  | Rand of a_exp                        (* rand(0, a) *)

let pp_p_exp (ppf: Format.formatter) (p: p_exp) : unit =
  let rec helper_c (ppf) (c: c_exp) : unit =
    match c with
      Skip -> Format.fprintf ppf "Skip"
    | Assignment (x, a) -> Format.fprintf ppf "%S := @[<h>%a@]" x helper_a a
    | Sequence (c1, c2) -> Format.fprintf ppf "@[<hv>Sequence (@;<1 2>%a,@;<1 0>%a@;<0 0>)@]" helper_c c1 helper_c c2
    | If (b, c1, c2) -> Format.fprintf ppf "@[<hv>If @[<h>%a@]@;<1 2>then (@[<hv>%a@])@;<1 2>else (@[<hv>%a@])@]" helper_b b helper_c c1 helper_c c2
    | While (b, c) -> Format.fprintf ppf "@[<hv>While @[<h>%a@] do@;<1 2>%a@]@;<0 0>" helper_b b helper_c c
    | For (c1, b, c2, c3) -> Format.fprintf ppf "@[<h>For (@;<0 2>%a,@;<1 2>@[<h>%a@],@;<1 2>%a) do@]@;<1 4>%a@;<0 0>" helper_c c1 helper_b b helper_c c2 helper_c c3
  and helper_b (ppf) (b: b_exp) =
    match b with
      Boolean (b) -> Format.fprintf ppf "%b" b
    | BAnd (b1, b2) -> Format.fprintf ppf "(%a &&@;<1 2>%a)" helper_b b1 helper_b b2
    | BOr (b1, b2) -> Format.fprintf ppf "(%a ||@;<1 2>%a)" helper_b b1 helper_b b2
    | BNot (b) -> Format.fprintf ppf "(not %a)" helper_b b
    | BCmp (a1, a2) -> Format.fprintf ppf "(%a ==@;<1 2>%a)" helper_a a1 helper_a a2
    | BCmpLess (a1, a2) -> Format.fprintf ppf "(%a <@;<1 2>%a)" helper_a a1 helper_a a2
    | BCmpLessEq (a1, a2) -> Format.fprintf ppf "(%a <=@;<1 2>%a)" helper_a a1 helper_a a2
    | BCmpGreater (a1, a2) -> Format.fprintf ppf "(%a >@;<1 2>%a)" helper_a a1 helper_a a2
    | BCmpGreaterEq (a1, a2) -> Format.fprintf ppf "(%a >=@;<1 2>%a)" helper_a a1 helper_a a2
  and helper_a (ppf) (a: a_exp) =
    match a with
      Variable v -> Format.fprintf ppf "%S" v
    | Integer n -> Format.fprintf ppf "%i" n
    | Plus (a1, a2) -> Format.fprintf ppf "%a +@;<1 2>%a" helper_a a1 helper_a a2
    | Minus (a1, a2) -> Format.fprintf ppf "%a -@;<1 2>%a" helper_a a1 helper_a a2
    | Times (a1, a2) -> Format.fprintf ppf "%a *@;<1 2>%a" helper_a a1 helper_a a2
    | Division (a1, a2) -> Format.fprintf ppf "%a /@;<1 2>%a" helper_a a1 helper_a a2
    | Modulo (a1, a2) -> Format.fprintf ppf "%a %%@;<1 2>%a" helper_a a1 helper_a a2
    | Power (a1, a2) -> Format.fprintf ppf "(%a ^@;<1 2>%a)" helper_a a1 helper_a a2
    | PowerMod (a1, a2, a3) -> Format.fprintf ppf  "(%a ^ %a %% %a)" helper_a a1 helper_a a2 helper_a a3
    | Rand (a) -> Format.fprintf ppf "Rand (%a)" helper_a a
  in
  match p with
  | Main (i, o, exp) ->
    Format.fprintf ppf "def main with (input %S) (output %S) as @.%a" i o helper_c exp

module VariableMap = Map.Make(String)

type memory = {
  assignments: int VariableMap.t
}

type error = [
    `AbsentAssignment of string
  | `DivisionByZero of string
]