p2m_Flattening_p3b

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
import Definitions.Def_p2m_Flattening_p3a

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 piTensorDistribMultilinear_exists {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)] :
    ∃ M : MultilinearMap K (fun i => V i ⊗[K] W i)
            ((PiTensorProduct K V) ⊗[K] (PiTensorProduct K W)),
      ∀ v w, M (fun i => v i ⊗ₜ[K] w i) = tprod K v ⊗ₜ[K] tprod K w := by
  induction h : Fintype.card ι using Nat.strong_induction_on generalizing ι V W with
  | _ n ih =>
    cases isEmpty_or_nonempty ι with
    | inl hι =>
      haveI : IsEmpty ι := hι
      refine ⟨⟨fun _ => (PiTensorProduct.isEmptyEquiv ι).symm 1 ⊗ₜ[K]
                        (PiTensorProduct.isEmptyEquiv ι).symm 1,
               fun _ i => isEmptyElim i, fun _ i => isEmptyElim i⟩, ?_⟩
      intro v w
      have hv : tprod K v = (PiTensorProduct.isEmptyEquiv ι).symm 1 := by
        have : v = isEmptyElim := funext (fun i => isEmptyElim i)
        rw [this, PiTensorProduct.isEmptyEquiv_symm_apply, one_smul]
      have hw : tprod K w = (PiTensorProduct.isEmptyEquiv ι).symm 1 := by
        have : w = isEmptyElim := funext (fun i => isEmptyElim i)
        rw [this, PiTensorProduct.isEmptyEquiv_symm_apply, one_smul]
      simp only [MultilinearMap.coe_mk, hv, hw]
    | inr hι =>
      haveI : Nonempty ι := hι
      haveI : Inhabited ι := Classical.inhabited_of_nonempty hι

      obtain ⟨α, e⟩ := Equiv.sigmaEquivOptionOfInhabited ι
      haveI : Fintype α := fintypeOfOptionEquiv e
      haveI : DecidableEq α := Classical.decEq α

      have hα_lt : Fintype.card α < n := by
        have h1 : Fintype.card ι = Fintype.card (Option α) := Fintype.card_congr e
        rw [Fintype.card_option] at h1
        omega

      let V' : α → Type _ := fun a => V (e.symm (some a))
      let W' : α → Type _ := fun a => W (e.symm (some a))
      obtain ⟨M', hM'⟩ := ih (Fintype.card α) hα_lt (V := V') (W := W') rfl

      let i₀ : ι := e.symm none

      -- Key equivalence: Option α ≃ α ⊕ PUnit (some ↔ inl, none ↔ inr)
      let optSum : Option α ≃ α ⊕ PUnit.{1} := Equiv.optionEquivSumPUnit α

      -- Define the type families over the sum type
      let Vsum : α ⊕ PUnit.{1} → Type _ := fun s => V (e.symm (optSum.symm s))
      let Wsum : α ⊕ PUnit.{1} → Type _ := fun s => W (e.symm (optSum.symm s))

      let eSum : ι ≃ α ⊕ PUnit.{1} := e.trans optSum

      let subsingleV := (PiTensorProduct.subsingletonEquiv (R := K) (s := fun _ : PUnit.{1} => V i₀) PUnit.unit).symm
      let subsingleW := (PiTensorProduct.subsingletonEquiv (R := K) (s := fun _ : PUnit.{1} => W i₀) PUnit.unit).symm

      let combineV_fun : V i₀ → PiTensorProduct K V' → PiTensorProduct K V := fun v₀ tV' =>
        (PiTensorProduct.reindex K V eSum).symm
          ((PiTensorProduct.tmulEquivDep K Vsum) (tV' ⊗ₜ[K] subsingleV v₀))

      let combineW_fun : W i₀ → PiTensorProduct K W' → PiTensorProduct K W := fun w₀ tW' =>
        (PiTensorProduct.reindex K W eSum).symm
          ((PiTensorProduct.tmulEquivDep K Wsum) (tW' ⊗ₜ[K] subsingleW w₀))

      have combineW_add : ∀ tW' x y, combineW_fun (x + y) tW' = combineW_fun x tW' + combineW_fun y tW' := by
        intro tW' x y
        show (PiTensorProduct.reindex K W eSum).symm
            ((PiTensorProduct.tmulEquivDep K Wsum) (tW' ⊗ₜ[K] subsingleW (x + y))) = _
        rw [subsingleW.map_add]
        erw [TensorProduct.tmul_add,
            (PiTensorProduct.tmulEquivDep K Wsum).map_add,
            (PiTensorProduct.reindex K W eSum).symm.map_add]
      have combineW_smul : ∀ tW' (c : K) x, combineW_fun (c • x) tW' = c • combineW_fun x tW' := by
        intro tW' c x
        show (PiTensorProduct.reindex K W eSum).symm
            ((PiTensorProduct.tmulEquivDep K Wsum) (tW' ⊗ₜ[K] subsingleW (c • x))) = _
        rw [subsingleW.map_smul]
        erw [show tW' ⊗ₜ[K] (c • subsingleW x) = c • (tW' ⊗ₜ[K] subsingleW x) from
          TensorProduct.tmul_smul c tW' (subsingleW x),
            (PiTensorProduct.tmulEquivDep K Wsum).map_smul,
            (PiTensorProduct.reindex K W eSum).symm.map_smul]
        rfl
      have combineV_add : ∀ tV' x y, combineV_fun (x + y) tV' = combineV_fun x tV' + combineV_fun y tV' := by
        intro tV' x y
        show (PiTensorProduct.reindex K V eSum).symm
            ((PiTensorProduct.tmulEquivDep K Vsum) (tV' ⊗ₜ[K] subsingleV (x + y))) = _
        rw [subsingleV.map_add]
        erw [TensorProduct.tmul_add,
            (PiTensorProduct.tmulEquivDep K Vsum).map_add,
            (PiTensorProduct.reindex K V eSum).symm.map_add]
      have combineV_smul : ∀ tV' (c : K) x, combineV_fun (c • x) tV' = c • combineV_fun x tV' := by
        intro tV' c x
        show (PiTensorProduct.reindex K V eSum).symm
            ((PiTensorProduct.tmulEquivDep K Vsum) (tV' ⊗ₜ[K] subsingleV (c • x))) = _
        rw [subsingleV.map_smul]
        erw [show tV' ⊗ₜ[K] (c • subsingleV x) = c • (tV' ⊗ₜ[K] subsingleV x) from
          TensorProduct.tmul_smul c tV' (subsingleV x),
            (PiTensorProduct.tmulEquivDep K Vsum).map_smul,
            (PiTensorProduct.reindex K V eSum).symm.map_smul]
        rfl

      -- Linearity in the second argument (tW' and tV')
      have combineW_add' : ∀ w₀ x y, combineW_fun w₀ (x + y) = combineW_fun w₀ x + combineW_fun w₀ y := by
        intro w₀ x y
        show (PiTensorProduct.reindex K W eSum).symm
            ((PiTensorProduct.tmulEquivDep K Wsum) ((x + y) ⊗ₜ[K] subsingleW w₀)) = _
        erw [TensorProduct.add_tmul,
            (PiTensorProduct.tmulEquivDep K Wsum).map_add,
            (PiTensorProduct.reindex K W eSum).symm.map_add]
      have combineW_smul' : ∀ w₀ (c : K) x, combineW_fun w₀ (c • x) = c • combineW_fun w₀ x := by
        intro w₀ c x
        show (PiTensorProduct.reindex K W eSum).symm
            ((PiTensorProduct.tmulEquivDep K Wsum) ((c • x) ⊗ₜ[K] subsingleW w₀)) = _
        erw [show (c • x) ⊗ₜ[K] subsingleW w₀ = c • (x ⊗ₜ[K] subsingleW w₀) from
          TensorProduct.smul_tmul' c x (subsingleW w₀),
            (PiTensorProduct.tmulEquivDep K Wsum).map_smul,
            (PiTensorProduct.reindex K W eSum).symm.map_smul]
        rfl
      have combineV_add' : ∀ v₀ x y, combineV_fun v₀ (x + y) = combineV_fun v₀ x + combineV_fun v₀ y := by
        intro v₀ x y
        show (PiTensorProduct.reindex K V eSum).symm
            ((PiTensorProduct.tmulEquivDep K Vsum) ((x + y) ⊗ₜ[K] subsingleV v₀)) = _
        erw [TensorProduct.add_tmul,
            (PiTensorProduct.tmulEquivDep K Vsum).map_add,
            (PiTensorProduct.reindex K V eSum).symm.map_add]
      have combineV_smul' : ∀ v₀ (c : K) x, combineV_fun v₀ (c • x) = c • combineV_fun v₀ x := by
        intro v₀ c x
        show (PiTensorProduct.reindex K V eSum).symm
            ((PiTensorProduct.tmulEquivDep K Vsum) ((c • x) ⊗ₜ[K] subsingleV v₀)) = _
        erw [show (c • x) ⊗ₜ[K] subsingleV v₀ = c • (x ⊗ₜ[K] subsingleV v₀) from
          TensorProduct.smul_tmul' c x (subsingleV v₀),
            (PiTensorProduct.tmulEquivDep K Vsum).map_smul,
            (PiTensorProduct.reindex K V eSum).symm.map_smul]
        rfl

      let innerBilin := mkInnerBilin e V' W' combineV_fun combineW_fun
          combineW_add combineW_smul combineV_add combineV_smul
          combineW_add' combineW_smul' combineV_add' combineV_smul'

      have hM : ∃ M : MultilinearMap K (fun i => V i ⊗[K] W i)
              ((PiTensorProduct K V) ⊗[K] (PiTensorProduct K W)),
            ∀ v w, M (fun i => v i ⊗ₜ[K] w i) = tprod K v ⊗ₜ[K] tprod K w := by
        refine ⟨⟨fun f => TensorProduct.lift
              (TensorProduct.lift innerBilin (M' (fun a => f (e.symm (some a))))) (f i₀),
            ?map_update_add', ?map_update_smul'⟩, ?spec⟩
        case map_update_add' =>
          intro _ f i x y
          simp only
          by_cases hi : i = i₀
          · subst hi
            simp only [Function.update_self]
            have hne : ∀ a, e.symm (some a) ≠ i₀ := by
              intro a h
              have : e (e.symm (some a)) = e i₀ := by rw [h]
              simp only [Equiv.apply_symm_apply, i₀] at this
              cases this
            have h_some : ∀ a, Function.update f i₀ (x + y) (e.symm (some a)) = f (e.symm (some a)) :=
              fun a => Function.update_of_ne (hne a) (x + y) f
            have h_some_x : ∀ a, Function.update f i₀ x (e.symm (some a)) = f (e.symm (some a)) :=
              fun a => Function.update_of_ne (hne a) x f
            have h_some_y : ∀ a, Function.update f i₀ y (e.symm (some a)) = f (e.symm (some a)) :=
              fun a => Function.update_of_ne (hne a) y f
            simp only [h_some, h_some_x, h_some_y]
            rw [map_add]
          · have h_none : Function.update f i (x + y) i₀ = f i₀ := Function.update_of_ne (Ne.symm hi) (x + y) f
            have h_none_x : Function.update f i x i₀ = f i₀ := Function.update_of_ne (Ne.symm hi) x f
            have h_none_y : Function.update f i y i₀ = f i₀ := Function.update_of_ne (Ne.symm hi) y f
            simp only [h_none, h_none_x, h_none_y]
            have hi' : ∃ a, i = e.symm (some a) := by
              cases he : e i with
              | none => exfalso; apply hi; simp only [i₀]; rw [← he]; simp
              | some a => exact ⟨a, by simp [← he]⟩
            obtain ⟨a₀, ha₀⟩ := hi'
            subst ha₀
            have hM'_add := M'.map_update_add (fun a => f (e.symm (some a))) a₀ x y
            simp only at hM'_add ⊢
            have h_update : ∀ z : V' a₀ ⊗[K] W' a₀, (fun a => Function.update f (e.symm (some a₀)) z (e.symm (some a))) =
                Function.update (fun a => f (e.symm (some a))) a₀ z := by
              intro z; ext a
              by_cases ha : a = a₀
              · subst ha; simp [Function.update_self]
              · rw [Function.update_of_ne ha, Function.update_of_ne]
                intro h; apply ha; exact Option.some_injective α (e.symm.injective h)
            simp only [h_update]
            rw [hM'_add, map_add]
            have lift_add : ∀ (g₁ g₂ : V i₀ →ₗ[K] W i₀ →ₗ[K] PiTensorProduct K V ⊗[K] PiTensorProduct K W),
                TensorProduct.lift (g₁ + g₂) = TensorProduct.lift g₁ + TensorProduct.lift g₂ := by
              intro g₁ g₂
              change TensorProduct.uncurry _ _ _ _ (g₁ + g₂) = TensorProduct.uncurry _ _ _ _ g₁ + TensorProduct.uncurry _ _ _ _ g₂
              rw [map_add]
            rw [lift_add, LinearMap.add_apply]
        case map_update_smul' =>
          intro _ f i c x
          simp only
          by_cases hi : i = i₀
          · subst hi
            simp only [Function.update_self]
            have hne : ∀ a, e.symm (some a) ≠ i₀ := by
              intro a h
              have : e (e.symm (some a)) = e i₀ := by rw [h]
              simp only [Equiv.apply_symm_apply, i₀] at this
              cases this
            have h_some : ∀ a, Function.update f i₀ (c • x) (e.symm (some a)) = f (e.symm (some a)) :=
              fun a => Function.update_of_ne (hne a) (c • x) f
            have h_some_x : ∀ a, Function.update f i₀ x (e.symm (some a)) = f (e.symm (some a)) :=
              fun a => Function.update_of_ne (hne a) x f
            simp only [h_some, h_some_x]
            rw [map_smul]
          · have h_none : Function.update f i (c • x) i₀ = f i₀ := Function.update_of_ne (Ne.symm hi) (c • x) f
            have h_none_x : Function.update f i x i₀ = f i₀ := Function.update_of_ne (Ne.symm hi) x f
            simp only [h_none, h_none_x]
            have hi' : ∃ a, i = e.symm (some a) := by
              cases he : e i with
              | none => exfalso; apply hi; simp only [i₀]; rw [← he]; simp
              | some a => exact ⟨a, by simp [← he]⟩
            obtain ⟨a₀, ha₀⟩ := hi'
            subst ha₀
            have hM'_smul := M'.map_update_smul (fun a => f (e.symm (some a))) a₀ c x
            simp only at hM'_smul ⊢
            have h_update : ∀ z : V' a₀ ⊗[K] W' a₀, (fun a => Function.update f (e.symm (some a₀)) z (e.symm (some a))) =
                Function.update (fun a => f (e.symm (some a))) a₀ z := by
              intro z; ext a
              by_cases ha : a = a₀
              · subst ha; simp [Function.update_self]
              · rw [Function.update_of_ne ha, Function.update_of_ne]
                intro h; apply ha; exact Option.some_injective α (e.symm.injective h)
            simp only [h_update]
            rw [hM'_smul, map_smul]
            have lift_smul : ∀ (c : K) (g : V i₀ →ₗ[K] W i₀ →ₗ[K] PiTensorProduct K V ⊗[K] PiTensorProduct K W),
                TensorProduct.lift (c • g) = c • TensorProduct.lift g := by
              intro c g
              change TensorProduct.uncurry _ _ _ _ (c • g) = c • TensorProduct.uncurry _ _ _ _ g
              rw [map_smul]
            rw [lift_smul, LinearMap.smul_apply]
        case spec =>
          intro v w
          simp only [MultilinearMap.coe_mk]
          let v' : ∀ a : α, V' a := fun a => v (e.symm (some a))
          let w' : ∀ a : α, W' a := fun a => w (e.symm (some a))
          have ih_eq : M' (fun a => v' a ⊗ₜ[K] w' a) = (PiTensorProduct.tprod K) v' ⊗ₜ[K] (PiTensorProduct.tprod K) w' := hM' v' w'
          conv_lhs => simp only [TensorProduct.lift.tmul, LinearMap.coe_mk, AddHom.coe_mk, ih_eq]
          rw [ih_eq, TensorProduct.lift.tmul, LinearMap.coe_mk, AddHom.coe_mk]
          show combineV_fun (v i₀) ((PiTensorProduct.tprod K) v') ⊗ₜ[K] combineW_fun (w i₀) ((PiTensorProduct.tprod K) w') =
               (PiTensorProduct.tprod K) v ⊗ₜ[K] (PiTensorProduct.tprod K) w
          have hcombineV : combineV_fun (v i₀) ((PiTensorProduct.tprod K) v') = (PiTensorProduct.tprod K) v := by
            simp only [combineV_fun, subsingleV, PiTensorProduct.subsingletonEquiv_symm_apply']
            rw [show ((PiTensorProduct.tprod K) v' : PiTensorProduct K V') =
                (PiTensorProduct.tprod K (fun a => v' a) : ⨂[K] (a : α), Vsum (Sum.inl a)) from rfl]
            rw [show ((PiTensorProduct.tprod K) (fun _ : PUnit => v i₀) : ⨂[K] (_ : PUnit), V i₀) =
                (PiTensorProduct.tprod K (fun _ : PUnit => v i₀) : ⨂[K] (u : PUnit), Vsum (Sum.inr u)) from rfl]
            rw [PiTensorProduct.tmulEquivDep_apply, LinearEquiv.symm_apply_eq, PiTensorProduct.reindex_tprod]
            congr 1; ext s
            cases s with
            | inl a => simp only [eSum, Equiv.symm_trans_apply, optSum, Equiv.optionEquivSumPUnit_symm_inl, v']
            | inr u => simp only [eSum, Equiv.symm_trans_apply, optSum, Equiv.optionEquivSumPUnit_symm_inr, i₀]
          have hcombineW : combineW_fun (w i₀) ((PiTensorProduct.tprod K) w') = (PiTensorProduct.tprod K) w := by
            simp only [combineW_fun, subsingleW, PiTensorProduct.subsingletonEquiv_symm_apply']
            rw [show ((PiTensorProduct.tprod K) w' : PiTensorProduct K W') =
                (PiTensorProduct.tprod K (fun a => w' a) : ⨂[K] (a : α), Wsum (Sum.inl a)) from rfl]
            rw [show ((PiTensorProduct.tprod K) (fun _ : PUnit => w i₀) : ⨂[K] (_ : PUnit), W i₀) =
                (PiTensorProduct.tprod K (fun _ : PUnit => w i₀) : ⨂[K] (u : PUnit), Wsum (Sum.inr u)) from rfl]
            rw [PiTensorProduct.tmulEquivDep_apply, LinearEquiv.symm_apply_eq, PiTensorProduct.reindex_tprod]
            have heq : (fun (s : α ⊕ PUnit) => Sum.rec w' (fun _ => w i₀) s) =
                       (fun (s : α ⊕ PUnit) => w (eSum.symm s)) := by funext s; cases s <;> rfl
            congr 1
          rw [hcombineV, hcombineW]
      exact hM

end AsymptoticSpectra.Tensor