(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

open Univ
open UVars
open UnivConstraint

module G = AcyclicGraph.Make(struct
    type t = Level.t
    module Set = Level.Set
    module Map = Level.Map

    let equal = Level.equal
    let compare = Level.compare
    let root = Some Level.set

    let raw_pr = Level.raw_pr

    let anomaly_err u = Pp.(str"Universe " ++ Level.raw_pr u ++ str" undefined.")
  end)
(* Do not include G to make it easier to control universe specific
   code (eg add_universe with a constraint vs G.add with no
   constraint) *)

type t = {
  graph: G.t;
  type_in_type: bool;
  above_prop_qvars: Sorts.QVar.Set.t;
}

(* Universe inconsistency: error raised when trying to enforce a relation
   that would create a cycle in the graph of universes. *)

type path_explanation = G.explanation Lazy.t

type explanation =
  | Path of path_explanation
  | Other of Pp.t

type univ_inconsistency = Sorts.printer option * (UnivConstraint.kind * Sorts.t * Sorts.t * explanation option)

exception UniverseInconsistency of univ_inconsistency

type 'a check_function = t -> 'a -> 'a -> bool

let set_type_in_type b g = {g with type_in_type=b}

let type_in_type g = g.type_in_type

let check_smaller_expr g (u,n) (v,m) =
  let diff = n - m in
    match diff with
    | 0 -> G.check_leq g.graph u v
    | 1 -> G.check_lt g.graph u v
    | x when x < 0 -> G.check_leq g.graph u v
    | _ -> false

let exists_bigger g ul l =
  Universe.exists (fun ul' ->
    check_smaller_expr g ul ul') l

let real_check_leq g u v =
  Universe.for_all (fun ul -> exists_bigger g ul v) u

let check_leq g u v =
  type_in_type g || Universe.equal u v || (real_check_leq g u v)

let check_eq g u v =
  type_in_type g || Universe.equal u v ||
    (real_check_leq g u v && real_check_leq g v u)

let check_eq_level g u v =
  u == v || type_in_type g || G.check_eq g.graph u v

let empty_universes = {
  graph=G.empty;
  type_in_type=false;
  above_prop_qvars=Sorts.QVar.Set.empty;
}

let initial_universes = empty_universes

let initial_universes_with g = {g with graph=initial_universes.graph}

let enforce_constraint (u,d,v) g =
  match d with
  | Le -> G.enforce_leq u v g
  | Lt -> G.enforce_lt u v g
  | Eq -> G.enforce_eq u v g

let enforce_constraint0 cst g = match enforce_constraint cst g.graph with
| None -> None
| Some g' ->
  if g' == g.graph then Some g
  else Some { g with graph = g' }

let enforce_constraint cst g = match enforce_constraint0 cst g with
| None ->
  if not (type_in_type g) then
    let (u, c, v) = cst in
    let e = lazy (G.get_explanation cst g.graph) in
    let mk u = Sorts.sort_of_univ @@ Universe.make u in
    raise (UniverseInconsistency (None, (c, mk u, mk v, Some (Path e))))
  else g
| Some g -> g

let merge_constraints csts g = UnivConstraints.fold enforce_constraint csts g

let check_constraint { graph = g; type_in_type; _ } (u,d,v) =
  type_in_type
  || match d with
  | Le -> G.check_leq g u v
  | Lt -> G.check_lt g u v
  | Eq -> G.check_eq g u v

let check_constraints csts g = UnivConstraints.for_all (check_constraint g) csts

let is_above_prop ugraph q =
  Sorts.QVar.Set.mem q ugraph.above_prop_qvars

let check_eq_sort qeq univs s1 s2 =
  let u1 = Sorts.univ_of_sort s1 in
  let u2 = Sorts.univ_of_sort s2 in
  let q1 = Sorts.quality s1 in
  let q2 = Sorts.quality s2 in
  qeq q1 q2 && (type_in_type univs || check_eq univs u1 u2)

let check_leq_sort qeq univs s1 s2 =
  if type_in_type univs then
    let q1 = Sorts.quality s1 in
    let q2 = Sorts.quality s2 in
    let open Sorts.Quality in
    match q1, q2 with
    | QConstant QProp, QConstant QType -> true
    | _ -> qeq q1 q2
  else
    match s1, s2 with
    | (SProp, SProp) | (Prop, Prop) | (Set, Set) -> true
    | (Prop, (Set | Type _)) -> true
    | (Prop, VSort (q, _)) -> is_above_prop univs q
    | (Type _ | Set), (Set | Type _) -> check_leq univs (Sorts.univ_of_sort s1) (Sorts.univ_of_sort s2)
    | (GSort (s1, u1), GSort (s2, u2)) ->
      qeq (QGlobal s1) (QGlobal s2) && check_leq univs u1 u2
    | (VSort (s1, u1), VSort (s2, u2)) ->
      qeq (QVar s1) (QVar s2) && check_leq univs u1 u2
    | (GSort (s1, u1), VSort (s2, u2)) ->
      qeq (QGlobal s1) (QVar s2) && check_leq univs u1 u2
    | (VSort (s1, u1), GSort (s2, u2)) ->
      qeq (QVar s1) (QGlobal s2) && check_leq univs u1 u2
    | (VSort (q, u1), Type u2) -> is_above_prop univs q && check_leq univs u1 u2
    | ((SProp | Prop | Set | Type _ | GSort _ | VSort _), _) -> false

let leq_expr (u,m) (v,n) =
  let d = match m - n with
    | 1 -> Lt
    | diff -> assert (diff <= 0); Le
  in
  (u,d,v)

let enforce_leq_alg u v g =
  let open Util in
  let enforce_one (u,v) = function
    | Inr _ as orig -> orig
    | Inl (cstrs,g) as orig ->
      if check_smaller_expr g u v then orig
      else
        (let c = leq_expr u v in
         match enforce_constraint0 c g with
         | Some g -> Inl (UnivConstraints.add c cstrs,g)
         | None -> Inr (c, g))
  in
  (* max(us) <= max(vs) <-> forall u in us, exists v in vs, u <= v *)
  let c = List.map (fun u -> List.map (fun v -> (u,v)) (Universe.repr v)) (Universe.repr u) in
  let c = List.cartesians enforce_one (Inl (UnivConstraints.empty,g)) c in
  (* We pick a best constraint: smallest number of constraints, not an error if possible. *)
  let order x y = match x, y with
    | Inr _, Inr _ -> 0
    | Inl _, Inr _ -> -1
    | Inr _, Inl _ -> 1
    | Inl (c,_), Inl (c',_) ->
      Int.compare (UnivConstraints.cardinal c) (UnivConstraints.cardinal c')
  in
  match List.min order c with
  | Inl x -> x
  | Inr ((u, c, v), g) ->
    let e = lazy (G.get_explanation (u, c, v) g.graph) in
    let mk u = Sorts.sort_of_univ @@ Universe.make u in
    let e = UniverseInconsistency (None, (c, mk u, mk v, Some (Path e))) in
    raise e

exception AlreadyDeclared = G.AlreadyDeclared
let add_universe u ~strict g =
  let graph = G.add u g.graph in
  let d = if strict then Lt else Le in
  enforce_constraint (Level.set, d, u) { g with graph }

let check_declared_universes g l =
  G.check_declared g.graph l

let constraints_of_universes g =
  let add cst accu = UnivConstraints.add cst accu in
  G.constraints_of g.graph add UnivConstraints.empty
let constraints_for ~kept g =
  let add cst accu = UnivConstraints.add cst accu in
  G.constraints_for ~kept g.graph add UnivConstraints.empty

(** Instances *)

let check_eq_instances qeq univs t1 t2 =
  let qt1, ut1 = Instance.to_array t1 in
  let qt2, ut2 = Instance.to_array t2 in
  CArray.equal qeq qt1 qt2
  && CArray.equal (check_eq_level univs) ut1 ut2

let domain g = G.domain g.graph

let check_universes_invariants g = G.check_invariants ~required_canonical:Level.is_set g.graph

(** Pretty-printing *)

let pr_pmap sep pr map =
  let cmp (u,_) (v,_) = Level.compare u v in
  Pp.prlist_with_sep sep pr (List.sort cmp (Level.Map.bindings map))

let pr_arc prl = let open Pp in
  function
  | u, G.Node ltle ->
    if Level.Map.is_empty ltle then mt ()
    else
      prl u ++ str " " ++
      v 0
        (pr_pmap spc (fun (v, strict) ->
              (if strict then str "< " else str "<= ") ++ prl v)
            ltle) ++
      fnl ()
  | u, G.Alias v ->
    prl u  ++ str " = " ++ prl v ++ fnl ()

type node = G.node =
| Alias of Level.t
| Node of bool Level.Map.t

let repr g = G.repr g.graph

let pr_universes prl g = pr_pmap Pp.mt (pr_arc prl) g

open Pp

let explain_universe_inconsistency default_printer (printer, (o,u,v,p) : univ_inconsistency) =
  let printer = Option.default default_printer printer in
  let pr_uni u = match u with
  | Sorts.Set -> str "Set"
  | Sorts.Prop -> str "Prop"
  | Sorts.SProp -> str "SProp"
  | Sorts.Type u -> Universe.pr printer.pru u
  | Sorts.VSort _ | GSort _ -> Sorts.pr printer u
  in
  let pr_rel = function
    | Eq -> str"=" | Lt -> str"<" | Le -> str"<="
  in
  let reason = match p with
    | None -> mt()
    | Some (Other p) -> spc() ++ p
    | Some (Path p) ->
      let pstart, p = Lazy.force p in
      if p = [] then mt ()
      else
        str " because" ++ spc() ++ printer.pru pstart ++
        prlist (fun (r,v) -> spc() ++ pr_rel r ++ str" " ++ printer.pru v) p
  in
    str "Cannot enforce" ++ spc() ++ pr_uni u ++ spc() ++
      pr_rel o ++ spc() ++ pr_uni v ++ reason

module Internal = struct
  let add_template_qvars qvars g =
    assert (Sorts.QVar.Set.is_empty g.above_prop_qvars);
    {g with above_prop_qvars=qvars}

  let is_above_prop = is_above_prop
end