From mathcomp Require Import ssreflect.all_ssreflect.
Require Import Base.axioms Base.fintype Framework.graded.
Set Implicit Arguments.
Unset Strict Implicit.

Local Ltac asimpl := repeat progress (cbn; fsimpl; gsimpl).

Module TyTerms.
  Section Defs.

  Inductive ty (n : nat) :=
  | var (i : fin n) : ty n
  | arr (A B : ty n) : ty n
  | all (A : ty n.+1) : ty n.

  Structure traversal (V D : nat -> Type) := Traversal
    { tvar : forall n, V n -> D n
    ; tarr : forall n, D n -> D n -> D n
    ; tall : forall m, ☐ (fun m' => V m' -> D m') m -> D m
    }.
  Arguments tvar {V D} t {n} i.
  Arguments tarr {V D} t {n} A B.
  Arguments tall {V D} t {m} F.

  Fixpoint eval {V : graded1} {D : nat -> Type} (t : traversal V D)
           {m n : nat} (rho : fin m -> V n) (A : ty m) : D n :=
    match A with
    | var i => tvar t (rho i)
    | arr A B => tarr t (eval t rho A) (eval t rho B)
    | all A => tall t (fun _ xi v => eval t (v .: rho >> th1 xi) A)
    end.

  Definition ren_traversal : traversal fin ty :=
    Traversal (@var) (@arr) (fun m F => all (F m.+1 shift bound)).
  Local Notation ty_ren := (eval ren_traversal).
  Local Notation "s ◁ f" := (ty_ren f s) (at level 50).
  Canonical ty_graded1 := Eval hnf in Graded1 ty (@eval _ _ ren_traversal).

  Lemma th1_tyE {m n} (f : iren m n) (A : ty m) : A ⋅ f = A ◁ f.
  Proof. by []. Qed.

  Local Ltac gsimpl' := rewrite ?th1_tyE; gsimpl.

  Section EvalRen.
    Context {V : graded1} {D : nat -> Type} (t : traversal V D).

    Theorem eval_ren {m n k : nat} (A : ty m) (f : ren m n) (rho : fin n -> V k) :
      eval t rho (A ◁ f) = eval t (f >> rho) A.
    Proof.
      elim: A n k f rho => {m}//=[m A ih1 B ih2|m A ih] n k f rho.
        by rewrite ih1 ih2. congr tall; fext=> l g i. rewrite ih. by fsimpl.
    Qed.
  End EvalRen.

  Lemma ren_comp {m n k : nat} (f : ren m n) (g : ren n k) (A : ty m) :
    A ◁ f ◁ g = A ◁ (f >> g).
  Proof. exact: eval_ren. Qed.

  Lemma ren_inj {m n : nat} (f : ren m n) : injective f -> injective (ty_ren f).
  Proof.
    move=> H A B. elim: A B n f H => {m}[m i|m A1 ih1 A2 ih2|m A ih] [j|B1 B2|B] n f H //=[].
    - by move=>/H->.
    - move=>/ih1 e1 /ih2 e2. by rewrite e1 // e2.
    - move=>/ih e. rewrite e //. by hnf=> /= -[i|] [j|] //=[/H->].
  Qed.

  Lemma ren_id {n : nat} (A : ty n) : A ◁ id = A.
  Proof. apply: (ren_inj (f := id)) => //. exact: ren_comp. Qed.

  Canonical ty_cgraded1 := Eval hnf in CGraded1 ty (@ren_comp).
  Canonical ty_igraded1 := Eval hnf in IGraded1 ty (@ren_id).
  Canonical ty_pgraded1 := Eval hnf in PGraded1 ty (fun m n f => ren_inj (@iren_inj m n f)).

  Definition lift {V D : graded1} (t : traversal V D) : traversal D D :=
    Traversal (fun n => id) (@tarr _ _ t)
              (fun m F => tall t (fun n f v => F n f (tvar t v))).

  Notation inst_traversal := (lift ren_traversal).
  Notation inst := (eval inst_traversal).
  Notation "A .[ f ]" := (inst f A) : ty_inst_scope.
  Delimit Scope ty_inst_scope with ty.

  Section Embedding.
    Context {V D : graded1} (t : traversal V D).

    Hypothesis var_is_natural : forall m n (f : iren m n) (i : V m),
        (tvar t i) ⋅ f = tvar t (i ⋅ f).

    Theorem lift_embed' {m n : nat} (rho : fin m -> V n) (A : ty m) :
      eval t rho A = eval (lift t) (rho >> tvar t) A.
    Proof.
      elim: A n rho => {m} //=[m A ihA B ihB|m A ih] n rho.
      - by rewrite ihA ihB.
      - congr tall; fext => k xi v. rewrite ih. congr eval.
        fsimpl=>/=. congr scons. fext=> i /=. symmetry.
        exact: var_is_natural.
    Qed.
  End Embedding.

  Structure is_natural {V D : graded1} (t : traversal V D) := IsNatural
    { nvar : forall m n (f : iren m n) (i : V m),
      (tvar t i) ⋅ f = tvar t (i ⋅ f)
    ; narr : forall m n (f : iren m n) (A B : D m),
      (tarr t A B) ⋅ f = tarr t (A ⋅ f) (B ⋅ f)
    ; nall : forall m (F : ☐ (fun m' => V m' -> D m') m),
      (forall n k (f : iren m n) (g : iren n k) (i : V n),
         (F n f i) ⋅ g = F k (f >>> g) (i ⋅ g)) ->
       forall n (f : iren m n),
        (tall t F) ⋅ f = tall t (F ⋅ f)
    }.

  Structure model (V D : graded1) := Model
    { base : traversal V D
    ; _ : is_natural base
    }.
  Local Coercion base : model >-> traversal.

  Definition model_is_natural {V D : graded1} (t : model V D) : is_natural t :=
    let: Model _ c := t in c.
  Local Coercion model_is_natural : model >-> is_natural.

  Lemma naturality {V : cgraded1} {D : graded1} (t : model V D)
           {m n k : nat} (rho : fin m -> V n) (f : iren n k) (A : ty m) :
    (eval t rho A) ⋅ f = eval t (rho >> th1 f) A.
  Proof.
    elim: A n k rho f => {m}[m i|m A ihA B ihB|m A ih] n k rho f.
    - exact: (nvar t).
    - cbn. rewrite (narr t). congr tarr. exact: ihA. exact: ihB.
    - cbn. rewrite (nall t).
      + congr tall; fext=> l g v. by asimpl.
      + move=> {k f} k l f g v. rewrite ih. by asimpl.
  Qed.

  Theorem eval_inst {V : cgraded1} {D : graded1} (t : model V D)
           {m n k : nat} (sigma : fin m -> ty n) (rho : fin n -> V k) (A : ty m) :
    eval t rho A.[sigma]%ty =
    eval (lift t) (sigma >> eval t rho) A.
  Proof.
    elim: A n k sigma rho => {m}//=[m A ihA B ihB|m A ih] n k sigma rho.
    - by rewrite ihA ihB.
    - congr tall; fext=> l f v. rewrite ih. asimpl. congr eval.
      fext=>/=;case=>//=i. rewrite eval_ren /=. symmetry. exact: naturality.
  Qed.

  Definition lift_is_natural {V D : graded1} (t : model V D) :
    is_natural (lift t).
  Proof.
    apply: IsNatural => //[m n f A B/=|m F FP n f]. exact: (narr t).
    apply: (nall t) => {n f} n k f g v. by rewrite FP (nvar t).
  Qed.

  Canonical lift_model {V D : graded1} (t : model V D) : model D D :=
    Eval hnf in Model (lift_is_natural t).

  Lemma lift_embed {V D : graded1} (t : model V D)
      {m n : nat} (rho : fin m -> V n) (A : ty m) :
    eval t rho A = eval (lift t) (rho >> tvar t) A.
  Proof. apply: lift_embed'. exact: (nvar t). Qed.

  Definition ty_ren_inst {m n k : nat} (A : ty m) (f : ren m n) (g : fin n -> ty k) :
    (A ◁ f).[g]%ty = A.[f >> g]%ty.
  Proof. exact: eval_ren. Qed.

  Lemma ren_is_natural : is_natural ren_traversal.
  Proof.
    apply: IsNatural => //= m F H n f. gsimpl' => /=.
      by rewrite -shift_up -(H m.+1 n.+1 shift (up f) bound).
  Qed.
  Canonical ren_model : model fin_graded1 ty_graded1 :=
    Eval hnf in Model ren_is_natural.

  Corollary ren_is_inst {m n} (A : ty m) (f : ren m n) :
    A ◁ f = A.[f >> @var n]%ty.
  Proof. exact: lift_embed. Qed.

  Corollary inst_id {n : nat} (A : ty n) :
    A.[@var n]%ty = A.
  Proof.
    by rewrite -{2}[A]th_id1/= th1_tyE ren_is_inst.
  Qed.

  Definition inst_comp {m n k : nat} (A : ty m) (f : fin m -> ty n) (g : fin n -> ty k) :
    A.[f]%ty.[g]%ty = A.[f >> inst g]%ty.
  Proof. exact: eval_inst. Qed.

  Lemma ty_inst_ren {m n k : nat} (A : ty m) (f : fin m -> ty n) (g : ren n k) :
    A.[f]%ty ◁ g = A.[f >> ty_ren g]%ty.
  Proof.
    rewrite lift_embed inst_comp. do 2 f_equal. symmetry. fext. exact: lift_embed.
  Qed.

  Structure smodel (V D : Type) := SModel
    { svar : V -> D
    ; sarr : D -> D -> D
    ; sall : (V -> D) -> D
    }.
  Arguments SModel {V D} svar sarr sall.

  Definition smodel_to_traversal (V D : Type) (t : smodel V D) :
    traversal (const1 V) (const1 D) :=
    {| tvar := fun _ => svar t
     ; tarr := fun _ => sarr t
     ; tall := fun n F => sall t (fun v => F n idren v)
    |}.
  Local Coercion smodel_to_traversal : smodel >-> traversal.

  Lemma smodel_is_natural {V D : Type} (t : smodel V D) : is_natural t.
  Proof. apply: IsNatural => //= m F FP n f. congr sall. fext=> v. exact: FP. Qed.

  Canonical smodel_to_model {V D : Type} (t : smodel V D) :
    model (const_graded1 V) (const_graded1 D) :=
    Eval hnf in Model (smodel_is_natural t).

  Definition lift_smodel {V D : Type} (t : smodel V D) : smodel D D :=
    {| svar := id
     ; sarr := sarr t
     ; sall := fun F => sall t (fun x => F (svar t x))
    |}.

  Definition seval {V D : Type} (t : smodel V D) {n : nat} (rho : fin n -> V) (A : ty n) :=
    @eval _ _ t n 1 rho A.

  (* not needed but nice to know *)
  Corollary seval_irrelevance {V D : Type} (t : smodel V D)
      {m n : nat} (A : ty m) (rho : fin m -> V) :
    eval (n := n) t rho A = seval t rho A.
  Proof.
    case: n.
    - refine (@naturality _ _ (smodel_to_model t) m 0 1 rho _ A).
      exact: IRen (fun (f : fin 0) => match f with end) _.
    - symmetry. apply: (naturality (smodel_to_model t) rho). apply (IRen (fun=> bound)).
      by move=>/=[[]|]//[[]|]//.
  Qed.

  Corollary seval_ren {V D : Type} {t : smodel V D} {m n : nat}
      (A : ty m) (f : ren m n) (rho : fin n -> V) :
    seval t rho (A ◁ f) = seval t (f >> rho) A.
  Proof. exact: eval_ren. Qed.

  Corollary seval_inst {V D : Type} {t : smodel V D} {m n : nat}
      (A : ty m) (sigma : fin m -> ty n) (rho : fin n -> V) :
    seval t rho A.[sigma]%ty =
    seval (lift_smodel t) (sigma >> seval t rho) A.
  Proof. exact: eval_inst. Qed.

  End Defs.

  Module Exports.
    Notation ty := ty.
    Canonical ty_graded1.
    Canonical ty_cgraded1.
    Canonical ty_igraded1.
    Canonical ty_pgraded1.

    Canonical ren_model.
    Canonical lift_model.
    Canonical smodel_to_model.

    Coercion base : model >-> traversal.
    Coercion model_is_natural : model >-> is_natural.
    Coercion smodel_to_traversal : smodel >-> traversal.

    Arguments tvar {V D} t {n} i.
    Arguments tarr {V D} t {n} A B.
    Arguments tall {V D} t {m} F.

    Arguments var {n}, n.

    Module Ty.
      (* Types *)
      Notation var := var.
      Notation arr := arr.
      Notation all := all.

      (* Traversal *)
      Notation traversal := traversal.
      Notation Traversal := Traversal.
      Notation tvar := tvar.
      Notation tarr := tarr.
      Notation tall := tall.

      Notation ren_traversal := ren_traversal.
      Notation inst_traversal := (lift ren_traversal).
      Notation lift := lift.

      (* Models *)
      Notation is_natural := is_natural.
      Notation IsNatural := IsNatural.
      Notation nvar := nvar.
      Notation narr := narr.
      Notation nall := nall.

      Notation model := model.
      Notation Model := Model.

      (* Evaluation *)
      Notation eval := eval.

      Notation eval_ren := eval_ren.
      Notation eval_inst := eval_inst.
      Notation naturality := naturality.

      Notation embed := lift_embed.

      (* Renaming *)
      Notation ren := (eval ren_traversal).

      Notation ren_id := ren_id.
      Notation ren_comp := ren_comp.
      Notation ren_inj := ren_inj.

      (* Instantiation *)
      Notation inst := (eval inst_traversal).

      Notation ren_inst := ty_ren_inst.
      Notation inst_ren := ty_inst_ren.

      Notation inst_id := inst_id.
      Notation inst_comp := inst_comp.

      (* Simple models *)
      Notation smodel := smodel.
      Notation SModel := SModel.
      Notation svar := svar.
      Notation sarr := sarr.
      Notation sall := sall.

      Notation seval := seval.
      Notation seval_irrelevance := seval_irrelevance.
      Notation seval_ren := seval_ren.
      Notation seval_inst := seval_inst.
    End Ty.
    Notation "A ◁ f" := (Ty.ren f A) (at level 50) : ty_inst_scope.
    Notation "A .[ f ]" := (Ty.inst f A) : ty_inst_scope.
    Delimit Scope ty_inst_scope with ty.
  End Exports.
End TyTerms.
Export TyTerms.Exports.

(* Automation *)
Lemma ty_ren_idX {n} : Ty.ren (id : fin n -> fin n) = id.
Proof. fext=> A. exact: Ty.ren_id. Qed.

Lemma ty_ren_compX {m n k} (f : ren m n) (g : ren n k) :
  Ty.ren f >> Ty.ren g = Ty.ren (f >> g).
Proof. fext=> A /=. exact: Ty.ren_comp. Qed.

Lemma ty_ren_compR {T m n k} (f : ren m n) (g : ren n k) (h : _ -> T) :
  Ty.ren f >> (Ty.ren g >> h) = Ty.ren (f >> g) >> h.
Proof. fext=> A /=. by rewrite Ty.ren_comp. Qed.

Lemma ty_ren_instX {m n k} (f : ren m n) (g : fin n -> ty k) :
  Ty.ren f >> Ty.inst g = Ty.inst (f >> g).
Proof. fext=> A /=. exact: Ty.ren_inst. Qed.

Lemma ty_ren_instR {T m n k} (f : ren m n) (g : fin n -> ty k) (h : _ -> T) :
  Ty.ren f >> (Ty.inst g >> h) = Ty.inst (f >> g) >> h.
Proof. fext=> A /=. by rewrite Ty.ren_inst. Qed.

Lemma ty_inst_renX {m n k} (f : fin m -> ty n) (g : ren n k) :
  Ty.inst f >> Ty.ren g = Ty.inst (f >> Ty.ren g).
Proof. fext=> A /=. exact: Ty.inst_ren. Qed.

Lemma ty_inst_renR {T m n k} (f : fin m -> ty n) (g : ren n k) (h : _ -> T) :
  Ty.inst f >> (Ty.ren g >> h) = Ty.inst (f >> Ty.ren g) >> h.
Proof. fext=> A /=. by rewrite Ty.inst_ren. Qed.

Lemma ty_id_instX {m n} (f : fin m -> ty n) : Ty.var >> Ty.inst f = f.
Proof. by []. Qed.

Lemma ty_inst_idX {n} : Ty.inst (@Ty.var n) = id.
Proof. fext=> A. exact: Ty.inst_id. Qed.

Lemma ty_inst_idR {T m n} (f : fin m -> ty n) (g : _ -> T) : Ty.var >> (Ty.inst f >> g) = (f >> g).
Proof. by []. Qed.

Lemma ty_inst_compX {m n k} (f : fin m -> ty n) (g : fin n -> ty k) :
  Ty.inst f >> Ty.inst g = Ty.inst (f >> Ty.inst g).
Proof. fext=> A /=. exact: Ty.inst_comp. Qed.

Lemma ty_inst_compR {T m n k} (f : fin m -> ty n) (g : fin n -> ty k) (h : _ -> T) :
  Ty.inst f >> (Ty.inst g >> h) = Ty.inst (f >> Ty.inst g) >> h.
Proof. fext=> A /=. by rewrite Ty.inst_comp. Qed.

Ltac ty_simpl :=
  rewrite ?ty_ren_idX ?ty_ren_compX ?ty_ren_compR
          ?ty_ren_instX ?ty_ren_instR
          ?ty_inst_renX ?ty_inst_renR
          ?ty_id_instX ?ty_inst_idX ?ty_inst_idR
          ?ty_inst_compX ?ty_inst_compR
          ?Ty.ren_id ?Ty.ren_comp ?Ty.ren_inst ?Ty.inst_ren ?Ty.inst_id ?Ty.inst_comp
          ?Ty.eval_ren ?Ty.eval_inst ?Ty.naturality ?Ty.seval_ren ?Ty.seval_inst
          ?TyTerms.th1_tyE.