import Mathlib

section

variable {ι : Type*} [Preorder ι] {M : ι → Type*} {P : Type*}
  (f : (i j : ι) → (h : i ≤ j) → M i → M j) (g : ∀ i, M i → P)

@[mk_iff] class IsDirectLimit [DirectedSystem M f] : Prop where
  surj : ∀ m : P, ∃ i, ∃ mi : M i, g i mi = m
  inj :  ∀ i j, ∀ mi : M i, ∀ mj : M j, g i mi = g j mj → ∃ (k : ι) (hik : i ≤ k) (hjk : j ≤ k),
      f i k hik mi = f j k hjk mj
  compatibility : ∀ i j hij x, g j (f i j hij x) = g i x

variable [DirectedSystem M f] [IsDirectLimit f g]

namespace IsDirectLimit

theorem compatibility' {i j hij} (x : M i) : g j (f i j hij x) = g i x :=
  IsDirectLimit.compatibility i j hij x

theorem is_injective {i j} {mi : M i} {mj : M j} (h : g i mi = g j mj) :
    ∃ (k : ι) (hik : i ≤ k) (hjk : j ≤ k), f i k hik mi = f j k hjk mj :=
  IsDirectLimit.inj i j mi mj h

include f in
theorem is_surjective : ∀ m : P, ∃ i, ∃ mi : M i, g i mi = m :=
  IsDirectLimit.surj f

noncomputable def preimage_index (p : P) : ι := (is_surjective f g p).choose

noncomputable def preimage (p : P) : M (preimage_index f g p) :=
  (is_surjective f g p).choose_spec.choose

theorem image_preimage (p : P) :
    g (preimage_index f g p) (preimage f g p) = p :=
  (is_surjective f g p).choose_spec.choose_spec

variable {P₁ P₂ : Type*} (g₁ : ∀ i, M i → P₁) (g₂ : ∀ i, M i → P₂)

noncomputable def lift [IsDirectLimit f g₁] (p : P₁) : P₂ :=
  g₂ (preimage_index f g₁ p) (preimage f g₁ p)

@[simp]
theorem lift_of [IsDirectLimit f g₁] (Hg : ∀ i j hij x, g₂ j (f i j hij x) = g₂ i x) {i} (x) :
    (lift f g₁ g₂) (g₁ i x) = g₂ i x := by
  dsimp [lift]
  have ⟨k, hpk, hik, h_eq⟩ := is_injective f g₁ (image_preimage f g₁ (g₁ i x))
  rw [← Hg i k hik x, ← Hg (preimage_index f g₁ (g₁ i x)) k hpk, h_eq]

noncomputable def Equiv [h₁ : IsDirectLimit f g₁] [h₂ : IsDirectLimit f g₂] :
    P₁ ≃ P₂ where
  toFun := lift f g₁ g₂
  invFun := lift f g₂ g₁
  left_inv x := by
    obtain ⟨i, mi, hmi⟩ := h₁.surj x
    rw [← hmi]
    simp only [compatibility', implies_true, lift_of]
  right_inv x := by
    obtain ⟨i, mi, hmi⟩ := h₂.surj x
    rw [← hmi]
    simp only [compatibility', implies_true, lift_of]

@[simp]
lemma Equiv_apply (i : ι) (x : M i) [IsDirectLimit f g₁] [IsDirectLimit f g₂] :
  Equiv f g₁ g₂ (g₁ i x) = g₂ i x := by
  simp [Equiv, Equiv.coe_fn_mk, compatibility', implies_true, lift_of]

@[simp]
lemma linearEquiv_symm_apply (i : ι) (x : M i) [IsDirectLimit f g₁]
  [IsDirectLimit f g₂] : (Equiv f g₁ g₂).symm (g₂ i x) = g₁ i x := by
  simp [Equiv, compatibility', implies_true, lift_of]

namespace Module

variable {R : Type*} [Semiring R] [∀ i, AddCommMonoid (M i)] [∀ i, Module R (M i)]
  [AddCommMonoid P] [Module R P] (f : (i j : ι) → i ≤ j → M i →ₗ[R] M j)
  [DirectedSystem M (f · · ·)] [IsDirectLimit (fun ⦃i j⦄ h => f i j h) g] [IsDirected ι (· ≤ ·)]

instance [DecidableEq ι] [Nonempty ι] :
    IsDirectLimit (f · · ·) (Module.DirectLimit.of R ι M f ·) where
  surj := Module.DirectLimit.exists_of
  inj i j mi mj h := by
    apply_fun Module.DirectLimit.linearEquiv _ _ at h
    simp_rw [Module.DirectLimit.linearEquiv_of, Quotient.eq] at h
    exact h
  compatibility i j hij x := Module.DirectLimit.of_f

variable [AddCommMonoid P₁] [Module R P₁] [AddCommMonoid P₂] [Module R P₂]
  (g₁ : ∀ i, M i →ₗ[R] P₁) (g₂ : ∀ i, M i →ₗ[R] P₂) [h₁ : IsDirectLimit (f · · ·) (g₁ · ·)]

omit [IsDirected ι (· ≤ ·)] in
theorem compatibility_module {i j hij} (x : M i) : g₁ j (f i j hij x) = g₁ i x :=
  compatibility' (f · · ·) (g₁ · ·) x

noncomputable def lift
    (Hg : ∀ i j hij x, g₂ j (f i j hij x) = g₂ i x) : P₁ →ₗ[R] P₂ where
  toFun := IsDirectLimit.lift (f · · ·) (g₁ · ·) (g₂ · · )
  map_add' x y := by
    obtain ⟨k, hxk, hyk⟩ :=
      IsDirected.directed (r := (· ≤ ·)) (preimage_index (f · · ·) (g₁ · ·) x)
      (preimage_index (f · · ·) (g₁ · ·) y)
    obtain ⟨k', hxyk', hkk'⟩ := IsDirected.directed (r := (· ≤ ·))
      (preimage_index (f · · ·) (g₁ · ·) (x+y)) k
    have sum_eq : g₁ k' (f (preimage_index (f · · ·) (g₁ · ·) x) k' (le_trans hxk hkk')
      (preimage (f · · ·) (g₁ · ·) x) + f (preimage_index (f · · ·) (g₁ · ·) y) k'
      (le_trans hyk hkk') (preimage (f · · ·) (g₁ · ·) y)) =
        g₁ (preimage_index (f · · ·) (g₁ · ·) (x+y)) (preimage (f · · ·) (g₁ · ·) (x+y)) := by
      simp only [LinearMap.map_add, image_preimage]
      repeat rw [compatibility' (f · · ·) (g₁ · ·),
        image_preimage (f := (f · · ·)) (g := (g₁ · ·))]
    obtain ⟨k'', hk'k'', hxyk'', h'''⟩ := is_injective (f · · ·) (g₁ · ·) sum_eq
    simpa [Hg, IsDirectLimit.lift] using congr_arg (g₂ k'') h'''.symm
  map_smul' r x := by
    have smul_eq : g₁ (preimage_index (f · · ·) (g₁ · ·) (r • x)) (preimage (f · · ·) (g₁ · ·)
      (r • x)) = g₁ (preimage_index (f · · ·) (g₁ · ·) x) (r • preimage (f · · ·) (g₁ · ·) x) := by
      simp only [image_preimage, map_smul]
    obtain ⟨k, hixk, hirxk, h_smul⟩ := is_injective (f · · ·) (g₁ · ·) smul_eq
    simpa [Hg, IsDirectLimit.lift] using congr_arg (g₂ k) h_smul

@[simp]
theorem lift_of (Hg : ∀ i j hij x, g₂ j (f i j hij x) = g₂ i x) {i} (x) :
    (Module.lift f g₁ g₂ Hg) (g₁ i x) = g₂ i x := by
  dsimp [lift, IsDirectLimit.lift_of]
  exact IsDirectLimit.lift_of (f · · ·) (g₁ · ·) (g₂ · ·) Hg x

theorem lift_unique (F : P₁ →ₗ[R] P₂) (x) : F x = (Module.lift f g₁ (fun i ↦ F.comp <| g₁ i)
    (fun i j hij x ↦ by
      simp only [LinearMap.coe_comp, Function.comp_apply]
      exact LinearMap.congr_arg (IsDirectLimit.compatibility (f := (f · · ·))
      (g := (g₁ · ·)) i j hij x))) x := by
  simp only [lift, IsDirectLimit.lift,
    LinearMap.coe_comp, Function.comp_apply, LinearMap.coe_mk, AddHom.coe_mk, image_preimage]

variable [h₂ : IsDirectLimit (f · · ·) (g₂ · ·)]

noncomputable def linearEquiv : P₁ ≃ₗ[R] P₂ :=
  { IsDirectLimit.Equiv (f · · ·) (g₁ · ·) (g₂ · ·) with
    map_add' := (lift f g₁ g₂ h₂.compatibility).map_add'
    map_smul' := (lift f g₁ g₂ h₂.compatibility).map_smul'}

@[simp]
lemma linearEquiv_apply (i : ι) (x : M i) :
  linearEquiv f g₁ g₂ (g₁ i x) = g₂ i x :=
  IsDirectLimit.Equiv_apply (f · · ·) (g₁ · ·) (g₂ · ·) i x

@[simp]
lemma linearEquiv_symm_apply (i : ι) (x : M i) :
  (linearEquiv f g₁ g₂).symm (g₂ i x) = g₁ i x :=
  IsDirectLimit.linearEquiv_symm_apply (f · · ·) (g₁ · ·) (g₂ · ·) i x

end Module

end IsDirectLimit