/-
Copyright (c) 2024 Judith Ludwig, Christian Merten. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Judith Ludwig, Christian Merten
-/
import BruhatTits.Graph.Graph
import BruhatTits.Utils.GraphAction

/-!
# Group actions on the Bruhat-Tits graph

In this file we equip the Bruhat-Tits graph with group actions of `GL₂(K)` and compute stabilizers.

# Main results

- `MulAction.IsPretransitive (GL (Fin 2) K) (Vertices R)` : We show that `GL₂(K)` acts transitively
on the vertices.
- `GraphAction (GL (Fin 2) K) (BTgraph (R:= R))` : The group `GL₂(K)` acts on `BTgraph`.
- `stabilizer_mk_standard_eq_sup` : The stabilizer of the standard vertex of the `GL₂(K)`-action
is `Kˣ * GL₂(R)`.

-/

-- Let R be a discrete valuation ring and K its field of fractions
variable {K : Type*} [Field K]
variable {R : Subring K}

namespace BruhatTits

section Action

lemma isSimilar_smul_of_isSimilar (g : GL (Fin 2) K) (L M : Lattice R) (h : L.IsSimilar R M) :
    (g • L).IsSimilar R (g • M) := by
  obtain ⟨a, rfl⟩ := h
  have : g • a • L = a • g • L := by
    apply Lattice.ext
    simp only [Lattice.smul_M, Lattice.smul_module, Units.smul_def]
    ext x
    simp only [Matrix.GeneralLinearGroup.mem_smul]
    constructor
    · rintro ⟨y, ⟨z, hz, rfl⟩, rfl⟩
      simp only [DistribMulAction.toLinearMap_apply, Matrix.mulVec_smul]
      refine ⟨g.val.mulVec z, ?_, rfl⟩
      simp only [SetLike.mem_coe, Matrix.GeneralLinearGroup.mem_smul]
      use z, hz
    · rintro ⟨y, hy, rfl⟩
      simp only [SetLike.mem_coe, Matrix.GeneralLinearGroup.mem_smul] at hy
      obtain ⟨z, hz, rfl⟩ := hy
      refine ⟨a.val • z, ?_, ?_⟩
      · use z, hz
        rfl
      · simp [Matrix.mulVec_smul]
  rw [this]
  exact Lattice.isSimilar_smul R (g • L) a

def smulGL (g : GL (Fin 2) K) : Vertices R → Vertices R :=
  Quotient.lift (fun L ↦ ⟦g • L⟧) <| by
    intro L M h
    apply Quotient.sound
    exact isSimilar_smul_of_isSimilar g L M h

lemma smulGL_mk (g : GL (Fin 2) K) (L : Lattice R) :
    smulGL g ⟦L⟧ = ⟦g • L⟧ :=
  rfl

instance : SMul (GL (Fin 2) K) (Vertices R) where
  smul := smulGL

lemma Vertices.smul_def (g : GL (Fin 2) K) (x : Vertices R) :
    g • x = smulGL g x :=
  rfl

/-- The action of `GL₂(K)` on vertices. -/
instance : MulAction (GL (Fin 2) K) (Vertices R) where
  one_smul x := by
    refine Quotient.inductionOn x (fun L ↦ ?_)
    rw [Vertices.smul_def, smulGL_mk, one_smul]
  mul_smul g h x := by
    refine Quotient.inductionOn x (fun L ↦ ?_)
    rw [Vertices.smul_def, smulGL_mk]
    rw [mul_smul]
    rw [← smulGL_mk, ← Vertices.smul_def]
    rw [← smulGL_mk, ← Vertices.smul_def]

variable [IsDiscreteValuationRing R] [IsFractionRing R K]

/-- `GL₂(K)` acts transitively on the vertices. -/
instance : MulAction.IsPretransitive (GL (Fin 2) K) (Vertices R) where
  exists_smul_eq L M := by
    classical
    induction' L using Quotient.inductionOn with L
    induction' M using Quotient.inductionOn with M
    obtain ⟨g, rfl⟩ := MulAction.exists_smul_eq (GL (Fin 2) K) L M
    use g
    rfl

@[simp]
lemma inv_smul_smul_eq_inv (g : GL (Fin 2) K) (x y : Vertices R) :
    inv (g • x) (g • y) = inv x y := by
  refine Quotient.inductionOn₂ x y (fun L M ↦ ?_)
  rw [Vertices.smul_def, smulGL_mk]
  rw [Vertices.smul_def, smulGL_mk]
  show dist (g • L) (g • M) = dist L M
  simp

@[simp]
lemma adj_smul_smul_iff_adj (g : GL (Fin 2) K) (x y : Vertices R) :
    BTgraph.Adj (g • x) (g • y) ↔ BTgraph.Adj x y := by
  show inv (g • x) (g • y) = 1 ↔ inv x y = 1
  simp

/-- `GL₂(K)` acts by graph isomorphisms on the Bruhat-Tits tree. -/
def _root_.Matrix.GeneralLinearGroup.toGraphIso (g : GL (Fin 2) K) :
    BTgraph (R := R) ≃g BTgraph (R := R) where
  toEquiv := MulAction.toPerm g
  map_rel_iff' {x y} := by simp

/-- The action of `GL₂(K)` on the Bruhat-Tits graph. -/
instance : GraphAction (GL (Fin 2) K) (BTgraph (R:= R)) where
  smul_adj_smul g x y := (adj_smul_smul_iff_adj g x y).mpr

omit [IsDiscreteValuationRing ↥R] [IsFractionRing (↥R) K] in
lemma cartanDiag_smul_standard
    (ϖ : R) (hϖ : Irreducible ϖ) (f : Fin 2 → ℤ) :
    cartanDiag ϖ hϖ f • Lattice.standard R =
      ((Pi.basisFun K (Fin 2)).twist hϖ f).toLattice := by
  simp only [Lattice.standard]
  have : cartanDiag ϖ hϖ f • ((Pi.basisFun K (Fin 2))) =
      (((Pi.basisFun K (Fin 2)).twist hϖ f)) := by
    ext i j
    simp only [Basis.smulGL_apply, val_cartanDiag, Basis.twist_apply, Pi.smul_apply, smul_eq_mul]
    fin_cases i <;> fin_cases j <;> simp
  rw [← this, Basis.smulGL_toLattice]

open Pointwise in
lemma pow_smul_range_eq_range_iff {ι : Type*} [Finite ι] [Nonempty ι] {ϖ : R}
    (hϖ : Irreducible ϖ) (n : ℤ) :
    ϖ.val ^ n • Set.range ((Algebra.linearMap R K).compLeft ι) =
      Set.range ((Algebra.linearMap R K).compLeft ι) ↔ n = 0 := by
  classical
  refine ⟨fun h ↦ ?_, fun h ↦ by subst h; simp⟩
  let i : ι := ‹Nonempty ι›.some
  by_cases hn : n ≥ 0
  · rw [Set.ext_iff] at h
    have : (Pi.basisFun K ι) i ∈ ϖ.val ^ n • Set.range ⇑((Algebra.linearMap (↥R) K).compLeft ι) := by
      rw [h]
      use Pi.basisFun R ι i
      ext
      simp [Algebra.algebraMap_ofSubring_apply, Pi.single_apply]
    simp at this
    rw [Set.mem_smul_set] at this
    obtain ⟨-, ⟨z, rfl⟩, hy⟩ := this
    rw [funext_iff] at hy
    have := hy i
    simp [Algebra.algebraMap_ofSubring_apply] at this
    have heq : (ValuationRing.valuation R K) (z i).val = 1 := by
      rw [← valuation_eq_one_iff, isUnit_iff_exists_inv']
      use ⟨ϖ.val ^ n, (irreducible_zpow_mem_subring_iff ϖ hϖ n).mpr hn⟩
      ext
      simpa
    apply_fun (ValuationRing.valuation R K ·) at this
    simp [-map_zpow₀] at this
    rw [heq] at this
    simp [-map_zpow₀] at this
    rwa [valuation_irreducible_zpow_eq_one_iff _ hϖ] at this
  · have : Pi.single i (ϖ.val ^ n) ∈ Set.range ((Algebra.linearMap (↥R) K).compLeft ι) := by
      rw [← h]
      rw [Set.mem_smul_set]
      refine ⟨Pi.single i 1, ?_, ?_⟩
      use Pi.single i 1
      ext
      simp [Algebra.algebraMap_ofSubring_apply, Pi.single_apply]
      ext j
      simp [Pi.single_apply]
    simp at this
    obtain ⟨z, hz⟩ := this
    rw [funext_iff] at hz
    have := hz i
    simp [Algebra.algebraMap_ofSubring_apply] at this
    have : ϖ.val ^ n ∈ R := by
      rw [← this]
      exact (z i).property
    rw [irreducible_zpow_mem_subring_iff _ hϖ] at this
    omega

open Pointwise in
/-- The stabilizer of the standard lattice of the `GL₂(K)`-action on lattices is `GL₂(R)`. -/
lemma stabilizer_standard_eq_range_map_subtype :
    MulAction.stabilizer (GL (Fin 2) K) (Lattice.standard R) =
      (Matrix.GeneralLinearGroup.map R.subtype).range := by
  obtain ⟨ϖ, hϖ⟩ := IsDiscreteValuationRing.exists_irreducible R
  ext g
  simp only [MulAction.mem_stabilizer_iff, MonoidHom.mem_range]
  constructor
  · intro hg
    obtain ⟨(k₁ : GL (Fin 2) R), (k₂ : GL (Fin 2) R), (f : Fin 2 → ℤ), hf, h⟩ :=
      cartan_decomposition' (k := 2) ϖ hϖ g
    rw [← h] at hg
    erw [mul_smul, GL_map_eq, GL_map_eq, map_subtype_smul_standard_eq_standard k₂, mul_smul] at hg
    apply_fun (Matrix.GeneralLinearGroup.map R.subtype k₁⁻¹ • ·) at hg
    simp only [map_inv, PNat.val_ofNat, inv_smul_smul] at hg
    rw [← map_inv, map_subtype_smul_standard_eq_standard] at hg
    have hdist : dist (Lattice.standard R) (cartanDiag ϖ hϖ f • Lattice.standard R) = 0 := by
      simp [hg]
    simp only [cartanDiag_smul_standard] at hdist
    simp only [Lattice.standard] at hdist
    apply_fun Int.ofNat at hdist
    simp only [Int.ofNat_eq_coe, dist_twist _ _ hf, Fin.isValue, CharP.cast_eq_zero, sub_eq_zero] at hdist
    rw [Lattice.ext_iff, Lattice.smul_M] at hg
    simp only [cartanDiag] at hg
    rw [Matrix.GeneralLinearGroup.diagonal_smul _ _ _ 0] at hg
    · simp only [Fin.isValue, zpow_neg, Units.smul_mk_apply] at hg
      have : ϖ.val ^ f 0 • (Lattice.standard R).M = (Lattice.standard R).M := hg
      rw [Lattice.standard_M, SetLike.ext'_iff] at this
      simp only [Fin.isValue, Submodule.coe_pointwise_smul, LinearMap.range_coe,
        pow_smul_range_eq_range_iff hϖ] at this
      have : f = 0 := by
        ext i
        fin_cases i
        · simpa
        · simpa [← hdist]
      subst this
      use k₁ * k₂
      rw [cartanDiag_zero, GL_map_eq, GL_map_eq] at h
      simpa using h
    · intro i j
      fin_cases i <;> fin_cases j <;> simp [hdist]
  · rintro ⟨k, rfl⟩
    rw [map_subtype_smul_standard_eq_standard]

open Pointwise in
/-- The stabilizer of the `GL₂(K)`-action on the standard vertex is `Kˣ * GL₂(R)`. -/
lemma stabilizer_mk_standard_eq_sup :
    MulAction.stabilizer (GL (Fin 2) K) (Vertices.mk <| Lattice.standard R) =
      (Matrix.GeneralLinearGroup.embDiagonal K (Fin 2)).range ⊔
        (Matrix.GeneralLinearGroup.map (n := Fin 2) R.subtype).range := by
  apply le_antisymm
  · intro g (hg : ⟦g • Lattice.standard R⟧ = ⟦Lattice.standard R⟧)
    simp only [Quotient.eq] at hg
    obtain ⟨u, hu⟩ := hg
    have (L : Lattice R) : u • L = Matrix.GeneralLinearGroup.embDiagonal _ (Fin 2) u • L := by
      simp only [Matrix.GeneralLinearGroup.embDiagonal_apply]
      ext : 1
      rw [Lattice.smul_M, Matrix.GeneralLinearGroup.diagonal_smul _ _ _ 1]
      rfl
      tauto
    rw [this, ← mul_smul] at hu
    have : ((Matrix.GeneralLinearGroup.embDiagonal K (Fin 2)) u * g) ∈ MulAction.stabilizer
        (GL (Fin 2) K) (Lattice.standard R) := hu
    rw [stabilizer_standard_eq_range_map_subtype] at this
    obtain ⟨a, ha⟩ := this
    have : g = ((Matrix.GeneralLinearGroup.embDiagonal K (Fin 2)) u)⁻¹ *
        (Matrix.GeneralLinearGroup.map R.subtype) a := by simp [ha]
    rw [this]
    apply Subgroup.mul_mem_sup
    apply Subgroup.inv_mem
    use u
    use a
  · rw [sup_le_iff]
    constructor
    · rintro - ⟨u, rfl⟩
      simp only [Matrix.GeneralLinearGroup.embDiagonal_apply, MulAction.mem_stabilizer_iff]
      show ⟦_⟧ = _
      have : (Matrix.GL.diagonal fun x : Fin 2 ↦ u) • Lattice.standard R = u • Lattice.standard R := by
        ext : 1
        rw [Lattice.smul_M, Matrix.GeneralLinearGroup.diagonal_smul _ _ _ 1]
        rfl
        tauto
      simp only [this, Quotient.eq]
      apply Lattice.smul_isSimilar
    · rintro - ⟨a, rfl⟩
      simp only [MulAction.mem_stabilizer_iff]
      show ⟦_⟧ = _
      rw [map_subtype_smul_standard_eq_standard]