-- src/Iris/Algebra/COFESolver.lean

import Iris.Algebra.OFE

namespace Iris.COFE.OFunctor

open OFE

variable {F : ∀ α β [OFE α] [OFE β], Type u} [OFunctorContractive F]

variable [∀ α [COFE α], IsCOFE (F α α)]

variable [inh : Inhabited (F (ULift Unit) (ULift Unit))]

namespace Fix.Impl

variable (F) in

def A' : Nat → Σ α : Type u, COFE α
  | 0 => ⟨ULift Unit, inferInstance⟩
  | n+1 => let ⟨A, _⟩ := A' n; ⟨F A A, inferInstance⟩

variable (F) in

def A (n : Nat) : Type u := (A' F n).1

variable (F) in

def up : ∀ k, A F k -n> A F (k+1)
  | 0 => ⟨fun _ => inh.default, ⟨fun _ _ _ _ => .rfl⟩⟩
  | k+1 => map (down k) (up k)

def down : ∀ k, A F (k+1) -n> A F k
  | 0 => ⟨fun _ => ⟨()⟩, ⟨fun _ _ _ _ => .rfl⟩⟩
  | k+1 => map (up k) (down k)

end

  open OFunctorContractive in exact match k with
  | 0 => map_contractive.zero (x := (_, _)) (y := (_, _)) _ _
  | k+1 => map_contractive.succ (x := (_, _)) (y := (_, _)) _ ⟨up_down, up_down⟩ _

variable (F) in

@[ext] structure Tower : Type u where
  val k : A F k
  protected down {k} : down F k (val (k+1)) ≡ val k

variable (F) in

def upN {k} : ∀ n, A F k -n> A F (k + n)
  | 0 => .id
  | n+1 => (up F (k + n)).comp (upN n)

variable (F) in

def downN {k} : ∀ n, A F (k + n) -n> A F k
  | 0 => .id
  | n+1 => (downN n).comp (down F (k + n))

def eqToHom (e : i = k) : A F i -n> A F k := e ▸ .id

def embed : A F k -n> A F i :=
  if h : k ≤ i then (eqToHom (Nat.add_sub_cancel' h)).comp (upN ..)
  else (downN ..).comp (eqToHom (Nat.add_sub_cancel' (Nat.le_of_not_ge h)).symm)

protected def Tower.embed (k) : A F k -n> Tower F :=

-- Repository definitions used:
content: scoped notation:40 x " ≡{" n "}≡ " y:41 => OFE.Dist n x y

content: instance COFE.OFunctor.constOF_RFunctor [CMRA B] : RFunctor (constOF B) where
  map f g := by admit /- proof elided -/

content: instance OFE.ContractiveHom.fixpoint_ne [COFE α] [Inhabited α] :
    NonExpansive (ContractiveHom.fixpoint (α := α)) where
  ne n f1 f2 H := by admit /- proof elided -/

content: class IsCOFE (α : Type _) [OFE α] where
  compl : Chain α → α
  conv_compl {c : Chain α} : compl c ≡{n}≡ c n

content: structure Chain (α : Type _) [OFE α] where
  chain : Nat → α
  cauchy : n ≤ i → chain i ≡{n}≡ chain n

content: @[ext] structure Hom (α β : Type _) [OFE α] [OFE β] where
  f : α → β
  ne : NonExpansive f

content: @[ext] structure ContractiveHom (α β : Type _) [OFE α] [OFE β] extends Hom α β where
  [contractive : Contractive f]
  ne := ne_of_contractive f

content: class NonExpansive [OFE α] [OFE β] (f : α → β) where
  ne : ∀ ⦃n x₁ x₂⦄, x₁ ≡{n}≡ x₂ → f x₁ ≡{n}≡ f x₂

content: instance OFE.Option.some.ne [OFE α] : OFE.NonExpansive (some : α → Option α) := ⟨fun _ _ _ => id⟩

content: instance Option.merge_ne [OFE α] {op : α → α → α} [NonExpansive₂ op] :
    NonExpansive₂ (Option.merge op) where
  ne n x1 x2 Hx y1 y2 Hy := by admit /- proof elided -/

content: class NonExpansive₂ [OFE α] [OFE β] [OFE γ] (f : α → β → γ) where
  ne : ∀ ⦃n x₁ x₂⦄, x₁ ≡{n}≡ x₂ → ∀ ⦃y₁ y₂⦄, y₁ ≡{n}≡ y₂ → f x₁ y₁ ≡{n}≡ f x₂ y₂

content: protected def Hom.comp [OFE α] [OFE β] [OFE γ] (g : β -n> γ) (f : α -n> β) : α -n> γ where
  f := g.f ∘ f.f
  ne.1 _ _ _ h := g.ne.1 (f.ne.1 h)

content: scoped infix:40 " ≡ " => OFE.Equiv

content: scoped notation:40 x " ≡{" n "}≡ " y:41 => OFE.Dist n x y

content: @[inherit_doc]
infixr:25 " -n> " => Hom

content: infixr:25 " -c> " => ContractiveHom

content: @[inherit_doc]
infixr:25 " -C> " => Hom

-- Repository lemmas used:
content: theorem exact [BI PROP] (Q : PROP) : Q ⊢ Q

-- Local definitions used:
content: def A' : Nat → Σ α : Type u, COFE α
  | 0 => ⟨ULift Unit, inferInstance⟩
  | n+1 => let ⟨A, _⟩ := A' n; ⟨F A A, inferInstance⟩

content: def A (n : Nat) : Type u := (A' F n).1

content: def up : ∀ k, A F k -n> A F (k+1)
  | 0 => ⟨fun _ => inh.default, ⟨fun _ _ _ _ => .rfl⟩⟩
  | k+1 => map (down k) (up k)

content: def down : ∀ k, A F (k+1) -n> A F k
  | 0 => ⟨fun _ => ⟨()⟩, ⟨fun _ _ _ _ => .rfl⟩⟩
  | k+1 => map (up k) (down k)

content: @[ext] structure Tower : Type u where
  val k : A F k
  protected down {k} : down F k (val (k+1)) ≡ val k

content: def upN {k} : ∀ n, A F k -n> A F (k + n)
  | 0 => .id
  | n+1 => (up F (k + n)).comp (upN n)

content: def downN {k} : ∀ n, A F (k + n) -n> A F k
  | 0 => .id
  | n+1 => (downN n).comp (down F (k + n))

content: def eqToHom (e : i = k) : A F i -n> A F k := e ▸ .id

content: def embed : A F k -n> A F i :=
  if h : k ≤ i then (eqToHom (Nat.add_sub_cancel' h)).comp (upN ..)
  else (downN ..).comp (eqToHom (Nat.add_sub_cancel' (Nat.le_of_not_ge h)).symm)

content: protected def Tower.embed (k) : A F k -n> Tower F :=

-- Local lemmas used:
content: theorem down_up : ∀ {k} x, down F k (up F k x) ≡ x
  | 0, ⟨()⟩ => .rfl
  | _+1, _ => (map_comp ..).symm.trans <|
    (map_ne.eqv down_up down_up _).trans (map_id _)

content: theorem eqToHom_up {k k'} {x : A F k} (e : k = k') :
    eqToHom (congrArg Nat.succ e) (up F k x) = up F k' (eqToHom e x)

-- Target theorem and proof:
theorem Tower.embed_up (x : A F k) :
    Tower.embed (k+1) (up F k x) ≡ Tower.embed k x
:= by
  refine equiv_dist.2 fun n i => ?_
  dsimp [Tower.embed, embed]; split <;> rename_i h₁
  · simp [Nat.le_of_succ_le h₁]
    suffices ∀ a b (e₁ : k + 1 + a = i) (e₂ : k+b = i),
      eqToHom e₁ (upN F a (up F k x)) = eqToHom e₂ (upN F b x) from this .. ▸ .rfl
    rintro a b eq rfl
    rw [Nat.add_right_comm, Nat.add_assoc, Nat.add_left_cancel_iff] at eq; subst b
    show _ = up F (k + a) (upN F a x); clear h₁
    induction a with
    | zero => rfl
    | succ a ih =>
      dsimp [upN, Hom.comp]
      rw [eqToHom_up (by omega : k + 1 + a = k + (a + 1))]; congr 1; apply ih
  · split <;> rename_i h₂
    · cases Nat.le_antisymm h₂ (Nat.not_lt.1 h₁)
      have {a b} {e₁ : k+1 = k+a} {e₂ : k+b = k+0} :
        downN F a (eqToHom e₁ (up F k x)) ≡{n}≡ eqToHom e₂ (upN F b x) := by
        cases Nat.add_left_cancel e₁; cases Nat.add_left_cancel e₂
        exact (down_up ..).dist
      exact this
    · dsimp [Hom.comp]
      suffices ∀ a b (e₁ : k + 1 = i + a) (e₂ : k = i + b),
        downN F a (eqToHom e₁ (up F k x)) ≡{n}≡
        downN F b (eqToHom e₂ x) from this ..
      rintro a b eq rfl; cases Nat.add_left_cancel (m := b+1) eq
      exact (downN ..).ne.1 (down_up x).dist