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. I think the code is
 * self-explanatory and therefore needs no more documentation *)
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 it contains spaces. I'm not
         * sure if this is required in this implementation, but I found that is
         * is tidy when trying to represent OCaml syntax (see the 'fac'
         * function in 'test_expr.ml') *)
        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. Where no rules apply,
 * concatenat the free variables of all sub-expressions *)
let freevars e =
    let rec free_vars bound_vars = function
          Num _ | Bool _ -> []
        | Var v when List.mem v bound_vars -> []
        (* FV(x) = {x} *)
        | Var v -> [v]
        | MonopAp (_, e) -> free_vars bound_vars e
        | BinopAp (_, e1, e2)
        (* FV((e1 e2)) = FV(e1) UNION FV(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
        (* FV(lamda x.e) = FV(e) \ {x} *)
        | 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 using the substitution rules
 * provided in the lecture.
 * NOTE: I have decided to resolve name clashes here instead of in the eval
 * function of assignment 13. The reason for this is that both in the lecture
 * and in lambda calculus wiki's on the internet, name clashes are resolved by
 * applying the substution rule that I've implemented in lines 112-120 below.
 * This function is called 'subs', so it should implement this substitution
 * rule *)
let rec subs exp y a = match exp with
      Var x when x = y -> a
    (* substitute in operator arguments *)
    | MonopAp (op, e) -> MonopAp (op, subs e y a)
    | BinopAp (op, e1, e2) -> BinopAp (op, subs e1 y a, subs e2 y a)
    (* (lambda y.e)[a/x] --> lambda y.e
     * if y = x *)
    | Fun (x, _) when x = y -> exp
    (* (lambda y.e)[a/x] --> lambda w.(e[w/y][a/x])
     * if y != x, x free in e, y free in a, new variable w
     * lambda y.(e[a/x])
     * otherwise *)
    | 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 a name clash *)
            let all_free_vars = free_a @ free_e in
            (* Add primes to the new variable until it is free *)
            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)
    (* (e e')[a/x] --> (e[a/x])(e'[a/x]) *)
    | FunAp (e1, e2) -> FunAp (subs e1 y a, subs e2 y a)
    (* Conditional expression has no substirution rules, just substitute in all
     * sub-expressions recursively *)
    | Cond (cond, e1, e2) ->  Cond (subs cond y a, subs e1 y a, subs e2 y a)
    (* Only substitue in a 'let'/'let rec' expression if there is no name clash
     * with the assigned variable *)
    | Let (x, e1, e2) when x != y -> Let (x, subs e1 y a, subs e2 y a)
    | LetRec (x, e1, e2) when x != y -> LetRec (x, subs e1 y a, subs e2 y a)
    (* No substitition rule applies, just return the expression *)
    | _ -> exp
