module Utility = Utility
open Types;;

Random.self_init ()

let (let*) = Result.bind

(* -------------------------------------------------------------------------- *)
(*                         polimporphic type checking                         *)
(* -------------------------------------------------------------------------- *)

let global_type_id = ref 0

let new_global_id () =
  let id = !global_type_id in
  incr global_type_id;
  id


let rec unify type_a type_b =
  if type_a == type_b then Ok () else
    match (type_a, type_b) with
    | IntegerTypeP, IntegerTypeP
    | BooleanTypeP, BooleanTypeP ->
      Ok ()

    | TupleTypeP (a1, a2), TupleTypeP (b1, b2)
    | ApplicationP (a1, a2), ApplicationP (b1, b2)
    | FunctionTypeP (a1, a2), FunctionTypeP (b1, b2) ->
      let* _ = unify a1 b1 in
      unify a2 b2

    | VariableTypeP ({contents = Link a1}),
      VariableTypeP ({contents = Link b1}) ->
      unify a1 b1

    | VariableTypeP ({contents = Link ty_link}), ty_rest
    | ty_rest, VariableTypeP ({contents = Link ty_link})
      when ty_link = ty_rest ->
      Ok ()

    | VariableTypeP ({contents = Unbound (a1, _)}),
      VariableTypeP ({contents = Unbound (b1, _)})
      when a1 = b1 ->
      Error (`WrongType "Only a single instance of a type should be available.")

    | type_ab, VariableTypeP ({contents = Unbound (_id, _level)} as tvar)
    | VariableTypeP ({contents = Unbound (_id, _level)} as tvar), type_ab ->
      tvar := Link type_ab;
      Ok ()

    | _, _ ->
      Error (`WrongType "Cannot unify types.")

let rec unifyable type_a type_b =
  if type_a == type_b then Ok () else
    match (type_a, type_b) with
    | IntegerTypeP, IntegerTypeP
    | BooleanTypeP, BooleanTypeP ->
      Ok ()

    | TupleTypeP (a1, a2), TupleTypeP (b1, b2)
    | ApplicationP (a1, a2), ApplicationP (b1, b2)
    | FunctionTypeP (a1, a2), FunctionTypeP (b1, b2) ->
      let* _ = unifyable a1 b1 in
      unifyable a2 b2

    | VariableTypeP ({contents = Link a1}),
      VariableTypeP ({contents = Link b1}) ->
      unifyable a1 b1

    | VariableTypeP ({contents = Link ty_link}), ty_rest
    | ty_rest, VariableTypeP ({contents = Link ty_link})
      when ty_link = ty_rest ->
      Ok ()

    | VariableTypeP ({contents = Unbound (a1, _)}),
      VariableTypeP ({contents = Unbound (b1, _)})
      when a1 = b1 ->
      Error (`WrongType "Only a single instance of a type should be available.")

    | _type_ab, VariableTypeP ({contents = Unbound (_id, _level)})
    | VariableTypeP ({contents = Unbound (_id, _level)}), _type_ab ->
      Ok ()

    | _, _ ->
      Error (`WrongType "Cannot unify types.")

let rec generalize level ty =
  match ty with
  | VariableTypeP {contents = Unbound (id, o_level)} when o_level > level ->
    VariableTypeP (ref (Generic id))
  | ApplicationP (ty, ty_arg) ->
    ApplicationP (generalize level ty, generalize level ty_arg)
  | FunctionTypeP (ty_arg, ty) ->
    FunctionTypeP (generalize level ty_arg, generalize level ty)
  | TupleTypeP (ty1, ty2) ->
    TupleTypeP (generalize level ty1, generalize level ty2)
  | VariableTypeP {contents = Link ty} ->
    generalize level ty
  | VariableTypeP {contents = Generic _}
  | VariableTypeP {contents = Unbound _}
  | IntegerTypeP
  | BooleanTypeP ->
    ty


let instantiate level ty =
  let var_map = ref IntegerMap.empty in
  let rec aux ty =
    match ty with
    | IntegerTypeP
    | BooleanTypeP ->
      ty
    | TupleTypeP (ty1, ty2) ->
      TupleTypeP (aux ty1, aux ty2)
    | VariableTypeP {contents = Link ty} ->
      aux ty
    | VariableTypeP {contents = Generic id} -> (
      match IntegerMap.find_opt id !var_map with
      | Some ty -> ty
      | None ->
        let var = VariableTypeP (ref (Unbound (new_global_id (), level))) in
        var_map := IntegerMap.add id var !var_map;
        var
    )
    | VariableTypeP {contents = Unbound _} ->
      ty
    | ApplicationP (ty, ty_arg) ->
      ApplicationP (aux ty, aux ty_arg)
    | FunctionTypeP (ty_arg, ty) ->
      FunctionTypeP (aux ty_arg, aux ty)
  in
  aux ty


let rec evaluate_type_polimorphic program (env: env) level =
  match program with
  | Integer _ -> Ok (IntegerTypeP)
  | Boolean _ -> Ok (BooleanTypeP)
  | Tuple (a, b) ->
    let* type_a = evaluate_type_polimorphic a env level in
    let* type_b = evaluate_type_polimorphic b env level in
    Ok (TupleTypeP (type_a, type_b))
  | Variable (x) -> (
      match VariableMap.find_opt x env with
      | Some (ty) ->
        Ok (instantiate level ty)
      | None ->
        Error (`AbsentAssignment ("Variable " ^ x ^ " is not assigned."))
    )
  | Function (x, _typef, fbody) ->
    let param_type = VariableTypeP (ref (Unbound (new_global_id (), level))) in
    let fn_env = VariableMap.add x param_type env in
    let* body_type = evaluate_type_polimorphic fbody fn_env level in
    Ok (FunctionTypeP (param_type, body_type))
  | Application (f, xs) ->
    let* type_f = evaluate_type_polimorphic f env level in
    let rec aux =
      function
      | FunctionTypeP (type_f_arg, type_f) ->
        Ok (type_f_arg, type_f)
      | VariableTypeP {contents = Link ty} ->
        aux ty
      | VariableTypeP ({contents = Unbound(_id, level)} as tvar) ->
        let param_ty = VariableTypeP (ref (Unbound (new_global_id (), level)))
        in
        let f_ty = VariableTypeP (ref (Unbound (new_global_id (), level))) in
        tvar := Link ( FunctionTypeP (param_ty, f_ty) );
        Ok (param_ty, f_ty)
      | _ -> Error (`WrongType "Expecting a function to apply.")
    in
    let* param_ty, f_ty = aux type_f in
    let* type_xs = evaluate_type_polimorphic xs env level in
    let* _ = unify param_ty type_xs in
    Ok f_ty
  | Plus (a, b)
  | Minus (a, b)
  | Times (a, b)
  | Division (a, b)
  | Modulo (a, b)
  | Power (a, b) ->
    let* type_a = evaluate_type_polimorphic a env level in
    let* type_b = evaluate_type_polimorphic b env level in
    let* _ = unify type_a IntegerTypeP in
    let* _ = unify type_b IntegerTypeP in
    Ok (IntegerTypeP)
  | First a -> (
      let* type_a = evaluate_type_polimorphic a env level in
      let* _ = unify type_a
          (TupleTypeP(VariableTypeP (ref (Unbound (new_global_id (), level))),
                      VariableTypeP (ref (Unbound (new_global_id (), level)))))
      in
      match type_a with
      | TupleTypeP (ty_a, _)
      | VariableTypeP {contents = Link TupleTypeP (ty_a, _)} -> Ok ty_a
      | _ -> Error (`WrongType "Applying First to non tuple type.")
    )
  | Second a -> (
      let* type_a = evaluate_type_polimorphic a env level in
      let* _ = unify type_a
          (TupleTypeP(VariableTypeP (ref (Unbound (new_global_id (), level))),
                      VariableTypeP (ref (Unbound (new_global_id (), level)))))
      in
      match type_a with
      | TupleTypeP (_, ty_a)
      | VariableTypeP {contents = Link TupleTypeP (_, ty_a)} -> Ok ty_a
      | _ -> Error (`WrongType "Applying Second to non tuple type.")
    )
  | PowerMod (x, y, z) ->
    let* type_x = evaluate_type_polimorphic x env level in
    let* type_y = evaluate_type_polimorphic y env level in
    let* type_z = evaluate_type_polimorphic z env level in
    let* _ = unify type_x IntegerTypeP in
    let* _ = unify type_y IntegerTypeP in
    let* _ = unify type_z IntegerTypeP in
    Ok (IntegerTypeP)
  | Rand (x) ->
    let* type_x = evaluate_type_polimorphic x env level in
    let* _ = unify type_x IntegerTypeP in
    Ok (IntegerTypeP)
  | BAnd (a, b)
  | BOr (a, b) ->
    let* type_a = evaluate_type_polimorphic a env level in
    let* type_b = evaluate_type_polimorphic b env level in
    let* _ = unify type_a BooleanTypeP in
    let* _ = unify type_b BooleanTypeP in
    Ok (BooleanTypeP)
  | BNot (x) ->
    let* type_x = evaluate_type_polimorphic x env level in
    let* _ = unify type_x BooleanTypeP in
    Ok (BooleanTypeP)
  | Cmp (a, b)
  | CmpLess (a, b)
  | CmpLessEq (a, b)
  | CmpGreater (a, b)
  | CmpGreaterEq (a, b) ->
    let* type_a = evaluate_type_polimorphic a env level in
    let* type_b = evaluate_type_polimorphic b env level in
    let* _ = unify type_a IntegerTypeP in
    let* _ = unify type_b IntegerTypeP in
    Ok (BooleanTypeP)
  | IfThenElse (guard, if_exp, else_exp) ->
    let* type_guard = evaluate_type_polimorphic guard env level in
    let* type_if_exp = evaluate_type_polimorphic if_exp env level in
    let* type_else_exp = evaluate_type_polimorphic else_exp env level in
    let* _ = unify type_guard BooleanTypeP in
    let* _ = unify type_if_exp type_else_exp in
    Ok (type_if_exp)
  | LetIn (x, xval, rest) ->
    let* var_ty = evaluate_type_polimorphic xval env (level + 1) in
    let generalized_ty = generalize level var_ty in
    evaluate_type_polimorphic rest (VariableMap.add x generalized_ty env) level
  | LetFun (f, xs, _typef, fbody, rest) ->
    let* _ = Error (`RecusrionsNotImplemented "Let Rec is not implemented.") in
    let tmp_type_f = VariableTypeP (ref (Unbound (new_global_id (), level))) in
    let new_env = VariableMap.add f tmp_type_f env in
    let param_type = VariableTypeP (ref (Unbound (new_global_id (), level))) in
    let fn_env = VariableMap.add xs param_type new_env in
    let* body_type = evaluate_type_polimorphic fbody fn_env level in
    let* _ = unify tmp_type_f body_type in
    evaluate_type_polimorphic rest (VariableMap.add f tmp_type_f env) level


(* -------------------------------------------------------------------------- *)


let rec evaluate_type (program: t_exp) (context: ftype VariableMap.t) :
  (ftype, [> typechecking_error]) result =
  match program with
    Integer _ -> Ok IntegerType
  | Boolean _ -> Ok BooleanType
  | Variable x -> ( (* check for the type in the context *)
      match VariableMap.find_opt x context with
        None -> Error (`AbsentAssignment
                         ("The variable " ^ x ^ " is not defined."))
      | Some t -> Ok t
    )
  | Tuple (x, y) -> (
      let* xtype = evaluate_type x context in
      let* ytype = evaluate_type y context in
      Ok (TupleType (xtype, ytype))
    )
  | Function (x, typef, fbody) -> (
      (* first check that the function has the right specified type then check
         the type of the body using the bindings for the input *)
      match typef with
        FunctionType (tin, tout) -> (
          let* typefbody = evaluate_type fbody (VariableMap.add x tin context)
          in
          if (typefbody = tout) then
            Ok typef
          else
            Error (`WrongTypeSpecification
                     ("Function does not return specified type."))
        )
      | _ -> Error (`WrongTypeSpecification
                      ("Specification of function is not a function type."))
    )
  | Application (f, x) -> (
      let* evalf = evaluate_type f context in
      let* evalx = evaluate_type x context in
      match evalf with
        FunctionType (tin, tout) -> (
          if tin = evalx then
            Ok tout
          else
            Error (`WrongType "Appling function with wrong input type to value")
        )
      | _ -> Error (`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) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (IntegerType, IntegerType) -> Ok IntegerType
      | (IntegerType, _) -> Error (`WrongType "Second term is not an integer.")
      | (_, _) -> Error (`WrongType "First term is not an integer.")
    )
  | PowerMod (x, y, z) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      let* typez = evaluate_type z context in
      match typex, typey, typez with
      | (IntegerType, IntegerType, IntegerType) -> Ok IntegerType
      | (IntegerType, IntegerType, _) -> Error (`WrongType ("Third term is " ^
                                                            "not an integer."))
      | (IntegerType, _, _) -> Error (`WrongType
                                        ("Second term is not an integer."))
      | (_, _, _) -> Error (`WrongType "First term is not an integer.")
    )
  | Rand (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (IntegerType) -> Ok IntegerType
      | (_) -> Error (`WrongType "Term is not an integer.")
    )
  | BAnd (x, y)
  | BOr (x, y) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (BooleanType, BooleanType) -> Ok BooleanType
      | (BooleanType, _) -> Error (`WrongType "Second term is not a boolean.")
      | (_, _) -> Error (`WrongType "First term is not a boolean.")
    )
  | BNot (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (BooleanType) -> Ok BooleanType
      | (_) -> Error (`WrongType "Term is not a boolean.")
    )
  | First (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (TupleType (x, _)) -> Ok x
      | (_) -> Error (`WrongType "Term is not a tuple.")
    )
  | Second (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (TupleType (_, x)) -> Ok x
      | (_) -> Error (`WrongType "Term is not a tuple.")
    )
  | Cmp (x, y)
  | CmpLess (x, y)
  | CmpLessEq (x, y)
  | CmpGreater (x, y)
  | CmpGreaterEq (x, y) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (IntegerType, IntegerType) -> Ok BooleanType
      | (IntegerType, _) -> Error (`WrongType "Second term is not an integer.")
      | (_, _) -> Error (`WrongType "First term is not an integer.")
    )
  | IfThenElse (guard, if_exp, else_exp) -> (
      let* typeguard = evaluate_type guard context in
      let* typeif_exp = evaluate_type if_exp context in
      let* typeelse_exp = evaluate_type else_exp context in
      match typeguard, typeif_exp, typeelse_exp with
        (BooleanType, t1, t2) -> (
          if t1 = t2 then
            Ok t1
          else
            Error (`WrongType "If branches do not have the same type.")
        )
      | (_, _, _) -> Error (`WrongType "If guard is not a boolean.")
    )
  | LetIn (x, xval, rest) ->
    (* bind the type to the variable name in the context *)
    let* typex = evaluate_type xval context in
    evaluate_type rest (VariableMap.add x typex context)
  | LetFun (f, x, typef, fbody, rest) ->
    (* like with the function case, but also add f itself to the bindings *)
    match typef with
      FunctionType (tin, tout) -> (
        let newcontext = VariableMap.add f typef context in
        let newcontextwithx = VariableMap.add x tin newcontext in
        let* typefbody = evaluate_type fbody newcontextwithx in
        let* typerest = evaluate_type rest newcontext in
        match (typefbody = tout, typerest) with
          (false, _) -> Error (`WrongTypeSpecification
                                 "Function does not return specified type.")
        | (true, t) -> Ok t
      )
    | _ -> Error (`WrongTypeSpecification
                    "Specification of function is not a function type.")


let typecheck (program: t_exp) : (ftype, [> typechecking_error]) result =
  let* typeprogram = evaluate_type program VariableMap.empty in
  match typeprogram with
    FunctionType (IntegerType, IntegerType) -> Ok (typeprogram)
  | _ -> Error (`WrongType "Program is not a function from int to int.")


let typecheck_polymorphic (program: t_exp)
  : (type_f, [> typechecking_error]) result =
  global_type_id := 0;
  let* type_program = evaluate_type_polimorphic program VariableMap.empty 0 in
  let* _ = unifyable type_program (FunctionTypeP (IntegerTypeP, IntegerTypeP))
  in
  let generalized_ty = generalize (-1) type_program in
  Ok (generalized_ty)

let typecheck_polymorphic_unbound (program: t_exp)
  : (type_f, [> typechecking_error]) result =
  global_type_id := 0;
  let* type_program = evaluate_type_polimorphic program VariableMap.empty 0 in
  let generalized_ty = generalize (-1) type_program in
  Ok (generalized_ty)