open Bindlib

let ( let* ) m f = Option.bind m f
let ( let+ ) m f = Option.map  f m
let ( let$ ) m f = match m with None -> f () | Some x -> Some x

type ('a, 'b) eq = Eq : ('a, 'a) eq

module type ReflexiveType = sig
  type 'a t
  val equal : 'a t -> 'b t -> ('a, 'b) eq option
end



module type Parameters = sig

  module Sort : ReflexiveType
  type 'a sort = 'a Sort.t

  type ('a, 'b) axiom
  val is_valid_axiom : 'a sort -> 'b sort -> ('a, 'b) axiom option

  type ('a, 'b, 'c) rule
  val is_valid_rule : 'a sort -> 'b sort -> 'c sort -> ('a, 'b, 'c) rule option

end



module Make (Parameters : Parameters) = struct
  include Parameters

  type term =
    | Sort        : 'n  sort -> term
    | Variable    : term var -> term
    | Application : application -> term
    | Function    : abstraction -> term
    | Product     : abstraction -> term

  and application =
    { func : term ; parameter : term }

  and abstraction =
    { binder : (term, term) binder ; parameter_typ : term }



  let rec equal (l : term) (r : term) =
    match l, r with
    | Sort l, Sort r -> Option.is_some (Sort.equal l r)
    | Variable l, Variable r -> eq_vars l r
    | Application l, Application r ->
      equal l.parameter r.parameter
      && equal l.func r.func
    | Function l, Function r ->
      equal l.parameter_typ r.parameter_typ
      && eq_binder equal l.binder r.binder
    | Product l, Function r ->
      equal l.parameter_typ r.parameter_typ
      && eq_binder equal l.binder r.binder
    | _, _ -> false

  let rec beta_reduction = function
    | Sort     s -> Sort     s
    | Variable x -> Variable x
    | Application { func = Function f ; parameter } -> subst f.binder parameter
    | Application app ->
      let func = beta_reduction app.func in
      let parameter = beta_reduction app.parameter in
      Application { func ; parameter }
    | Function f ->
      let parameter_typ = beta_reduction f.parameter_typ in
      let binder = binder_compose f.binder beta_reduction in
      Function { binder ; parameter_typ }
    | Product p ->
      let parameter_typ = beta_reduction p.parameter_typ in
      let binder = binder_compose p.binder beta_reduction in
      Product { binder ; parameter_typ }



  module Context = Map.Make (struct
    type t = term var
    let compare = compare_vars
  end)

  type context  = term Context.t
  type judgment = { context : context ; term : term ; typ : term }

  type universe = Universe : _ sort -> universe
  type checked  = Checked



  let rec check universe { context ; term ; typ } =
    let typ = beta_reduction typ in
    let* Universe _ = universe context typ in
    let judgment = { context ; term ; typ } in
    let$ () = axiom       universe judgment in
    let$ () = product     universe judgment in
    let$ () = abstraction universe judgment in
    let$ () = application universe judgment in
    weakening universe judgment

  and axiom _ { context ; term ; typ } =
    match term, typ with
    | Sort s1, Sort s2 when Context.is_empty context ->
      let+ _ = is_valid_axiom s1 s2 in
      Checked
    | _, _ -> None

  and start universe { context ; term ; typ } =
    match term, typ with
    | Variable x, a ->
      let* a' = Context.find_opt x context in
      let* () = if equal a a' then Some () else None in
      let context = Context.remove x context in
      let+ Universe _ = universe context a in
      Checked
    | _, _ -> None

  and weakening universe { context ; term ; typ } =
    let* x, b = Context.choose_opt context in
    let* Universe _ = universe context b in
    let context = Context.remove x context in
    check universe { context ; term ; typ }

  and product universe { context ; term ; typ } =
    match term, typ with
    | Product p, Sort s3 ->
      let* Universe s1 = universe context p.parameter_typ in
      let x, b = unbind p.binder in
      let context = Context.add x p.parameter_typ context in
      let* Universe s2 = universe context b in
      let+ _ = is_valid_rule s1 s2 s3 in
      Checked
    | _, _ -> None

  and application universe { context ; term ; typ } =
    match term, typ with
    | Application { func ; parameter }, app_typ ->
      let* product = match func with Product p -> Some p | _ -> None in
      let result_typ = subst product.binder parameter in
      let* () = if equal app_typ result_typ then Some () else None in
      let term = parameter and typ = product.parameter_typ in
      check universe { context ; term ; typ }
    | _, _ -> None

  and abstraction universe { context ; term ; typ } =
    match term, typ with
    | Function f, Product p ->
      let x, m = unbind f.binder in
      let b = subst p.binder (Variable x) in
      let* Universe _ = universe context (Product p) in
      let context = Context.add x f.parameter_typ context in
      check universe { context ; term = m ; typ = b }
    | _, _ -> None

end



module Inductive_Constructions_Parameters = struct

  type    zero = |
  type 'n succ = |
  type 'n nat = Zero : zero nat | Succ : 'n nat -> 'n succ nat

  let rec eq : type a b. a nat -> b nat -> (a, b) eq option = fun l r ->
    match l, r with
    | Zero  , Zero   -> Some Eq
    | Zero  , Succ _ -> None
    | Succ _, Zero   -> None
    | Succ l, Succ r ->
      match eq l r with
      | Some Eq -> Some Eq
      | None -> None

  type ('a, 'b) leq =
    | LTZero : (zero, 'a) leq
    | LTSucc : ('a, 'b) leq -> ('a succ, 'b succ) leq

  let rec leq : type a b. a nat -> b nat -> (a, b) leq option = fun l r ->
    match l, r with
    | Zero  , Zero   -> Some LTZero
    | Zero  , Succ _ -> Some LTZero
    | Succ _, Zero   -> None
    | Succ l, Succ r -> let+ prf = leq l r in LTSucc prf



  type    set  = |
  type    prop = |
  type 'n typ  = |

  type 's sort =
    | Set  : set  sort
    | Prop : prop sort
    | Type : 'n nat -> 'n typ sort

  type ('a, 'b) axiom =
    | Set_is_Type_0     : (set , zero typ) axiom
    | Prop_is_Type_0    : (prop, zero typ) axiom
    | Type_Start : (zero typ, zero succ typ) axiom
    | Type_Incr  : ('a typ, 'b typ) axiom -> ('a succ typ, 'b succ typ) axiom

  type ('a, 'b, 'c) rule =
    | Prop_Prop : (prop, prop, prop) rule
    | Set_Prop  : (set , prop, prop) rule
    | Prop_Set  : (prop, set , set ) rule
    | Set_Set   : (set , set , set ) rule
    | Type_Prop : ('n typ, prop, prop) rule
    | Type_Right : ('l, 'r) leq -> ('l typ, 'r typ, 'r typ) rule
    | Type_Left  : ('r, 'l) leq -> ('l typ, 'r typ, 'l typ) rule



  type ('a, 'b) is_valid_axiom =
    'a sort -> 'b sort -> ('a, 'b) axiom option

  let rec is_valid_axiom : type a b. (a, b) is_valid_axiom = fun l r ->
    match l, r with
    | Set , Type Zero -> Some Set_is_Type_0
    | Prop, Type Zero -> Some Prop_is_Type_0
    | Type Zero, Type Succ Zero -> Some Type_Start
    | Type Succ l, Type Succ r ->
      let+ prf = is_valid_axiom (Type l) (Type r) in
      Type_Incr prf
    | _, _ -> None

  type ('a, 'b,' c) is_valid_rule =
    'a sort -> 'b sort -> 'c sort -> ('a, 'b, 'c) rule option

  let rec is_valid_rule : type a b c. (a, b, c) is_valid_rule = fun a b c ->
    match a, b, c with
    | Prop, Prop, Prop -> Some Prop_Prop
    | Set , Prop, Prop -> Some Set_Prop
    | Prop, Set , Set  -> Some Prop_Set
    | Set , Set , Set  -> Some Set_Set
    | Type _, Prop, Prop -> Some Type_Prop
    | Type a, Type b, Type c ->
      begin match leq a b, eq b c, leq b a, eq a c with
        | Some leq, Some Eq, _, _ -> Some (Type_Right leq)
        | _, _, Some geq, Some Eq -> Some (Type_Left  geq)
        | _, _, _, _ -> None
      end
    | _, _, _ -> None



  type ('a, 'b) eq_sort = 'a sort -> 'b sort -> ('a, 'b) eq option

  let equal : type a b. (a, b) eq_sort = fun l r ->
    match l, r with
    | Set, Set -> Some Eq
    | Prop, Prop -> Some Eq
    | Type l, Type r ->
      begin match eq l r with
        | Some Eq -> Some Eq
        | None -> None
      end
    | _, _ -> None

  module Sort = struct
    type 'a t = 'a sort
    let equal = equal
  end

end



module Inductive_Constructions = Make (Inductive_Constructions_Parameters)