p2m_Flattening_p3a

import Mathlib.LinearAlgebra.PiTensorProduct
import Mathlib.LinearAlgebra.TensorProduct.Basic
import Mathlib.RingTheory.Flat.Basic
import Mathlib.LinearAlgebra.PiTensorProduct.Basis
import Mathlib.LinearAlgebra.Basis.VectorSpace
import Mathlib.Data.Fintype.Basic
import Mathlib.Data.Finset.Basic
import Mathlib.Data.Real.Basic
import Mathlib.Logic.Function.Basic
import Mathlib.LinearAlgebra.Dual.Lemmas
import Mathlib.LinearAlgebra.Dimension.LinearMap
import Mathlib.LinearAlgebra.Contraction
import Mathlib.LinearAlgebra.FreeModule.Finite.Basic
import Mathlib.LinearAlgebra.Dimension.Finrank
import Definitions.Def_p2m_TensorObj_p1
import Definitions.Def_p2m_TensorObj_p2
import Definitions.Def_p2m_TensorObj_p3
import Definitions.Def_p2m_TensorObj_p4
import Definitions.Def_p2m_Tensor_quot
import Definitions.Def_p2m_Spectrum
import Definitions.Def_p2m_StrassenPreorder
import Definitions.Def_p2m_Flattening_p1
import Definitions.Def_p2m_Flattening_p2

open TensorProduct PiTensorProduct BigOperators Module

universe u v w

set_option maxHeartbeats 1000000

namespace AsymptoticSpectra.Tensor

variable {K : Type u} [Field K] {d : ℕ} [Fact (1 < d)]
variable (σ : Split (Fin d))



theorem flatteningRank_add
    (x y : Tensor K d) :
    flatteningRank σ (x + y)
      = flatteningRank σ x
      + flatteningRank σ y := by
  -- Lift to TensorObj representatives
  induction x using Quotient.inductionOn with | _ X =>
  induction y using Quotient.inductionOn with | _ Y =>
  -- Tensor.add is defined via Quotient.liftOn₂
  change flatteningRank σ (Tensor.add (Tensor.toTensor X) (Tensor.toTensor Y)) =
         flatteningRank σ (Tensor.toTensor X) + flatteningRank σ (Tensor.toTensor Y)
  simp only [Tensor.add, flatteningRank_mk]

  -- Setup FiniteDimensional instances
  haveI := X.finiteDimensional
  haveI := Y.finiteDimensional

  let BS := PiTensorProduct K (fun i : Sc σ => X.V i)
  let BT := PiTensorProduct K (fun i : Sc σ => Y.V i)

  haveI : FiniteDimensional K BS :=
    Module.Finite.of_basis (Basis.piTensorProduct (fun i => Module.Free.chooseBasis K (X.V i)))
  haveI : FiniteDimensional K BT :=
    Module.Finite.of_basis (Basis.piTensorProduct (fun i => Module.Free.chooseBasis K (Y.V i)))

  -- Rewrite flatteningRank as finrank of range
  -- First, simplify the LHS using the quotient structure
  show AsymptoticSpectra.flatteningRank σ (X + Y) =
       AsymptoticSpectra.flatteningRank σ X + AsymptoticSpectra.flatteningRank σ Y
  unfold AsymptoticSpectra.flatteningRank

  -- Apply flatteningMap_add_range to decompose the range
  rw [flatteningMap_add_range X Y]

  -- Define abbreviations for the mapped submodules
  let sX := Submodule.map (TensorObj.liftMap (fun i : Sc σ => LinearMap.inl K (X.V i) (Y.V i)))
                          (LinearMap.range (flatteningMap σ X))
  let sY := Submodule.map (TensorObj.liftMap (fun i : Sc σ => LinearMap.inr K (X.V i) (Y.V i)))
                          (LinearMap.range (flatteningMap σ Y))

  -- finrank of each mapped submodule equals the original finrank (injectivity of inl/inr)
  have hX : finrank K ↥sX = finrank K ↥(LinearMap.range (flatteningMap σ X)) :=
    finrank_map_of_injective _ (TensorObj.liftMap_injective_of_injective _ (fun i => LinearMap.inl_injective)) _

  have hY : finrank K ↥sY = finrank K ↥(LinearMap.range (flatteningMap σ Y)) :=
    finrank_map_of_injective _ (TensorObj.liftMap_injective_of_injective _ (fun i => LinearMap.inr_injective)) _

  -- The two ranges are disjoint, so finrank of sup equals sum of finranks
  have h_disjoint : Disjoint sX sY := flatteningMap_add_range_disjoint X Y
  have h_sup : finrank K ↥(sX ⊔ sY) = finrank K ↥sX + finrank K ↥sY :=
    finrank_sup_of_disjoint sX sY inferInstance inferInstance h_disjoint

  calc finrank K ↥(sX ⊔ sY)
      = finrank K ↥sX + finrank K ↥sY := h_sup
    _ = finrank K ↥(flatteningMap σ X).range + finrank K ↥sY := by rw [hX]
    _ = finrank K ↥(flatteningMap σ X).range + finrank K ↥(flatteningMap σ Y).range := by rw [hY]

/-- Auxiliary bilinear map for the distributivity construction.
    For fixed v' and w', the map (v, w) ↦ (tprod (update v' i v)) ⊗ (tprod (update w' i w))
    is bilinear. -/
private noncomputable def piTensorDistribAuxBilin {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)]
    (i : ι) (v' : ∀ j, V j) (w' : ∀ j, W j) :
    V i →ₗ[K] W i →ₗ[K] (PiTensorProduct K V) ⊗[K] (PiTensorProduct K W) where
  toFun v := {
    toFun := fun w => tprod K (Function.update v' i v) ⊗ₜ[K] tprod K (Function.update w' i w)
    map_add' := fun w₁ w₂ => by simp only [MultilinearMap.map_update_add, tmul_add]
    map_smul' := fun c w => by simp only [MultilinearMap.map_update_smul, tmul_smul, RingHom.id_apply]
  }
  map_add' v₁ v₂ := by
    ext w; simp only [LinearMap.coe_mk, AddHom.coe_mk, LinearMap.add_apply,
      MultilinearMap.map_update_add, add_tmul]
  map_smul' c v := by
    ext w; simp only [LinearMap.coe_mk, AddHom.coe_mk, RingHom.id_apply, LinearMap.smul_apply,
      MultilinearMap.map_update_smul, smul_tmul']

/-- For fixed v' and w', the linear map on V i ⊗ W i induced by the bilinear map. -/
private noncomputable def piTensorDistribAuxLin {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)]
    (i : ι) (v' : ∀ j, V j) (w' : ∀ j, W j) :
    V i ⊗[K] W i →ₗ[K] (PiTensorProduct K V) ⊗[K] (PiTensorProduct K W) :=
  TensorProduct.lift (piTensorDistribAuxBilin i v' w')

@[simp]
private theorem piTensorDistribAuxLin_tmul {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)]
    (i : ι) (v' : ∀ j, V j) (w' : ∀ j, W j) (v : V i) (w : W i) :
    piTensorDistribAuxLin i v' w' (v ⊗ₜ[K] w) =
      tprod K (Function.update v' i v) ⊗ₜ[K] tprod K (Function.update w' i w) := by
  simp only [piTensorDistribAuxLin, TensorProduct.lift.tmul, piTensorDistribAuxBilin,
    LinearMap.coe_mk, AddHom.coe_mk]

/-- The bilinear map (v, w) ↦ (tprod v) ⊗ (tprod w) is multilinear in v and w.
    This is a 2n-linear map when ι has n elements. -/
private noncomputable def tprodTensorBilin {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)] :
    MultilinearMap K V (MultilinearMap K W ((PiTensorProduct K V) ⊗[K] (PiTensorProduct K W))) where
  toFun v := {
    toFun := fun w => tprod K v ⊗ₜ[K] tprod K w
    map_update_add' := fun w i x y => by
      simp only [MultilinearMap.map_update_add, tmul_add]
    map_update_smul' := fun w i c x => by
      simp only [MultilinearMap.map_update_smul, tmul_smul]
  }
  map_update_add' v i x y := by
    ext w
    simp only [MultilinearMap.coe_mk, MultilinearMap.add_apply,
      MultilinearMap.map_update_add, add_tmul]
  map_update_smul' v i c x := by
    ext w
    simp only [MultilinearMap.coe_mk, MultilinearMap.smul_apply,
      MultilinearMap.map_update_smul, smul_tmul']

/-- Inner bilinear map used in piTensorDistribMultilinear_exists. -/
noncomputable def mkInnerBilin {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)]
    {α : Type*} [Fintype α] [DecidableEq α]
    (e : ι ≃ Option α)
    (V' : α → Type _) (W' : α → Type _)
    [∀ a, AddCommGroup (V' a)] [∀ a, Module K (V' a)]
    [∀ a, AddCommGroup (W' a)] [∀ a, Module K (W' a)]
    (combineV_fun : V (e.symm none) → PiTensorProduct K V' → PiTensorProduct K V)
    (combineW_fun : W (e.symm none) → PiTensorProduct K W' → PiTensorProduct K W)
    (combineW_add  : ∀ tW' x y, combineW_fun (x + y) tW' = combineW_fun x tW' + combineW_fun y tW')
    (combineW_smul : ∀ tW' (c : K) x, combineW_fun (c • x) tW' = c • combineW_fun x tW')
    (combineV_add  : ∀ tV' x y, combineV_fun (x + y) tV' = combineV_fun x tV' + combineV_fun y tV')
    (combineV_smul : ∀ tV' (c : K) x, combineV_fun (c • x) tV' = c • combineV_fun x tV')
    (combineW_add' : ∀ w₀ x y, combineW_fun w₀ (x + y) = combineW_fun w₀ x + combineW_fun w₀ y)
    (combineW_smul' : ∀ w₀ (c : K) x, combineW_fun w₀ (c • x) = c • combineW_fun w₀ x)
    (combineV_add' : ∀ v₀ x y, combineV_fun v₀ (x + y) = combineV_fun v₀ x + combineV_fun v₀ y)
    (combineV_smul' : ∀ v₀ (c : K) x, combineV_fun v₀ (c • x) = c • combineV_fun v₀ x) :
    (PiTensorProduct K V') →ₗ[K] (PiTensorProduct K W') →ₗ[K]
        (V (e.symm none) →ₗ[K] W (e.symm none) →ₗ[K]
            (PiTensorProduct K V) ⊗[K] (PiTensorProduct K W)) :=
  { toFun := fun tV' =>
      { toFun := fun tW' =>
          { toFun := fun v₀ =>
              { toFun := fun w₀ => combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun w₀ tW'
                map_add' := fun x y => by
                  rw [combineW_add]
                  erw [show combineV_fun v₀ tV' ⊗ₜ[K] (combineW_fun x tW' + combineW_fun y tW') =
                      combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun x tW' +
                      combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun y tW' from
                    TensorProduct.tmul_add _ _ _]
                map_smul' := fun c x => by
                  rw [combineW_smul]
                  simp only [RingHom.id_apply]
                  erw [show combineV_fun v₀ tV' ⊗ₜ[K] (c • combineW_fun x tW') =
                      c • (combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun x tW') from
                    TensorProduct.tmul_smul c _ _] }
            map_add' := fun x y => by
              ext w₀
              simp only [LinearMap.coe_mk, AddHom.coe_mk, LinearMap.add_apply]
              rw [combineV_add]
              erw [show (combineV_fun x tV' + combineV_fun y tV') ⊗ₜ[K] combineW_fun w₀ tW' =
                  combineV_fun x tV' ⊗ₜ[K] combineW_fun w₀ tW' +
                  combineV_fun y tV' ⊗ₜ[K] combineW_fun w₀ tW' from
                TensorProduct.add_tmul _ _ _]
            map_smul' := fun c x => by
              ext w₀
              simp only [LinearMap.coe_mk, AddHom.coe_mk, RingHom.id_apply, LinearMap.smul_apply]
              rw [combineV_smul]
              erw [show (c • combineV_fun x tV') ⊗ₜ[K] combineW_fun w₀ tW' =
                  c • (combineV_fun x tV' ⊗ₜ[K] combineW_fun w₀ tW') from
                TensorProduct.smul_tmul' c _ _] }
        map_add' := fun x y => by
          ext v₀ w₀
          simp only [LinearMap.coe_mk, AddHom.coe_mk, LinearMap.add_apply]
          rw [combineW_add']
          erw [show combineV_fun v₀ tV' ⊗ₜ[K] (combineW_fun w₀ x + combineW_fun w₀ y) =
              combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun w₀ x +
              combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun w₀ y from
            TensorProduct.tmul_add _ _ _]
        map_smul' := fun c x => by
          ext v₀ w₀
          simp only [LinearMap.coe_mk, AddHom.coe_mk, LinearMap.smul_apply, RingHom.id_apply]
          rw [combineW_smul']
          erw [show combineV_fun v₀ tV' ⊗ₜ[K] (c • combineW_fun w₀ x) =
              c • (combineV_fun v₀ tV' ⊗ₜ[K] combineW_fun w₀ x) from
            TensorProduct.tmul_smul c _ _] }
    map_add' := fun x y => by
      apply PiTensorProduct.ext
      ext tW' v₀ w₀
      simp only [LinearMap.compMultilinearMap_apply]
      show combineV_fun v₀ (x + y) ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW') =
        combineV_fun v₀ x ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW') +
        combineV_fun v₀ y ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW')
      rw [combineV_add']
      erw [TensorProduct.add_tmul]
    map_smul' := fun c x => by
      apply PiTensorProduct.ext
      ext tW' v₀ w₀
      simp only [LinearMap.compMultilinearMap_apply, RingHom.id_apply]
      show combineV_fun v₀ (c • x) ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW') =
        c • (combineV_fun v₀ x ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW'))
      rw [combineV_smul']
      erw [show (c • combineV_fun v₀ x) ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW') =
          c • (combineV_fun v₀ x ⊗ₜ[K] combineW_fun w₀ (PiTensorProduct.tprod K tW')) from
        TensorProduct.smul_tmul' c _ _] }


end AsymptoticSpectra.Tensor