(*******************************************************************************
 * WIP implementation of the Hindley-Milner type system primarily for learning
 * purposes.
 *
 * Wish List:
 * - TODO Debug this inference (let f (fn x x) f)
 ******************************************************************************)

open Types

(*******************************************************************************
 * Library
 ******************************************************************************)

let ( let* ) = Option.bind

let set_from_list (xs : string list) : set =
  xs |> List.fold_left (fun acc x -> FromString.add x true acc) FromString.empty

(* Map union that favors the rhs values (i.e. "last writer wins"). *)
let lww (xs : 'a FromString.t) (ys : 'a FromString.t) : 'a FromString.t =
  FromString.union (fun k x y -> Some y) xs ys

let emptyEnv : env = FromString.empty

let rec free_type_vars (t : _type) : set =
  match t with
  | TypeVariable k -> FromString.singleton k true
  | TypeInt -> FromString.empty
  | TypeBool -> FromString.empty
  | TypeArrow (a, b) -> lww (free_type_vars a) (free_type_vars b)

let i : int ref = ref 0

let make_type_var () : _type =
  let res = Printf.sprintf "a%d" !i in
  i := !i + 1;
  TypeVariable res

exception OccursCheck

let bind_var (k : string) (t : _type) : substitution =
  if t == TypeVariable k then FromString.empty
  else if FromString.exists (fun name _ -> name == k) (free_type_vars t) then
    raise OccursCheck
  else FromString.singleton k t

let rec instantiate (q : quantified_type) : _type =
  let (QuantifiedType (names, t)) = q in
  match t with
  | TypeInt -> TypeInt
  | TypeBool -> TypeBool
  | TypeVariable k ->
      if List.exists (( == ) k) names then make_type_var () else TypeVariable k
  | TypeArrow (a, b) ->
      TypeArrow
        (instantiate (QuantifiedType (names, a)), instantiate (QuantifiedType (names, b)))

let quantified_type_ftvs (q : quantified_type) : set =
  let (QuantifiedType (names, t)) = q in
  lww (free_type_vars t) (names |> set_from_list)

let generalize (env : env) (t : _type) : quantified_type =
  let envftv =
    env |> FromString.bindings
    |> List.map (fun (_, v) -> quantified_type_ftvs v)
    |> List.fold_left lww FromString.empty
  in
  let names =
    lww (free_type_vars t) envftv
    |> FromString.bindings
    |> List.map (fun (k, _) -> k)
  in
  QuantifiedType (names, t)

let rec substitute_type (s : substitution) (t : _type) : _type =
  match t with
  | TypeVariable k as tvar ->
     (match FromString.find_opt k s with
      | Some v -> substitute_type s v
      | None -> tvar)
  | TypeArrow (a, b) -> TypeArrow (substitute_type s a, substitute_type s b)
  | TypeInt -> TypeInt
  | TypeBool -> TypeBool

let substitute_quantified_type (s : substitution) (q : quantified_type) : quantified_type =
  let (QuantifiedType (names, t)) = q in
  let s1 =
    FromString.filter (fun k v -> List.exists (fun x -> k != x) names) s
  in
  QuantifiedType (names, substitute_type s1 t)

let substitute_env (s : substitution) (env : env) : env =
  FromString.map (fun q -> substitute_quantified_type s q) env

let compose_substitutions (xs : substitution list) : substitution =
  let do_compose_substitutions s1 s2 = lww s2 (FromString.map (substitute_type s2) s1) in
  List.fold_left do_compose_substitutions FromString.empty xs

let rec unify (a : _type) (b : _type) : substitution option =
  match (a, b) with
  | TypeInt, TypeInt -> Some FromString.empty
  | TypeBool, TypeBool -> Some FromString.empty
  | TypeVariable k, _ -> Some (bind_var k b)
  | _, TypeVariable k -> Some (bind_var k a)
  | TypeArrow (a, b), TypeArrow (c, d) ->
      let* s1 = unify a c in
      let* s2 = unify (substitute_type s1 b) (substitute_type s1 d) in
      let s3 = compose_substitutions [s1; s2] in
      s1 |> Types.debug_substitution |> Printf.sprintf "s1: %s\n" |> print_string;
      s2 |> Types.debug_substitution |> Printf.sprintf "s2: %s\n" |> print_string;
      s3 |> Types.debug_substitution |> Printf.sprintf "s3: %s\n" |> print_string;
      Some s3
  | _ -> None

let print_env (env : env) =
  Printf.sprintf "env: %s\n" (Types.debug_env env)
  |> print_string

let print_val (x : value) =
  Printf.sprintf "val: %s\n" (Types.debug_value x)
  |> print_string

let print_inference (x : inference option) =
  match x with
  | None -> "no inference\n" |> print_string
  | Some x ->
     Printf.sprintf "inf: %s\n" (Types.debug_inference x)
     |> print_string

let rec infer (env : env) (x : value) : inference option =
  print_env env;
  print_val x;
  let res = match x with
  | ValueLiteral lit -> (
      match lit with
      | LiteralInt _ -> Some (Inference (FromString.empty, TypeInt))
      | LiteralBool _ -> Some (Inference (FromString.empty, TypeBool)))
  | ValueVariable k ->
      let* v = FromString.find_opt k env in
      Some (Inference (FromString.empty, instantiate v))
  | ValueFunction (param, body) ->
      let typevar = make_type_var () in
      let env1 = FromString.remove param env in
      let env2 = lww (FromString.singleton param (QuantifiedType ([], typevar))) env1 in
      let* (Inference (s1, t1)) = infer env2 body in
      Some (Inference (s1, TypeArrow (substitute_type s1 typevar, t1)))
  | ValueApplication (f, x) ->
      let result = make_type_var () in
      let* (Inference (s1, t1)) = infer env f in
      let* (Inference (s2, t2)) = infer (substitute_env s1 env) x in
      let* s3 = unify (substitute_type s2 t1) (TypeArrow (t2, result)) in
      Some (Inference
              ( compose_substitutions [s3; s2; s1],
                substitute_type s3 result ))
  | ValueVarApplication (name, x) ->
      let* v = FromString.find_opt name env in
      let t1 = instantiate v in
      let typevar = make_type_var () in
      let* (Inference (s2, t2)) = infer env x in
      let* s3 = unify (substitute_type s2 t1) (TypeArrow (t2, typevar)) in
      Some (Inference
              ( compose_substitutions [s2; s3],
                substitute_type s3 typevar ))
  | ValueBinder (k, v, body) ->
      let* (Inference (s1, t1)) = infer env v in
      let env1 = FromString.remove k env in
      let tg = generalize (substitute_env s1 env) t1 in
      let env2 = FromString.add k tg env1 in
      let* (Inference (s2, t2)) = infer (substitute_env s1 env2) body in
      Some (Inference (compose_substitutions [s1; s2], t2)) in
  print_inference res;
  res

let do_infer (x : value) : _type option =
  let* Inference (_, t) = infer FromString.empty x in
  Some t