open Gui
open Fix

(* var_t *)
module VarS = struct
  type 'var_t t = V of string

  let map f v = match v with V (s) -> V (s)

  let unify f v1 v2 = match (v1, v2) with
      (V (str1), V (str2)) when str1 = str2 -> ()
    | _ -> raise (Unify_Error "of VarS")

  let draw v = match v with V (s) -> create_str s

  let search _ = false
end

module Var = Fix.F (VarS)

let var n = Var.make_fix (VarS.V (n))

(* term_t *)
module TermS = struct
  type 'term_t t = Var of Var.fix_t
		 | Abs of Var.fix_t * 'term_t
		 | App of 'term_t * 'term_t

  let search tm = match tm with
      Var (v) -> Var.search v
    | Abs (v, tm) ->
	(* let で書かないと実行されないかもしれないので、注意 *)
	let b1 = Var.search v in b1 || tm
	(* Var.search v || tm *)
    | App (tm1, tm2) -> tm1 || tm2

  let map f tm = match tm with
      Var (v) -> Var (v) (* もう一度書かないと、型エラー *)
    | Abs (v, f1) -> Abs (v, f f1)
    | App (f1, f2) -> App (f f1, f f2)

  let unify f tm1 tm2 = match (tm1, tm2) with
      (Var (v1), Var (v2)) -> Var.unify v1 v2
    | (Abs (v1, f1), Abs (v2, f2)) -> Var.unify v1 v2; f f1 f2
    | (App (f1, f2), App (f3, f4)) -> f f1 f3; f f2 f4
    | _ -> raise (Unify_Error "of Terms")

  let draw tm = match tm with
      Var (v) -> Var.draw v
    | Abs (v, f) ->
	let id = Var.draw v in
	combineH [create_str "(λ"; id; create_str ". "; f; create_str ")"]
    | App (fix1, fix2) ->
	combineH [create_str "("; fix1; create_str " "; fix2; create_str ")"]
end

module Term = Fix.F (TermS)

let term_var v = Term.make_fix (TermS.Var (v))
let term_abs x e = Term.make_fix (TermS.Abs (x, e))
let term_app e1 e2 = Term.make_fix (TermS.App (e1, e2))

(* type_t *)
module TypeS = struct
  type 'type_t t = TVar of string
		 | Fun of 'type_t * 'type_t

  let search t = match t with
      TVar (tv) -> false
    | Fun (tp1, tp2) -> tp1 || tp2

  let map f tp = match tp with
      TVar (str) -> TVar (str)
    | Fun (tp1, tp2) -> Fun (f tp1, f tp2)

  let unify f tp1 tp2 = match (tp1, tp2) with
      (TVar (s1), TVar (s2)) when s1 = s2 -> ()
    | (Fun (tp1, tp2), Fun (tp3, tp4)) -> f tp1 tp3; f tp2 tp4
    | _ -> raise (Unify_Error "of Types")

  let draw tp = match tp with
      TVar (str) -> create_str str
    | Fun (tp1, tp2) ->
        combineH [create_str "("; tp1; create_str "→"; tp2; create_str ")"]
end

module Type = Fix.F (TypeS)

let tvar str = Type.make_fix (TypeS.TVar (str))
let tfun tp1 tp2 = Type.make_fix (TypeS.Fun (tp1, tp2))

(* env_t *)
module EnvS = struct
  type 'env_t t = Empty
		| Cons of Var.fix_t * Type.fix_t * 'env_t

  let search env = match env with
      Empty -> false
    | Cons (v, tp, env) ->
	let b1 = Var.search v in
	let b2 = Type.search tp in
	b1 || b2 || env
	(* (Var.search v) || (Type.search tp) || env *)

  let map f env = match env with
      Empty -> Empty
    | Cons (v, tp, env) -> Cons (v, tp, f env)

  let unify f env1 env2 = match (env1, env2) with
      (Empty, Empty) -> ()
    | (Cons (v1, tp1, env1), Cons (v2, tp2, env2)) ->
	Var.unify v1 v2;
	Type.unify tp1 tp2;
	f env1 env2
    | _ -> raise (Unify_Error "of Env")

  let draw env = match env with
      Empty -> create_str "Φ"
    | Cons (v, tp, rest) ->
	combineH [Var.draw v; create_str ":"; Type.draw tp;
		  create_str ", "; rest]
end

module Env = Fix.F (EnvS)

let empty = Env.make_fix (EnvS.Empty)
let cons x ts e = Env.make_fix (EnvS.Cons (x, ts, e))

(* judge_t *)
module JudgeS = struct
  type 'judge_t t = Judge of Env.fix_t * Term.fix_t * Type.fix_t
		  | Equal of Type.fix_t * Type.fix_t

  let search j = match j with
      Judge (env, tm, tp) ->
	let b1 = Env.search env in
	let b2 = Term.search tm in
	let b3 = Type.search tp in
	b1 || b2 || b3
    | Equal (tp1, tp2) -> Type.search tp1 || Type.search tp2

  let map _ j = match j with
      Judge (e, tm, tp) -> Judge (e, tm, tp)
    | Equal (tp1, tp2) -> Equal (tp1, tp2)

  let unify f j1 j2 = match (j1, j2) with
      (Judge (env1,tm1,tp1), Judge (env2,tm2,tp2)) ->
	Env.unify env1 env2;
	Term.unify tm1 tm2;
	Type.unify tp1 tp2
    | (Equal (tp1,tp2), Equal (tp3,tp4)) ->
	Type.unify tp1 tp3;
	Type.unify tp2 tp4
    | _ -> raise (Unify_Error "of Envs")

  let draw j = match j with
      Judge (env, tm, tp) ->
	combineH[Env.draw env; create_str "├ "; Term.draw tm;
		 create_str " : "; Type.draw tp]
    | Equal (tp1, tp2) ->
	combineH[Type.draw tp1; create_str " = "; Type.draw tp2]
end

module Judge = Fix.F (JudgeS)

let judge e tm tp = Judge.make_fix (JudgeS.Judge (e, tm, tp))
let equal tp1 tp2 = Judge.make_fix (JudgeS.Equal (tp1, tp2))

(* 推論規則を定義する *)
(* 環境から変数に当たる型を取ってくる *)
let rec get env var = try ( (* try... with で，中身が取れなかった場合を書く *)
  let env_e = Env.take env in (* env の中身を取ってくる *)
  let var_e = Var.take var in (* var の中身を取ってくる *)
  (match env_e with
      EnvS.Empty -> raise Not_found
    | EnvS.Cons (var2, tp, rest) ->
	let var2 = Var.take var2 in
	if var_e = var2 then tp else get rest var))
with _ ->
  Type.make_code
    (fun () -> combineH[create_str "get("; Env.draw env; create_str ",";
			Var.draw var; create_str ")"])
    (fun () -> get env var)

(* TVAR *)
let t_var () =
  let x = Var.make_meta "X" in
  let t = Type.make_meta "T" in
  let e = Env.make_meta "E" in
  Infer (judge e (term_var x) t,
	 [Infer (equal t (get e x), [])])

(* TABS *)
let t_abs () =
  let e = Env.make_meta "E" in
  let x = Var.make_meta "X" in
  let tm = Term.make_meta "T" in
  let tp1 = Type.make_meta "T" in
  let tp2 = Type.make_meta "T" in
  Infer(judge e (term_abs x tm) (tfun tp1 tp2),
	[Infer(judge (cons x tp1 e) tm tp2, [])])

(* TAPP *)
let t_app () =
  let e = Env.make_meta "E" in
  let tm1 = Term.make_meta "T" in
  let tm2 = Term.make_meta "T" in
  let tp1 = Type.make_meta "T" in
  let tp2 = Type.make_meta "T" in
  Infer(judge e (term_app tm1 tm2) tp2,
	[Infer(judge e tm1 (tfun tp1 tp2),[]);
	 Infer(judge e tm2 tp1,[])])

(* TEQUAL *)
let t_equal () =
  let tp = Type.make_meta "T" in
  Axiom(equal tp tp)

(* 登録される推論規則 *)
let infer_rule_list =
  [("TVAR", t_var); ("TABS", t_abs); ("TAPP", t_app); ("TEQUAL", t_equal)]