module Utility = Utility
open Types;;

Random.self_init ()

let (let*) = Result.bind


module VariableSet = Set.Make(String)

let rec freevariables (t: intermediaryTypes) : VariableSet.t =
  match t with
    (IntermediaryBooleanType) -> VariableSet.empty
  | (IntermediaryIntegerType) -> VariableSet.empty
  | (IntermediaryVariableType x) -> VariableSet.singleton x
  | (IntermediaryFunctionType (Some name, tin)) ->
    let partres = List.fold_left (fun acc x -> VariableSet.union acc (freevariables x)) VariableSet.empty tin in
    VariableSet.add name partres
  | (IntermediaryFunctionType (None, tin)) ->
    List.fold_left (fun acc x -> VariableSet.union acc (freevariables x)) VariableSet.empty tin

let rec replaceunified (t: intermediaryTypes) (unifier: substitution) : intermediaryTypes =
  match t with
    IntermediaryIntegerType
  | IntermediaryBooleanType -> t
  | IntermediaryVariableType x -> (
      match VariableMap.find_opt x unifier with
        Some tnew -> tnew
      | None -> t
    )
  | IntermediaryFunctionType (Some name, tin) -> (
      match VariableMap.find_opt name unifier with
        Some tnew -> tnew
      | None -> IntermediaryFunctionType (Some name, List.map (fun x -> replaceunified x unifier) tin)
    )
  | IntermediaryFunctionType (None, tin) ->
    IntermediaryFunctionType (None, List.map (fun x -> replaceunified x unifier) tin)

let rec unify t1 t2 =
  match (t1, t2) with
  | (_, IntermediaryVariableType x2) ->
    if VariableSet.mem x2 (freevariables t1) then
        Error (`WrongType "Cannot unify")
    else
      Ok (VariableMap.add x2 t1 VariableMap.empty)
  | (IntermediaryVariableType x1, _) ->
    if VariableSet.mem x1 (freevariables t2) then
        Error (`WrongType "Cannot unify")
    else
      Ok (VariableMap.add x1 t2 VariableMap.empty)
  | (IntermediaryIntegerType, IntermediaryIntegerType)
  | (IntermediaryBooleanType, IntermediaryBooleanType) -> Ok VariableMap.empty
  | (IntermediaryFunctionType (Some name1, tin1), IntermediaryFunctionType (Some name2, tin2)) ->
    if name1 <> name2 then
      Error (`WrongType "Different functions cannot be unified")
    else
      List.fold_left2 (
        fun acc t1 t2 ->
          if Result.is_error acc then (* stop at the first error *)
            acc
          else
            let* partres = unify (replaceunified t1 (Result.get_ok acc)) (replaceunified t2 (Result.get_ok acc)) in
            try
              Ok (VariableMap.union (fun _ val1 val2 -> if val1 = val2 then Some val1 else failwith "Should not happen") partres acc)
            with Failure e ->
              Error (`WrongType e)
      ) (Ok VariableMap.empty) tin1 tin2
  | (IntermediaryFunctionType (None, _), IntermediaryFunctionType _)
  | (IntermediaryFunctionType _, IntermediaryFunctionType (None, _)) ->
    Error (`WrongType "Unnamed Functions cannot be unified")
  | _ -> failwith "Not implemented"

let rec evaluate_type (program: t_exp) (context: ftype VariableMap.t) : (ftype, 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
    )
  | Function (xs, typef, fbody) -> (
      (* first check that the function has the right specified type then check
         that the number of inputs are the right number, finally eval the type
         of the body using the bindings for the inputs *)
      match typef with
        FunctionType (tin, tout) -> (
          if List.length xs <> List.length tin then
            Error (`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
            let* typefbody = evaluate_type fbody context1 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, xs) -> (
      (* check that the type is actually a function, then checks that the
         supplied inputs are of the right type, returns the return type if all
         inputs are supplied, otherwise a function from the remaining inputs to
         the output types *)
      let* evalf = evaluate_type f context in
      match evalf with
        FunctionType (tin, tout) -> (
          let rec helper (params: t_exp list) (typeparams: ftype list) =
            (* consumes until params are exausted *)
            match (params, typeparams) with
              ([], _) -> Ok typeparams
            | (_, []) -> Error (`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 = Ok v then
                helper tlparams tltypeparams
              else
                Error (`WrongType "Argument with wrong type.")
          in
          let* typesremaining = helper xs tin in
          match typesremaining with
            [] -> Ok tout
          | t -> Ok (FunctionType (t, tout))
        )
      | _ -> 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.")
    )
  | 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, xs, typef, fbody, rest) ->
    (* like with the function type, but also add f itself to the bindings *)
    match typef with
      FunctionType (tin, tout) -> (
        if List.length xs <> List.length tin then
          Error (`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
          let* typefbody = evaluate_type fbody context2 in
          let* typerest = evaluate_type rest context1 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, error) result =
  let* typeprogram = evaluate_type program VariableMap.empty in
  match typeprogram with
    FunctionType ([IntegerType], IntegerType) -> (
      Ok (FunctionType ([IntegerType], IntegerType))
    )
  | _ -> Error (`WrongType "Program is not a function from int to int.")