module Utility = Utility
open Types;;

Random.self_init ()

let (let*) = Result.bind

let rec evaluate (mem: memory) (command: t_exp) :
  (permitted_values, [> error]) result =
  match command with
    Integer n -> Ok (IntegerPermitted n)
  | Boolean b -> Ok (BooleanPermitted b)
  | Variable v -> (
      match VariableMap.find_opt v mem.assignments with
      | None ->
        Error (`AbsentAssignment ("The variable " ^ v ^ " is not defined."))
      | Some a ->
        Ok a
    )
  | Tuple (x, y) -> (
      let* xval = evaluate mem x in
      let* yval = evaluate mem y in
      Ok (TuplePermitted (xval, yval))
    )
  | Function (x, _, f) ->
    Ok (FunctionPermitted
          {input = x;
           body = f;
           assignments = mem.assignments;
           recursiveness = None}
       )
  | Application (f, x) -> (
      let* evalf = evaluate mem f in
      let* func_closure = (
        match evalf with
          FunctionPermitted ff -> Ok ff
        | IntegerPermitted _ -> Error (`WrongType ("Function is not a function,"
                                                   ^ " it's an integer"))
        | BooleanPermitted _ -> Error (`WrongType ("Function is not a function,"
                                                   ^ " it's a boolean"))
        | TuplePermitted _ -> Error (`WrongType ("Function is not a function,"
                                                   ^ " it's a tuple"))
      ) in
      let* param = evaluate mem x in
      let mem2 =
        match func_closure.recursiveness with
          None ->  {assignments = (
            VariableMap.add func_closure.input param func_closure.assignments)}
        | Some nameF -> {assignments = (
            VariableMap.add func_closure.input param func_closure.assignments |>
            VariableMap.add nameF (FunctionPermitted func_closure)
          )}
      in
      evaluate mem2 func_closure.body
    )
  | Plus (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (aval + bval))
  | Minus (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (aval - bval))
  | Times (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (aval * bval))

  | Division (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in (
      try
        Ok (IntegerPermitted (aval / bval))
      with Division_by_zero -> Error (`DivisionByZero "Dividing by zero")
    )

  | Modulo (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (aval mod bval))
  | Power (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (Utility.pow aval bval))

  | PowerMod (a, b, c) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    let* cval = (
      match evaluate mem c with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (Utility.powmod aval bval cval))
  | Rand (a) ->
    let* aval = (
      match evaluate mem a with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    )
    in
    Ok (IntegerPermitted (Random.int aval))

  | BAnd (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok BooleanPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an boolean"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok BooleanPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an boolean"))
    )
    in
    Ok (BooleanPermitted (aval && bval))

  | BOr (a, b) ->
    let* aval = (
      match evaluate mem a with
        Ok BooleanPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an boolean"))
    )
    in
    let* bval = (
      match evaluate mem b with
        Ok BooleanPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an boolean"))
    )
    in
    Ok (BooleanPermitted (aval || bval))

  | BNot a ->
    let* aval = (
      match evaluate mem a with
        Ok BooleanPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an boolean"))
    )
    in
    Ok (BooleanPermitted (not aval))
  | First a ->
    let* aval = (
      match evaluate mem a with
        Ok TuplePermitted (x, _) -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not a tuple"))
    )
    in
    Ok (aval)
  | Second a ->
    let* aval = (
      match evaluate mem a with
        Ok TuplePermitted (_, x) -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not a tuple"))
    )
    in
    Ok (aval)
  | Cmp (exp_1, exp_2) ->
    let* exp_1val = match evaluate mem exp_1 with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    in
    let* exp_2val = match evaluate mem exp_2 with
        Ok IntegerPermitted x -> Ok x
      | Error e -> Error e
      | _ -> Error (`WrongType ("Value is not an integer"))
    in
    Ok (BooleanPermitted (exp_1val = exp_2val))

  | CmpLess (exp_1, exp_2) -> (
      let* exp_1val = match evaluate mem exp_1 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      let* exp_2val = match evaluate mem exp_2 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      Ok (BooleanPermitted (exp_1val < exp_2val))
    )
  | CmpLessEq (exp_1, exp_2) -> (
      let* exp_1val = match evaluate mem exp_1 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      let* exp_2val = match evaluate mem exp_2 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      Ok (BooleanPermitted (exp_1val <= exp_2val))
    )
  | CmpGreater (exp_1, exp_2) -> (
      let* exp_1val = match evaluate mem exp_1 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      let* exp_2val = match evaluate mem exp_2 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      Ok (BooleanPermitted (exp_1val > exp_2val))
    )
  | CmpGreaterEq (exp_1, exp_2) -> (
      let* exp_1val = match evaluate mem exp_1 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      let* exp_2val = match evaluate mem exp_2 with
          Ok IntegerPermitted x -> Ok x
        | Error e -> Error e
        | _ -> Error (`WrongType ("Value is not an integer"))
      in
      Ok (BooleanPermitted (exp_1val >= exp_2val))
    )
  | IfThenElse (guard, if_exp, else_exp) ->
    let* bguard = (
      match evaluate mem guard with
        Ok BooleanPermitted b -> Ok b
      | Error e -> Error e
      | _ -> Error (`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, x, _, fbody, rest) ->
    let mem2 = {
      assignments =
        VariableMap.add
          f
          (FunctionPermitted
             { input = x;
               body = fbody;
               assignments = mem.assignments;
               recursiveness = Some f})
          mem.assignments}
    in
    evaluate mem2 rest


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