module Int = struct
  type t = int
  let compare x y = x - y
end
module IntMap = Map.Make (Int)
module StringMap = Map.Make (String)

type term =
  | Const of int
  | Var of int
  | App of int * term * term
  | Lam of int * bool * term * term
  | Star
  | Squere

let hashed mk =
  let cache = Hashtbl.create 123 in 
  let id = ref 0 in
  fun n ->
    try Hashtbl.find cache n
    with Not_found ->
      incr id;
      let v = mk !id n in
      Hashtbl.add cache n v;
      v

let const = hashed (fun _ x -> Const x)
let var = hashed (fun _ x -> assert (x >= 0); Var x)
let app = hashed (fun id (m, n) -> App (id, m, n))
let lam = hashed (fun id (f, m, n) -> Lam (id, f, m, n))
let star = Star
let squere = Squere

let fresh_var =
  let count = ref 0 in
  fun () -> incr count; !count

let pre_typed = Hashtbl.create 123

let offset n t =
  let rec aux d = function
    | t when Hashtbl.mem pre_typed t -> t
    | Lam (_, pi, ty, m) ->
      lam (pi, aux d ty, aux (d + 1) m)
    | App (_, t1, t2) ->
      app (aux d t1, aux d t2)
    | Var k when k >= d -> var (k + n)
    | (Star | Squere | Const _ | Var _) as t -> t
  in if n == 0 then t else aux 0 t

(* m[x := n] where m is in \x.m *)
let beta_forms = Hashtbl.create 123
let beta m n = 
  let rec aux x = function
    | Var y when x = y -> offset x n
    | Var y when x < y -> var (y - 1)
    | (Star | Squere | Const _ | Var _) as t -> t
    | App (_, t1, t2) ->
      app (aux x t1, aux x t2)
    | Lam (_, pi, ty, m) ->
      lam (pi, aux x ty, aux (x + 1) m)
  in 
    try Hashtbl.find beta_forms (m, n)
    with Not_found ->
      let nf = aux 0 m in
      Hashtbl.add beta_forms (m, n) nf;
      nf

let whnfs = Hashtbl.create 123
let rec whnf t =
  let aux = function
    | App (_, t1, t2) ->
      (match whnf t1 with
        | Lam (_, false, t, m) -> whnf (beta m t2)
        | t1 -> app (t1, whnf t2))
    | t -> t
  in 
  try Hashtbl.find whnfs t
  with Not_found ->
    let nf = aux t in
    Hashtbl.add whnfs t nf;
    nf

let beta_normals = Hashtbl.create 123
let rec beta_normal t =
  let aux = function
    | App (_, t1, t2) ->
      (match beta_normal t1 with
        | Lam (_, false, t, m) -> beta_normal (beta m t2)
        | t1 -> app (t1, beta_normal t2))
    | Lam (_, pi, t1, t2) -> lam (pi, beta_normal t1, beta_normal t2)
    | t -> t
  in
  try Hashtbl.find beta_normals t
  with Not_found ->
    let nf = aux t in
    Hashtbl.add beta_normals t nf;
    nf

let eq t1 t2 = t1 == t2

let err s = raise (Failure s)

let rec occur v = function
  | Var x -> v = x
  | Const _ | Star | Squere -> false
  | App (_, m, n) -> occur v m || occur v n
  | Lam (_, _, m, n) -> occur v m || occur (v + 1) n

let const_names = Hashtbl.create 123 

let t2s t = 
  let b = Buffer.create 100 in
  let say = Buffer.add_string b in
  let gen_name vars =
    let len = List.length vars in
    let name = 
      if len < 25 then String.make 1 (char_of_int (int_of_char 'A' + len))
      else Printf.sprintf "X%d" len
    in
    (name :: vars, name)
  in
  let rec f vars = function
    | Const c -> say (Hashtbl.find const_names c)
    | Var x -> 
      (try say (List.nth vars x)
       with Invalid_argument _ | Failure _ -> Printf.bprintf b "x%d" x)
    | App (_, App (_, Const ex, t), Lam (_, false, t', m)) when 
      eq t t' && Hashtbl.find const_names ex = "ex" ->
      let (vars', name) = gen_name vars in
      Printf.bprintf b "(exists (%s : " name;
      f vars t; say (") "); f vars' m; say ")"
    | App (_, m, n) -> 
      say "("; f vars m; say " "; f vars n; say ")"
    | Lam (_, pi, m, n) -> 
      if pi && not (occur 0 n) then (
        say "("; f vars m; say " -> "; f ("_" :: vars) n; say ")"
      ) else (
        say (if pi then "(forall " else "(fun ");
        let (vars', name) = gen_name vars in
        Printf.bprintf b "%s : " name; f vars m;
        say ". "; f vars' n; say ")"
      )
    | Star -> say "*"
    | Squere -> say "[]"
  in f [] t; Buffer.contents b

let rec typ depth env = function
  | t when Hashtbl.mem pre_typed t -> Hashtbl.find pre_typed t
  | Star -> squere
  | Squere -> err "cannot type []"
  | Const x ->
    (try IntMap.find x env
     with Not_found -> err "cannot find const in env")
  | Var x -> 
    (try offset (x + 1) (IntMap.find (x - depth) env)
     with Not_found -> err "cannot find var in env")
  | App (_, m, n) ->
    (match whnf (typ depth env m) with
      | Lam (_, true, a, b) ->
        let a' = typ depth env n in
        (* Printf.printf "cmp %s %s\n" (t2s a) (t2s a'); *)
        if eq (beta_normal a) (beta_normal a') then 
          ( (* Printf.printf "%s <- %s\n" (t2s b) (t2s n);  *)
          beta b n)
        else err ("not-eq " ^ t2s a ^ " and " ^ t2s a' ^ " when trying to type " ^ t2s (app (m,n)))
      | t -> err ("cannot apply " ^ t2s t))
  | Lam (_, false, a, m) -> 
    let b = typ (depth + 1) (IntMap.add (- depth - 1) a env) m in
    let r = lam (true, a, b) in
    let _ = typ depth env r in r
  | Lam (_, true, a, b) ->
    (match whnf (typ depth env a), whnf (typ (depth + 1) (IntMap.add (- depth - 1) a env) b) with
      | Star, Star
      | Squere, Star (* F2 *)
      | Squere, Squere  (* Fomega *)
      | Star, Squere (*-> err "CC not allowed" *)
        as p -> (snd p)
      | (t, s) -> err ("evil type types: " ^ t2s t ^ " and " ^ t2s s))

type sexpr =
  | Appl of string * sexpr list
  | Atom of string

let rec s2s = function
  | Atom n -> n
  | Appl (n, lst) ->
    "(" ^ n ^ " " ^ (String.concat " " (List.map s2s lst)) ^ ")"

let sexpr_to_term env s = 
  let rec self depth env s =
    let get s = 
      try 
        match StringMap.find s env with
          | Var n -> var (depth - n)
          | t -> t
      with Not_found -> err ("unbound symbol " ^ s) in
    match s with
      | Atom "*" -> star
      | Atom s -> get s
      | Appl ("let", [Appl (name, [m]); n]) ->
        let m = self depth env m in
        let n = self (depth + 1) (StringMap.add name (var (depth + 1)) env) n in
        beta n m
      | Appl ("let", _) ->
        err ("malformed let " ^ s2s s)
      | Appl (hd, args) when hd = "forall" || hd = "fun" ->
        let rec loop depth env = function
          | [t] -> self depth env t
          | Appl (name, [ty]) :: rest ->
            lam (hd = "forall", self depth env ty, 
                 loop (depth + 1) (StringMap.add name (var (depth + 1)) env) rest)
          | [] -> err ("illformed " ^ hd ^ " got []")
          | t :: _ -> err ("illformed " ^ hd ^ " got " ^ s2s t)
        in loop depth env args
      | Appl ("exists", args) ->
        let rec loop = function
          | [t] -> t
          | Appl (name, [ty]) as hd :: rest ->
            Appl ("ex", [ty; Appl ("fun", [hd; loop rest])])
          | _ -> err ("ill formed exists")
        in self depth env (loop args)
      | Appl ("implies", args) ->
        let rec loop = function
          | [t] -> t
          | ty :: rest ->
            Appl ("forall", [Appl ("_", [ty]); loop rest])
          | _ -> err ("ill formed implies")
        in self depth env (loop args)
      | Appl ("", x :: xs) ->
        List.fold_left (fun acc arg -> app (acc, self depth env arg)) (self depth env x) xs
      | Appl (name, args) ->
        List.fold_left (fun acc arg -> app (acc, self depth env arg)) (get name) args
  in self 0 env s
  

type my_file = { handle : in_channel; mutable peeked : int; }

let rec peek f =
  if f.peeked >= 0 then char_of_int f.peeked
  else ( f.peeked <- input_byte f.handle; peek f )

let read f = let c = peek f in f.peeked <- -1; c
let putback f c = f.peeked <- int_of_char c

let rec non_ws f =
  match read f with
    | ' ' | '\t' | '\n' -> non_ws f
    | ';' -> let rec loop () = if read f = '\n' then non_ws f else loop () in loop ()
    | c -> c

let read_word f =
  let b = Buffer.create 10 in
  let rec loop () =
    match peek f with
      | ' ' | '\t' | '\n' | '(' | ')' | ';' -> Buffer.contents b
      | _ -> Buffer.add_char b (read f); loop ()
  in loop ()

let rec read_sexpr f =
  match peek f with
    | '(' ->
      let _ = read f in
      let hd = read_word f in
      let rec loop acc = function
        | ')' -> Appl (hd, List.rev acc)
        | c ->
          putback f c; 
          let inner = read_sexpr f in loop (inner :: acc) (non_ws f)
      in loop [] (non_ws f)
    | _ -> Atom (read_word f)

let top_read f =
  let has_more =
    try let c = non_ws f in putback f c; true with End_of_file -> false
  in if has_more then read_sexpr f else Atom ("quit")

let f = { handle = open_in Sys.argv.(1); peeked = -1 }
let verbose = ref false

let rec eval_loop (types, numbers) = 
  let cont name ty =
    let id = fresh_var () in
    Hashtbl.add const_names id name;
    match typ 0 types ty with
      | Star | Squere ->
        let ty = beta_normal ty in
        if !verbose then
          Printf.printf "val %s : %s\n" name (t2s ty)
        else
          Printf.printf "val %s\n" name;
        flush stdout;
        eval_loop (IntMap.add id ty types, 
                   StringMap.add name (Const id) numbers)
      | t -> err ("the type: " ^ t2s ty ^ " has type " ^ t2s t)
  in
  match top_read f with
    | Atom ("quit") -> exit 0
    | Appl ("axiom", [Atom (name); body]) ->
      (* Printf.printf "ax: %s\n" (t2s ((sexpr_to_term numbers body))); *)
      cont name (sexpr_to_term numbers body)
    | Appl ("let", [Atom (name); body]) ->
      let body = sexpr_to_term numbers body in
      let ty = typ 0 types body in
      Hashtbl.add pre_typed body ty;
      Printf.printf "let %s\n" name;
      flush stdout;
      eval_loop (types, StringMap.add name body numbers)
    | Appl ("lemma", [Atom (name); body]) ->
      (* Printf.printf "lem: %s\n" (t2s ((sexpr_to_term numbers body)));  *)
      cont name (typ 0 types (sexpr_to_term numbers body))
    | Appl ("lemma", [Appl (name, [ty]); body]) ->
      let real_ty = typ 0 types (sexpr_to_term numbers body) in
      let ty = sexpr_to_term numbers ty in
      if eq real_ty ty then
        cont name real_ty
      else
        err ("the lamma " ^ name ^ " was expected to have type " ^ t2s ty ^ " while it has " ^
             t2s real_ty)
    | t ->
      err ("bad magic word " ^ s2s t)
;;

try 
  eval_loop (IntMap.empty, StringMap.empty)
with Failure s ->
  Printf.eprintf "FATAL: %s\n" s;
  exit 1