\import Algebra.Meta
\import Category
\import Category.Adjoint
\import Category.Functor
\import Category.Limit
\import Data.Array
\import Data.Or
\import Equiv
\import Function.Meta
\import Logic
\import Logic.FirstOrder.Algebraic
\import Logic.FirstOrder.Algebraic.FreeModel
\import Logic.FirstOrder.Term
\import Logic.Meta
\import Meta
\import Paths
\import Paths.Meta
\import Set \hiding (#)
\import Set.SetCategory
\import Set.Fin

\record ModelHom {T : Theory} (Dom Cod : Model T) {
  | \coerce funcs {s : Sort} : Dom s -> Cod s
  | func-op {r : Sort} (h : Symb r) (d : DArray (\lam j => Dom (domain h j))) : funcs (operation h d) = operation h (\lam j => funcs (d j))
  | func-rel (P : PredSymb) (d : DArray (\lam j => Dom (predDomain P j))) : relation P d -> relation P (\lam j => funcs (d j))

  \lemma funcs-interpret {V : T -> \Set} {rho : Dom.Env V} {s : T} (t : Term V s) : funcs (Dom.interpret rho t) = Cod.interpret (\lam v => funcs (rho v)) t \elim t
    | var v => idp
    | apply f d => func-op f _ *> pmap (operation f) (exts \lam j => funcs-interpret (d j))
}

\instance ModelCat.{u} (T : Theory.{u,u,u}) : BicompleteCat.{u} (Model.{u} T)
  | Precat => ModelPrecat T
  | univalence => Cat.makeUnivalence $ later \lam (e : Iso) =>
    \have t => SIP (FamCat SetCat) (Model T) (\lam M N f => ModelHom M N f)
                   (\lam {E} {M} {N} (p1 : ModelHom M N) (p2 : ModelHom N M) => exts (\lam h d => p1.func-op h d, \lam P d => ext (p1.func-rel P d, p2.func-rel P d)))
        (\new Iso (\lam {j} => e.f {j}) (\lam {j} => e.hinv {j}) (path (\lam i {j} => (e.hinv_f @ i) {j})) (path (\lam i {j} => (e.f_hinv @ i) {j}))) e.dom e.cod e.f e.hinv
    \in (t., simp_coe t.)
  | limit {J} G => \new Limit {
    | apex => \new Model {
      | Structure => limitStructure T G
      | isModel S a rho h => propExt.conv (limitStructureTruth S.) (\lam j => isModel S a (\lam v => (rho v). j) (\lam k => propExt.dir (limitStructureTruth (S. k)) (h k) j))
    }
    | coneMap j => \new ModelHom {
      | funcs P => P. j
      | func-op _ _ => idp
      | func-rel P d h => h j
    }
    | coneCoh f => exts (\lam P => P. f)
    | limMap (c : Cone) => \new ModelHom {
      | funcs x => (\lam j => c.coneMap j x, \lam f => (c.coneCoh f @ __) x)
      | func-op h d => ext $ ext (\lam j => func-op h d)
      | func-rel P d x j => func-rel P d x
    }
    | limBeta c j => idp
    | limUnique p => exts (\lam e => exts (\lam j i => funcs {p j i} e))
  }
  | colimit G => \let data => \new ColimitData G \in \new Limit {
    | apex => Colimit
    | coneMap => colimitMap
    | coneCoh => colimitCone.coneCoh
    | isLimit => isColimit
  }
  \where {
    \open Structure
    \open Theory
    \open FreeStructure
    \open ColimitData \hiding (T,thExt,ColimitStr)

    \func emptyFin {S : \Set} : FinSet (\Sigma S Empty) \cowith
      | card => 
      | finEq => inP \new QEquiv {
        | f => \case __
        | ret => \case __.
        | ret_f => \case __
        | f_sec => \case __.
      }

    \instance ModelPrecat.{u} (T : Theory) : Precat (Model.{u} T)
      | Hom M N => ModelHom M N
      | id => \new ModelHom {
        | funcs x => x
        | func-op _ _ => idp
        | func-rel _ _ p => p
      }
      | o g f => \new ModelHom {
        | funcs x => g (f x)
        | func-op h d => rewrite (func-op {f} h d) (func-op {g} h _)
        | func-rel P d x => func-rel P (\lam j => f (d j)) (func-rel P d x)
      }
      | id-left => idp
      | id-right => idp
      | o-assoc => idp

    \func Forget.{u} {T : Theory} : Functor (ModelPrecat.{u} T) (FamPrecat SetCat.{u}) \cowith
      | F M => M
      | Func f => f.funcs
      | Func-id => idp
      | Func-o => idp

    \func ForgetAdjunction.{u} {T : Theory.{u,u,u}} : Adjunction.FromReflector (FamPrecat SetCat.{u}) (ModelPrecat.{u} T) Forget \cowith
      | LOb => FreeModel
      | etaMap x => qinj (var x)
      | eta-adjoint => inP \new QEquiv {
        | ret g => \new ModelHom {
          | funcs => qinterpret g
          | func-op h d => idp
          | func-rel P d (inP (d',Pd',d'd)) => transport (relation P) (exts \lam j => pmap (qinterpret g) (d'd j)) $ Model.theoremIsTrue Pd' g (\case __)
        }
        | ret_f f => exts \lam {s} qt => \case qinj-surj qt \with {
          | inP (t,t=qt) => rewriteI t=qt $ inv (f.funcs-interpret t) *> pmap (f {s}) (interpret=subst t *> pmap qinj (subst_var t))
        }
        | f_sec => idpe
      }

    \func ForgetToSetAdjunction.{u} {T : Theory.{u,u} (\Sigma)} : Adjunction SetCat.{u} (ModelPrecat.{u} T)
      => Adjunction.comp (later \new Adjunction.FromReflector {
        | RAdj {
          | F M => M ()
          | Func f => f
          | Func-id => idp
          | Func-o => idp
        }
        | LOb X _ => X
        | etaMap x => x
        | eta-adjoint => inP \new QEquiv {
          | ret f x => f x
          | ret_f => idpe
          | f_sec => idpe
        }
      }) ForgetAdjunction

    \class ColimitData.{u} \noclassifying {J : Precat.{u,u}} {T : Theory.{u,u,u}} (G : Functor J (ModelPrecat.{u} T)) {
      \func sigExt : Signature \cowith
        | Sort => T.Sort
        | Symb s => Or (T.Symb s) (Trunc0 (\Sigma (j : J) (G j s)))
        | domain => \case __ \with {
          | inl f => T.domain f
          | inr _ => nil
        }
        | PredSymb => T.PredSymb
        | predDomain => T.predDomain

      \func liftTerm {V : Sort -> \Set u} {s : Sort} (t : Term V s) : Term {sigExt} V s \elim t
        | var v => var v
        | apply f d => apply (inl f) (\new DArray { | at j => liftTerm (d j) })

      \func liftFormula {V : Sort -> \Set u} (phi : Formula V) : Formula {sigExt} V \elim phi
        | equality t t' => equality (liftTerm t) (liftTerm t')
        | predicate P d => predicate P (\new DArray { | at j => liftTerm (d j) })

      \func liftSequent (S : Sequent) : Sequent {sigExt}
        => (S., S., map liftFormula S., liftFormula S.)

      \instance thExt : Theory
        | Signature => sigExt
        | axioms S => OneOf (∃ (S' : Sequent) (T.axioms S') (S = liftSequent S'),
            ∃ (j : J) (s : Sort) (f : Symb s) (vars : DArray (\lam i => G j (domain f i))) (S = (\lam _ => Empty, emptyFin, nil, equality (apply (inl f) (\lam i => apply (inr (in0 (j, vars i))) nil)) (apply (inr (in0 (j, operation f vars))) nil))),
            ∃ (j : J) (P : PredSymb) (vars : DArray (\lam i => G j (predDomain P i))) (relation P vars) (S = (\lam _ => Empty, emptyFin, nil, predicate P (\new DArray { | at i => apply (inr (in0 (j, vars i))) nil }))),
            ∃ (j j' : J) (h : J.Hom j j') (s : Sort) (x : G j s) (S = (\lam _ => Empty, emptyFin, nil, equality (apply (inr (in0 (j', G.Func h x))) nil) (apply (inr (in0 (j, x))) nil))))

      \instance ColimitStr : Structure T
        | E => QTerm (\lam _ => Empty)
        | operation h args => qapply (inl h) args
        | relation P args => ∃ (args' : DArray (\lam i => Term (\lam _ => Empty) (predDomain P i)))
                               (IsTheorem nil (predicate P args'))
                               (\Pi (i : Fin (predDomain P).len) -> qinj (args' i) = args i)

      \lemma substInterpT {V : Sort -> \Set u} (rho : \Pi {s : Sort} -> V s -> Term (\lam _ => Empty) s) {s : Sort} (t : Term {T} V s)
        : ColimitStr.interpret (\lam v => qinj (rho v)) t = qinj (subst (liftTerm t) rho) \elim t
        | var v => idp
        | apply f d => unfold $ cong (ext (\lam j => substInterpT rho (d j))) *> inv (path (qmerge _))

      \lemma substInterpF {V : Sort -> \Set u} (rho : \Pi {s : Sort} -> V s -> Term (\lam _ => Empty) s) (phi : Formula {T} V)
        : IsFormulaTrue (\lam v => qinj (rho v)) phi <-> IsTheorem nil (substF (liftFormula phi) (\lam {s} v => rho {s} v))
        \elim phi
        | equality t t' => <->trans (<->_=. $ pmap2 (=) (substInterpT _ t) (substInterpT _ t')) qinj-equality
        | predicate P d => (\lam (inP (d',q,e)) => congruenceP (\lam j => qinj-equality. (e j *> substInterpT rho (d j))) q,
                            \lam q => inP (\new DArray { | at j => subst (liftTerm (d j)) (\lam {s} => rho) }, q, \lam i => inv (substInterpT rho (d i))))

      \func Colimit : Model T \cowith
        | Structure => ColimitStr
        | isModel S ax => \lam rho phiT =>
            \let | (inP (tau',p)) => S..liftDepSurj (\lam {v} => qinj-surj {_} {\lam _ => Empty} {v.}) (\lam v => rho v.)
                 | tau {s'} v => tau' (s',v)
                 | eq : \Pi (phi : Formula {T} S.) -> IsFormulaTrue rho phi <-> IsTheorem nil (substF (liftFormula phi) tau)
                      => rewrite (path (\lam i {s'} v => inv (p (s',v)) i)) (substInterpF tau)
            \in (eq S.). $ axiom (liftSequent S) (inP (, inP (S,ax,idp))) tau (\lam j => (eq (S. j)). (phiT j)) idp

      \func colimitMap (j : J) : ModelHom {T} (G j) Colimit \cowith
        | funcs x => qapply (inr (in0 (j,x))) nil
        | func-op {s} h d =>
            \have args => \new DArray _ (\lam k => apply (inr (in0 (j, d k))) nil)
            \in inv (path (qmerge nil)) *> inv (path (qquot (axiom (\lam _ => Empty, emptyFin, nil, _) (inP (, inP (j,s,h,_,idp))) (\lam {_} => absurd) (\case __) idp))) *>
                path (qmerge args) *> path (\lam i => qapply (inl h) (\new DArray _ (\lam k => qmerge {_} {_} {_} {inr (in0 (j, d k))} nil i)))
        | func-rel P d r => inP (\new DArray { | at i => apply (inr (in0 (j, d i))) nil }, axiom (\lam _ => Empty, emptyFin, nil, _) (inP (, inP (j,P,d,r,idp))) (\lam {_} => absurd) (\case __) idp, \lam i => path (qmerge _))

      \func colimitCone : Cone G.op \cowith
        | apex => Colimit
        | coneMap => colimitMap
        | coneCoh {j} {j'} h => exts (\lam {s} x => unfold (inv (path (qmerge nil)) *> path (qquot (axiom _ (inP (, inP (j',j,h,s,x,idp))) (absurd __) (\case __) idp)) *> path (qmerge nil)))

      \func strExt {M : Model.{u} T} (C : Cone G.op M) : Structure thExt M \cowith
        | operation => \case \elim __ \with {
          | inl h => operation h
          | inr (in0 (j,x)) => \lam _ => C.coneMap j x
        }
        | relation => relation

      \lemma strExtT {V : Sort -> \Set u} (C : Cone G.op) {rho : Env {C} V} {s : Sort} (t : Term {T} V s)
        : interpret {strExt C} rho (liftTerm t) = interpret {C} rho t \elim t
        | var v => idp
        | apply f d => unfold $ cong $ ext (\lam j => strExtT C (d j))

      \lemma strExtF {V : Sort -> \Set u} (C : Cone G.op) {rho : Env {C} V} (phi : Formula {T} V)
        : IsFormulaTrue {strExt C} rho (liftFormula phi) = IsFormulaTrue {C} rho phi \elim phi
        | equality t t' => pmap2 (=) (strExtT C t) (strExtT C t')
        | predicate P d => cong (ext (\lam j => strExtT C (d j)))

      \func modelExt {M : Model.{u} T} (C : Cone G.op M) : Model thExt M \cowith
        | Structure => strExt C
        | isModel => \case \elim __, __ \with {
          | _, inP (, inP (S,a,idp)) => \lam rho p => propExt.conv (strExtF C S.) (isModel S a rho (\lam j => propExt.dir (strExtF C (S. j)) (p j)))
          | _, inP (, inP (j,s,f,d,idp)) => \lam rho _ => inv (func-op f d)
          | _, inP (, inP (j,P,d,r,idp)) => \lam rho _ => func-rel P d r
          | _, inP (, inP (j,j',h,s,x,idp)) => \lam rho _ => cong (path (\lam i => funcs {C.coneCoh h  @ i} x))
        }

      \func modelExtHom {M : Model.{u} T} (C : Cone G.op M) : ModelHom {T} Colimit C \cowith
        | funcs => qinterpret0 {thExt} {modelExt C}
        | func-op h d => idp
        | func-rel P d (inP (d',t,p)) => coe (\lam i => relation P (\new DArray _ (\lam j => qinterpret0 {thExt} {modelExt C} (p j @ i)))) (Model.theoremIsTrue {modelExt C} t (\lam {s} => absurd {modelExt C s}) (\case __)) right

      \lemma interpret=map {M : Model.{u} T} (f : ModelHom {T} Colimit M) {s : Sort} (t : Term (\lam _ => Empty) s)
        : Model.interpret {modelExt (conePullback {_} {_} {G.op} colimitCone M f)} (\lam {_} => absurd) t = f (qinj t) \elim t
        | var ()
        | apply (inl h) d => unfold (cong (ext (\lam j => interpret=map f (d j)))) *> inv (f.func-op h (\new DArray _ (\lam j => qinj (d j)))) *> pmap (f __) (inv (path (qmerge d)))
        | apply (inr (in0 (j,x))) nil => pmap (f __) (inv (path (qmerge nil)))

      \lemma qinterpret=map {M : Model.{u} T} (f : ModelHom {T} Colimit M) {s : Sort} (t : QTerm (\lam _ => Empty) s) : qinterpret0 {thExt} {modelExt (conePullback {_} {_} {G.op} colimitCone M f)} t = f t \elim t
        | qinj t => interpret=map f t
        | qapply (inl h) d => unfold qinterpret0 (unfold $ cong $ ext \lam j => qinterpret=map f (d j)) *> inv (f.func-op h d)
        | qapply (inr (in0 (j,x))) d => idp

      \lemma isColimit (M : Model.{u} T) : IsEquiv (conePullback {_} {_} {G.op} colimitCone M) => inP \new QEquiv {
        | ret => modelExtHom
        | ret_f f => exts (qinterpret=map f)
        | f_sec C => exts (\lam j => exts (\lam x => idp))
      }
    }

    \func limitStructure.{u} (T : Theory.{u,u,u}) {J : Precat.{u,u}} (G : Functor J (ModelPrecat.{u} T)) : Structure T \cowith
      | E s => \Sigma (P : \Pi (j : J) -> G j s) (\Pi {j j' : J} (h : Hom j j') -> G.Func h (P j) = P j')
      | operation h d => (\lam j => operation h (\lam k => (d k). j), \lam {j} {j'} f => func-op {G.Func f} h (\lam k => (d k). j) *> cong (ext (\lam k => (d k). f)))
      | relation P d => \Pi (j : J) -> relation P (\lam k => (d k). j)

    \lemma limitStructureInterp.{u} {T : Theory.{u,u,u}} {J : Precat.{u,u}} {G : Functor J (ModelPrecat.{u} T)} {V : Sort -> \Set u} {rho : Env {limitStructure T G} V} {s : Sort} (t : Term V s) {j : J}
      : (interpret rho t). j = interpret (\lam v => (rho v). j) t \elim t
      | var v => idp
      | apply f d => unfold $ cong $ ext (\lam k => limitStructureInterp (d k))

    \lemma limitStructureTruth.{u} {T : Theory.{u,u,u}} {J : Precat.{u,u}} {G : Functor J (ModelPrecat.{u} T)} {V : Sort -> \Set u} {rho : Env {limitStructure T G} V} (phi : Formula V)
      : IsFormulaTrue rho phi = (\Pi (j : J) -> IsFormulaTrue (\lam v => (rho v). j) phi) \elim phi
      | equality t t' => ext (\lam ft j => inv (limitStructureInterp t) *> rewrite ft idp *> limitStructureInterp t', \lam ft => exts (\lam j => limitStructureInterp t *> ft j *> inv (limitStructureInterp t')))
      | predicate P d => ext (\lam ft j => transport (relation P) (exts (\lam k => limitStructureInterp (d k))) (ft j), \lam ft j => transport (relation P) (exts (\lam k => inv $ limitStructureInterp (d k))) (ft j))
}