type monop = Neg | Not
type binop = Add | Sub | Mul | Div | Mod    (* arithmetic operators *)
           | Eq | Ne | Lt | Le | Gt | Ge    (* relational operators *)
           | And | Or                       (* logic operators *)
type expr = Num of int                      (* integer constant *)
          | Bool of bool                    (* boolean constant *)
          | Var of string                   (* variable *)
          | MonopAp of monop * expr         (* unary operator application *)
          | BinopAp of binop * expr * expr  (* binary operator application *)
          | Cond of expr * expr * expr      (* conditional expression *)
          | Fun of string * expr            (* function abstraction *)
          | FunAp of expr * expr            (* function application *)
          | Let of string * expr * expr     (* variable binding *)
          | LetRec of string * expr * expr  (* recursive variable binding *)

(* Convert a given expression to a string representation *)
let rec expr2string =
    let monop2string = function
          Neg -> "-"
        | Not -> "not "
    in
    let binop2string = function
          Add -> "+"
        | Sub -> "-"
        | Mul -> "*"
        | Div -> "/"
        | Mod -> "mod"
        | Eq -> "="
        | Ne -> "!="
        | Lt -> "<"
        | Le -> "<="
        | Gt -> ">"
        | Ge -> ">="
        | And -> "&&"
        | Or -> "||"
    in
    function
      Num i -> string_of_int i
    | Bool b -> string_of_bool b
    | Var v -> v
    | MonopAp (op, e) -> monop2string(op) ^ expr2string(e)
    | BinopAp (op, e1, e2) -> expr2string(e1) ^ " " ^ binop2string(op)
                              ^ " " ^ expr2string(e2)
    | Cond (cond, e1, e2) -> "if " ^ expr2string(cond)
                             ^ " then " ^ expr2string(e1)
                             ^ " else " ^ expr2string(e2)
    | Fun (arg, e) -> "fun " ^ arg ^ " -> " ^ expr2string(e)
    | FunAp (e1, e2) ->
        (* Add parenthesis around argument when needed *)
        let eval_e2 = expr2string(e2) in
        let eval_e2 = (if (String.get eval_e2 0) != '('
                          && String.contains eval_e2 ' '
                       then "(" ^ eval_e2 ^ ")" else eval_e2) in
        "(" ^ expr2string(e1) ^ " " ^ eval_e2 ^ ")"
    | Let (var, e1, e2) -> "let " ^ var ^ " = " ^ expr2string(e1)
                           ^ " in " ^ expr2string(e2)
    | LetRec (var, e1, e2) -> "let rec " ^ var ^ " = " ^ expr2string(e1)
                              ^ " in " ^ expr2string(e2)

(* Find all free variables within an expression *)
let freevars e =
    let rec free_vars bound_vars = function
          Num _ | Bool _ -> []
        | Var v when List.mem v bound_vars -> []
        | Var v -> [v]
        | MonopAp (_, e) -> free_vars bound_vars e
        | BinopAp (_, e1, e2)
        | FunAp (e1, e2) -> free_vars bound_vars e1 @ free_vars bound_vars e2
        | Cond (cond, e1, e2) -> free_vars bound_vars cond
                                 @ free_vars bound_vars e1
                                 @ free_vars bound_vars e2
        | Fun (arg, e) -> free_vars (arg::bound_vars) e
        | Let (v, e1, e2) -> free_vars (v::bound_vars) e2
                             @ free_vars bound_vars e1
        | LetRec (v, e1, e2) -> let bound = v::bound_vars in
                                free_vars bound e2 @ free_vars bound e1
    in
    free_vars [] e

(* Substitue all occurences of y by a in exp *)
let rec subs exp y a = match exp with
      Var x when x = y -> a
    | MonopAp (op, e) -> MonopAp (op, subs e y a)
    | BinopAp (op, e1, e2) -> BinopAp (op, subs e1 y a, subs e2 y a)
    | Fun (x, _) when x = y -> exp
    | Fun (x, e) ->
        let free_a = freevars a in
        let free_e = freevars e in
        if List.mem y free_e && List.mem x free_a then
            (* Replace x with a new variable w to avoid name clash *)
            let all_free_vars = free_a @ free_e in
            let rec add_prime x =
                let new_x = x ^ "'" in
                if List.mem new_x all_free_vars then add_prime x else new_x
            in
            let w = add_prime x in
            Fun (w, subs (subs e x (Var w)) y a)
        else
            Fun (x, subs e y a)
    | FunAp (e1, e2) -> FunAp (subs e1 y a, subs e2 y a)
    | Cond (cond, e1, e2) ->  Cond (subs cond y a, subs e1 y a, subs e2 y a)
    | Let (x, e1, e2) when x != y -> Let (x, subs e1 y a, subs e2 y a)
    | _ -> exp
