import PaperXII.AuxiliaryGroupAlgebra
import Mathlib.LinearAlgebra.Basis.Basic
import Mathlib.LinearAlgebra.Quotient.Basic
import Mathlib.LinearAlgebra.Isomorphisms
import Mathlib.LinearAlgebra.StdBasis
import Mathlib.LinearAlgebra.TensorProduct.Basic

set_option autoImplicit false

/-! Actual augmentation relations and free coinvariants for an augmented algebra.
The algebra may be noncommutative. All quotients are coefficient modules. -/
namespace PaperXII.AugmentationCoinvariants

open scoped TensorProduct
open Module

variable {R A M : Type*} [CommRing R] [Ring A] [Algebra R A]
  [AddCommGroup M] [Module R M] [Module A M] [IsScalarTower R A M]

/-- The module generated by the actual balancing relations. -/
def relations (ε : A →ₐ[R] R) : Submodule R M :=
  Submodule.span R (Set.range fun am : A × M => am.1 • am.2 - ε am.1 • am.2)

/-- The actual augmentation ideal acting on the module, with finite linear sums. -/
def augmentationSMul (ε : A →ₐ[R] R) : Submodule R M :=
  Submodule.span R {v | ∃ a : A, ε a = 0 ∧ ∃ m : M, a • m = v}

omit [IsScalarTower R A M] in
/-- Each balancing difference belongs to the relation submodule. -/
theorem relation_mem (ε : A →ₐ[R] R) (a : A) (m : M) :
    a • m - ε a • m ∈ relations ε :=
  Submodule.subset_span (Set.mem_range_self (a, m))

/-- Balancing relations are exactly the augmentation ideal times the module. -/
theorem relations_eq_augmentationSMul (ε : A →ₐ[R] R) :
    relations (M := M) ε = augmentationSMul ε := by
  apply le_antisymm
  · apply Submodule.span_le.mpr
    rintro v ⟨⟨a,m⟩,rfl⟩
    apply Submodule.subset_span
    refine ⟨a - algebraMap R A (ε a), by simp, m, ?_⟩
    simp [sub_smul]
  · apply Submodule.span_le.mpr
    rintro v ⟨a,ha,m,rfl⟩
    simpa [ha] using relation_mem ε a m

/-- Coinvariants as the actual quotient by balancing relations. -/
abbrev Coinvariants (ε : A →ₐ[R] R) := M ⧸ relations (M := M) ε

/-- The quotient map. -/
def mk (ε : A →ₐ[R] R) : M →ₗ[R] Coinvariants (M := M) ε :=
  (relations ε).mkQ

omit [IsScalarTower R A M] in
/-- Augmented scalar multiplication becomes coefficient scalar multiplication. -/
theorem mk_smul (ε : A →ₐ[R] R) (a : A) (m : M) :
    mk ε (a • m) = ε a • mk ε m := by
  rw [← map_smul]
  exact (Submodule.Quotient.eq _).mpr (relation_mem ε a m)

variable {ι : Type*} [Fintype ι]

/-- Apply augmentation to every coordinate of a supplied free basis. -/
noncomputable def coordinateMap (ε : A →ₐ[R] R) (b : Basis ι A M) :
    M →ₗ[R] (ι → R) :=
  (LinearMap.pi fun i => ε.toLinearMap.comp
    ((LinearMap.proj i).comp (b.equivFun.restrictScalars R).toLinearMap))

/-- The coordinate map on an element. -/
theorem coordinateMap_apply (ε : A →ₐ[R] R) (b : Basis ι A M) (m : M) (i : ι) :
    coordinateMap ε b m i = ε (b.equivFun m i) := rfl

/-- The coordinate map has the augmented scalar action. -/
theorem coordinateMap_smul (ε : A →ₐ[R] R) (b : Basis ι A M) (a : A) (m : M) :
    coordinateMap ε b (a • m) = ε a • coordinateMap ε b m := by
  ext i
  simp [coordinateMap_apply, map_smul, smul_eq_mul]

/-- The relation module is exactly the kernel of augmentation on free coordinates. -/
theorem ker_coordinateMap (ε : A →ₐ[R] R) (b : Basis ι A M) :
    LinearMap.ker (coordinateMap ε b) = relations ε := by
  apply le_antisymm
  · intro m hm
    have hz : ∀ i, ε (b.equivFun m i) = 0 := fun i => congrFun hm i
    rw [← b.sum_equivFun m]
    apply Submodule.sum_mem
    intro i _
    simpa only [hz i, zero_smul, sub_zero] using relation_mem ε (b.equivFun m i) (b i)
  · apply Submodule.span_le.mpr
    rintro v ⟨⟨a,m⟩,rfl⟩
    change coordinateMap ε b (a • m - ε a • m) = 0
    simp only [map_sub, coordinateMap_smul, map_smul, sub_self]

/-- Free coordinate augmentation is onto, including for the empty basis. -/
theorem coordinateMap_surjective (ε : A →ₐ[R] R) (b : Basis ι A M) :
    Function.Surjective (coordinateMap ε b) := by
  intro y
  refine ⟨b.equivFun.symm (fun i => algebraMap R A (y i)), ?_⟩
  ext i
  change ε (b.equivFun (b.equivFun.symm (fun i => algebraMap R A (y i))) i) = y i
  rw [b.equivFun.apply_symm_apply]
  exact ε.commutes (y i)

/-- Free coinvariants recover exactly one coefficient coordinate per basis vector. -/
noncomputable def freeEquiv (ε : A →ₐ[R] R) (b : Basis ι A M) :
    Coinvariants (M := M) ε ≃ₗ[R] (ι → R) :=
  (Submodule.quotEquivOfEq _ _ (ker_coordinateMap ε b).symm).trans
    ((coordinateMap ε b).quotKerEquivOfSurjective (coordinateMap_surjective ε b))

/-- Recovery agrees with augmentation of the original coordinates. -/
theorem freeEquiv_mk (ε : A →ₐ[R] R) (b : Basis ι A M) (m : M) :
    freeEquiv ε b (mk ε m) = coordinateMap ε b m := by
  simp [freeEquiv, mk]

/-- The recovered coefficient module has a basis indexed by the original basis. -/
noncomputable def recoveredBasis (ε : A →ₐ[R] R) (b : Basis ι A M) :
    Basis ι R (Coinvariants (M := M) ε) :=
  (Pi.basisFun R ι).map (freeEquiv ε b).symm

end PaperXII.AugmentationCoinvariants