open Gui

(* 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
		 | Shift of Var.fix_t * 'term_t
		 | Reset of '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
    | Shift (v, tm) -> let b1 = Var.search v in b1 || tm
    | Reset (tm) -> tm

  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)
    | Shift (v, f1) -> Shift (v, f f1)
    | Reset (f1) -> Reset (f f1)

  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
    | (Shift (v1, f1), Shift (v2, f2)) -> Var.unify v1 v2; f f1 f2
    | (Reset (f1), Reset (f2)) -> f f1 f2
    | _ -> 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 [fix1; create_str " "; fix2]
    | Shift (v, f) ->
	let id = Var.draw v in
	combineH [create_str "(ξ"; id; create_str ". "; f; create_str ")"]
    | Reset (f) ->
	combineH [create_str "<"; f; 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))
let term_shift x e = Term.make_fix (TermS.Shift (x, e))
let term_reset e = Term.make_fix (TermS.Reset (e))

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

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

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

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

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

module Type = Fix.F (TypeS)

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

(* 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 * Type.fix_t *
			     Term.fix_t * Type.fix_t * Type.fix_t
		  | Equal of Type.fix_t * Type.fix_t

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

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

  let unify f j1 j2 = match (j1, j2) with
      (Judge (env1,tp11,tm1,tp12,tp13), Judge (env2,tp21,tm2,tp22,tp23)) ->
	Env.unify env1 env2;
	Term.unify tm1 tm2;
	Type.unify tp11 tp21;
	Type.unify tp12 tp22;
	Type.unify tp13 tp23
    | (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, tp1, tm, tp2, tp3) ->
	combineH[Env.draw env; create_str ";"; Type.draw tp1;
		 create_str "├ "; Term.draw tm;
		 create_str " : "; Type.draw tp2;
		 create_str ";"; Type.draw tp3]
    | Equal (tp1, tp2) ->
	combineH[Type.draw tp1; create_str " = "; Type.draw tp2]
end

module Judge = Fix.F (JudgeS)

let judge e tp1 tm tp2 tp3 =
		    Judge.make_fix (JudgeS.Judge (e, tp1, tm, tp2, tp3))
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 () =
  (* x : t; e, a |- x : t, a *)
  let x = Var.make_meta "X" in
  let t = Type.make_meta "T" in
  let e = Env.make_meta "E" in
  let a = Type.make_meta "A" in
  Infer (judge e a (term_var x) t a,
	 [Infer (equal t (get e x), [])])

(* TABS *)
let t_abs () =
  (* x : t1; e, t3 |- m : t2, t4       *)
  (* --------------------------------- *)
  (* e, t5 |- λx.m : t1/t3->t2/t4, t5 *)
  let e = Env.make_meta "E" in
  let x = Var.make_meta "X" in
  let m = Term.make_meta "M" in
  let t1 = Type.make_meta "T" in
  let t2 = Type.make_meta "T" in
  let t3 = Type.make_meta "A" in
  let t4 = Type.make_meta "A" in
  let t5 = Type.make_meta "A" in
  Infer(judge e t5 (term_abs x m) (tfun t1 t3 t2 t4) t5,
	[Infer(judge (cons x t1 e) t3 m t2 t4, [])])

(* TAPP *)
let t_app () =
  (* e, t4 |- m : t1/t6->t2/t5, t3 *)
  (* e, t5 |- n : t1, t4           *)
  (* ----------------------------- *)
  (* e, t6 |- m n : t2, t3         *)
  let e = Env.make_meta "E" in
  let m = Term.make_meta "M" in
  let n = Term.make_meta "N" in
  let t1 = Type.make_meta "T" in
  let t2 = Type.make_meta "T" in
  let t3 = Type.make_meta "A" in
  let t4 = Type.make_meta "A" in
  let t5 = Type.make_meta "A" in
  let t6 = Type.make_meta "A" in
  Infer(judge e t6 (term_app m n) t2 t3,
	[Infer(judge e t4 m (tfun t1 t6 t2 t5) t3, []);
	 Infer(judge e t5 n t1 t4, [])])

(* TSHIFT *)
let t_shift () =
  (* k : t / d -> a / d; e, s |- m : s, b *)
  (* ------------------------------------ *)
  (* e, a |- Sk. m : t, b                 *)
  let e = Env.make_meta "E" in
  let k = Var.make_meta "K" in
  let m = Term.make_meta "M" in
  let t = Type.make_meta "T" in
  let a = Type.make_meta "A" in
  let b = Type.make_meta "A" in
  let d = Type.make_meta "D" in
  let s = Type.make_meta "S" in
  Infer(judge e a (term_shift k m) t b,
	[Infer(judge (cons k (tfun t d a d) e) s m s b, [])])

(* TRESET *)
let t_reset ()=
  (* e, s |- m : s, t   *)
  (* ------------------ *)
  (* e, a |- <m> : t, a *)
  let e = Env.make_meta "E" in
  let m = Term.make_meta "M" in
  let s = Type.make_meta "T" in
  let t = Type.make_meta "T" in
  let a = Type.make_meta "A" in
  Infer(judge e a (term_reset m) t a,
	[Infer(judge e s m s t, [])])

(* 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);
   ("TSHIFT", t_shift); ("TRESET", t_reset); ("TEQUAL", t_equal)]