type variable = string

module VariableMap = Map.Make(String)

type ftype =
    IntegerType
  | BooleanType
  | FunctionType of ftype list * ftype

type t_exp =
    Integer of int
  | Boolean of bool
  | Variable of variable
  | Function of variable list * ftype * t_exp
  | Application of t_exp * t_exp list
  | Plus of t_exp * t_exp
  | Minus of t_exp * t_exp
  | Times of t_exp * t_exp
  | Division of t_exp * t_exp
  | Modulo of t_exp * t_exp
  | Power of t_exp * t_exp
  | PowerMod of t_exp * t_exp * t_exp
  | Rand of t_exp
  | BAnd of t_exp * t_exp
  | BOr of t_exp * t_exp
  | BNot of t_exp
  | Cmp of t_exp * t_exp
  | CmpLess of t_exp * t_exp
  | CmpLessEq of t_exp * t_exp
  | CmpGreater of t_exp * t_exp
  | CmpGreaterEq of t_exp * t_exp
  | IfThenElse of t_exp * t_exp * t_exp
  | LetIn of variable * t_exp * t_exp
  | LetFun of variable * variable list * ftype * t_exp * t_exp

type permittedValues =
    IntegerPermitted of int
  | BooleanPermitted of bool
  | FunctionPermitted of closure
and closure = {
  inputList: variable list;
  body: t_exp;
  assignments: permittedValues VariableMap.t;
  recursiveness: variable option
}

type memory = {
  assignments: permittedValues VariableMap.t
}

exception AbsentAssignment of string
exception WrongType of string
exception DivisionByZero of string
exception WrongArity of string
exception WrongTypeSpecification of string

module Utility = Utility;;

Random.self_init ()

let rec evaluate (mem: memory) (command: t_exp) =
  match command with
    Integer n -> (IntegerPermitted n)
  | Boolean b -> (BooleanPermitted b)
  | Variable v -> (
      match VariableMap.find_opt v mem.assignments with
        None -> raise (AbsentAssignment ("The variable " ^
                                         v ^
                                         " is not defined."))
      | Some a -> a
    )
  | Function (xs, _, f) ->
    (FunctionPermitted
       {inputList = xs;
        body = f;
        assignments = mem.assignments;
        recursiveness = None}
    )
  | Application (f, xs) -> (
      let funcClosure = (
        match (evaluate mem f) with
          FunctionPermitted ff -> ff
        | IntegerPermitted _ -> raise (WrongType ("Function is not a function, it's an integer"))
        | BooleanPermitted _ -> raise (WrongType ("Function is not a function, it's a boolean"))
      ) in
      let parmList = List.map (fun k -> evaluate mem k) xs in
      let rec helper m params values =
        match (params, values) with
          (_, []) -> (m, params)
        | ([], _) ->
          raise (WrongArity ("Function application has arity " ^
                              (List.length funcClosure.inputList
                               |> string_of_int) ^
                              ", but was applied to " ^
                              (List.length xs |> string_of_int) ^
                              " parameters"))
        | (p::tlparams, v::tlvalues) -> helper
                                          (VariableMap.add p v m)
                                          tlparams
                                          tlvalues
      in
      let (mem2assignments, params) = helper
          funcClosure.assignments
          funcClosure.inputList
          parmList
      in (* helper funcClosure or helper mem ??? *)
      let mem2 = (
        match funcClosure.recursiveness with
          None -> {assignments = mem2assignments}
        | Some nameF -> {
            assignments =
              VariableMap.add
                nameF
                (FunctionPermitted funcClosure)
                mem2assignments
          }
      ) in
      match params with
        [] -> evaluate mem2 funcClosure.body
      | _ -> (
          FunctionPermitted {funcClosure with inputList = params;
                                              assignments = mem2assignments})
    )
  | Plus (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (aval + bval))
  | Minus (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (aval - bval))
  | Times (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (aval * bval))
  | Division (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in (
      try
        (IntegerPermitted (aval / bval))
      with Division_by_zero -> raise (DivisionByZero "Dividing by zero")
    )
  | Modulo (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (aval mod bval))
  | Power (a, b) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (Utility.pow aval bval))
  | PowerMod (a, b, c) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let bval = (
      match evaluate mem b with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    let cval = (
      match evaluate mem c with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    (IntegerPermitted (Utility.powmod aval bval cval))
  | Rand (a) ->
    let aval = (
      match evaluate mem a with
        IntegerPermitted x -> x
      | _ -> raise (WrongType ("Value is not an integer"))
    )
    in
    IntegerPermitted (Random.int aval)
  | BAnd (a, b) ->
    let aval = (
      match evaluate mem a with
        BooleanPermitted x -> x
      | _ -> raise (WrongType ("Value is not an boolean"))
    )
    in
    let bval = (
      match evaluate mem b with
        BooleanPermitted x -> x
      | _ -> raise (WrongType ("Value is not an boolean"))
    )
    in
    (BooleanPermitted (aval && bval))
  | BOr (a, b) ->
    let aval = (
      match evaluate mem a with
        BooleanPermitted x -> x
      | _ -> raise (WrongType ("Value is not an boolean"))
    )
    in
    let bval = (
      match evaluate mem b with
        BooleanPermitted x -> x
      | _ -> raise (WrongType ("Value is not an boolean"))
    )
    in
    (BooleanPermitted (aval || bval))
  | BNot a ->
    let aval = (
      match evaluate mem a with
        BooleanPermitted x -> x
      | _ -> raise (WrongType ("Value is not an boolean"))
    )
    in
    (BooleanPermitted (not aval))
  | Cmp (exp_1, exp_2) -> (
      let exp_1val = match evaluate mem exp_1 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      let exp_2val = match evaluate mem exp_2 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      BooleanPermitted (exp_1val = exp_2val)
    )
  | CmpLess (exp_1, exp_2) -> (
      let exp_1val = match evaluate mem exp_1 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      let exp_2val = match evaluate mem exp_2 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      BooleanPermitted (exp_1val < exp_2val)
    )
  | CmpLessEq (exp_1, exp_2) -> (
      let exp_1val = match evaluate mem exp_1 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      let exp_2val = match evaluate mem exp_2 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      BooleanPermitted (exp_1val <= exp_2val)
    )
  | CmpGreater (exp_1, exp_2) -> (
      let exp_1val = match evaluate mem exp_1 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      let exp_2val = match evaluate mem exp_2 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      BooleanPermitted (exp_1val > exp_2val)
    )
  | CmpGreaterEq (exp_1, exp_2) -> (
      let exp_1val = match evaluate mem exp_1 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      let exp_2val = match evaluate mem exp_2 with
          IntegerPermitted x -> x
        | _ -> raise (WrongType ("Value is not an integer"))
      in
      BooleanPermitted (exp_1val >= exp_2val)
    )
  | IfThenElse (guard, if_exp, else_exp) ->
    let bguard = (
      match evaluate mem guard with
        BooleanPermitted b -> b
      | _ -> raise (WrongType ("Value in if guard is not a boolean"))
    ) in
    if bguard then
      evaluate mem if_exp
    else
      evaluate mem else_exp
  | LetIn (x, xval, rest) ->
    let evalxval = evaluate mem xval in
    let mem2 = {assignments = VariableMap.add x evalxval mem.assignments} in
    evaluate mem2 rest
  | LetFun (f, xs, _, fbody, rest) ->
    let mem2 = {
      assignments =
        VariableMap.add
          f
          (FunctionPermitted
             { inputList = xs;
               body = fbody;
               assignments = mem.assignments;
               recursiveness = Some f})
          mem.assignments}
    in
    evaluate mem2 rest


let reduce (program: t_exp) (iin : int) =
  let program' = (Application (program, [(Integer iin)])) in
  let mem : memory = {assignments = VariableMap.empty} in
  match (evaluate mem program') with
    IntegerPermitted a -> a
  | _ -> raise (WrongType ("Main function doesn't return an integer"))




let rec evaluate_type (program: t_exp) context =
  match program with
    Integer _ -> IntegerType
  | Boolean _ -> BooleanType
  | Variable x -> (match VariableMap.find_opt x context with
        None -> raise (AbsentAssignment ("The variable " ^ x ^ " is not defined."))
      | Some t -> t)
  | Function (xs, typef, fbody) -> (
      match typef with
        FunctionType (tin, tout) -> (
          if List.length xs != List.length tin then
            raise (WrongTypeSpecification "Type specification for function has wrong arity.")
          else
            let context1 = List.fold_left2 (fun acc x t -> VariableMap.add x t acc) context xs tin in
            match (evaluate_type fbody context1 = tout) with
              (false) -> raise (WrongTypeSpecification "Function does not return specified type.")
            | (true) -> typef
        )
      | _ -> raise (WrongTypeSpecification "Specification of function is not a function type.")
    )
  | Application (f, xs) -> (
      match evaluate_type f context with
        FunctionType (tin, tout) -> (
          let rec helper params typeparams =
            match (params, typeparams) with
              ([], _) -> typeparams
            | (_, []) -> raise (WrongArity ("Function application has arity " ^
                                             (List.length tin |> string_of_int) ^
                                             ", but was applied to " ^
                                             (List.length xs |> string_of_int) ^
                                             " parameters"))
            | (p::tlparams, v::tltypeparams) ->
              if evaluate_type p context = v then
                helper tlparams tltypeparams
              else
                raise (WrongType "Argument with wrong type.")
          in
          match helper xs tin with
            [] -> tout
          | t -> FunctionType (t, tout)
        )
      | _ -> raise (WrongType "Applying to a non function type")
    )
  | Plus (x, y)
  | Minus (x, y)
  | Times (x, y)
  | Division (x, y)
  | Modulo (x, y)
  | Power (x, y) -> (
      match (evaluate_type x context, evaluate_type y context) with
      | (IntegerType, IntegerType) -> IntegerType
      | (IntegerType, _) -> raise (WrongType "Second term is not an integer.")
      | (_, _) -> raise (WrongType "First term is not an integer.")
    )
  | PowerMod (x, y, z) -> (
      match (evaluate_type x context, evaluate_type y context, evaluate_type z context) with
      | (IntegerType, IntegerType, IntegerType) -> IntegerType
      | (IntegerType, IntegerType, _) -> raise (WrongType "Third term is not an integer.")
      | (IntegerType, _, _) -> raise (WrongType "Second term is not an integer.")
      | (_, _, _) -> raise (WrongType "First term is not an integer.")
    )
  | Rand (x) -> (
      match (evaluate_type x context) with
      | (IntegerType) -> IntegerType
      | (_) -> raise (WrongType "Term is not an integer.")
    )
  | BAnd (x, y)
  | BOr (x, y) -> (
      match (evaluate_type x context, evaluate_type y context) with
      | (BooleanType, BooleanType) -> BooleanType
      | (BooleanType, _) -> raise (WrongType "Second term is not a boolean.")
      | (_, _) -> raise (WrongType "First term is not a boolean.")
    )
  | BNot (x) -> (
      match (evaluate_type x context) with
      | (BooleanType) -> BooleanType
      | (_) -> raise (WrongType "Term is not a boolean.")
    )
  | Cmp (x, y)
  | CmpLess (x, y)
  | CmpLessEq (x, y)
  | CmpGreater (x, y)
  | CmpGreaterEq (x, y) -> (
      match (evaluate_type x context, evaluate_type y context) with
      | (IntegerType, IntegerType) -> BooleanType
      | (IntegerType, _) -> raise (WrongType "Second term is not an integer.")
      | (_, _) -> raise (WrongType "First term is not an integer.")
    )
  | IfThenElse (guard, if_exp, else_exp) -> (
      match (evaluate_type guard context, evaluate_type if_exp context, evaluate_type else_exp context) with
        (BooleanType, t1, t2) -> (
          if t1 = t2 then
            t1
          else
            raise (WrongType "If branches do not have the same type.")
        )
      | (_, _, _) -> raise (WrongType "If guard is not a boolean.")
    )
  | LetIn (x, xval, rest) ->
    let typex = evaluate_type xval context in
    evaluate_type rest (VariableMap.add x typex context)
  | LetFun (f, xs, typef, fbody, rest) ->
    match typef with
      FunctionType (tin, tout) -> (
        if List.length xs != List.length tin then
          raise (WrongArity "Type specification for function has wrong arity.")
        else
          let context1 = VariableMap.add f typef context in
          let context2 = List.fold_left2 (fun acc x t -> VariableMap.add x t acc) context1 xs tin in
          match (evaluate_type fbody context2 = tout, evaluate_type rest context1) with
            (false, _) -> raise (WrongTypeSpecification "Function does not return specified type."
            )
          | (true, t) -> t
      )
    | _ -> raise (WrongTypeSpecification "Specification of function is not a function type.")

let typecheck (program: t_exp) =
  match evaluate_type program VariableMap.empty with
    FunctionType ([IntegerType], IntegerType) -> true
  | _ -> raise (WrongType "Program is not a function from int to int.")