p2m_TensorObj_p3

import Mathlib.LinearAlgebra.PiTensorProduct
import Mathlib.LinearAlgebra.FiniteDimensional.Basic
import Mathlib.RingTheory.TensorProduct.Finite
import Mathlib.LinearAlgebra.TensorProduct.Finiteness
import Mathlib.Logic.Basic
import Mathlib.Algebra.Module.ULift
import Mathlib.LinearAlgebra.Prod
import Mathlib.Algebra.Module.PUnit
import Mathlib.LinearAlgebra.TensorProduct.Prod
import Mathlib.LinearAlgebra.TensorProduct.Tower
import Definitions.Def_p2m_TensorObj_p1
import Definitions.Def_p2m_TensorObj_p2

universe u v w

open BigOperators TensorProduct
open PiTensorProduct

set_option maxHeartbeats 1000000

open TensorObj

namespace TensorObj

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

/-- Multiplication respects mutual restriction. -/
theorem mul_isomorphic {X Y Z W : TensorObj K d}
    (h1 : Restrict X Y ∧ Restrict Y X) (h2 : Restrict Z W ∧ Restrict W Z) :
    Restrict (X * Z) (Y * W) ∧ Restrict (Y * W) (X * Z) := by
  obtain ⟨⟨f, hf⟩, ⟨f', hf'⟩⟩ := h1
  obtain ⟨⟨g, hg⟩, ⟨g', hg'⟩⟩ := h2
  constructor
  · exact ⟨fun i => TensorProduct.map (f i) (g i), by
      simp only [mul_t]
      erw [liftMap_interchange, hf, hg]⟩
  · exact ⟨fun i => TensorProduct.map (f' i) (g' i), by
      simp only [mul_t]
      erw [liftMap_interchange, hf', hg']⟩

private theorem mul_comm_restrict_aux (X Y : TensorObj K d)
    (x : PiTensorProduct K X.V) (y : PiTensorProduct K Y.V) :
    (liftMap (fun i => (TensorProduct.comm K (Y.V i) (X.V i)).toLinearMap)) (interchange y x) =
    interchange x y := by
  induction x using PiTensorProduct.induction_on with
  | smul_tprod c v =>
    simp only [map_smul]
    induction y using PiTensorProduct.induction_on with
    | smul_tprod c' v' =>
      dsimp only [interchange, interchangeMap, liftMap, interchangeAux]
      simp only [map_smul, LinearMap.smul_apply, smul_smul]
      simp only [PiTensorProduct.lift.tprod, MultilinearMap.compLinearMap_apply, MultilinearMap.coe_mk]
      rw [mul_comm c c']; congr 1
    | add y1 y2 ih1 ih2 =>
      simp only [map_add, LinearMap.add_apply, smul_add, ih1, ih2]
  | add x1 x2 ih1 ih2 =>
    simp only [LinearMap.add_apply, map_add, ih1, ih2]

private theorem mul_comm_restrict (X Y : TensorObj K d) : Restrict (X * Y) (Y * X) :=
  ⟨fun i => (TensorProduct.comm K (Y.V i) (X.V i)).toLinearMap, by
    simp only [mul_t]
    exact mul_comm_restrict_aux X Y X.t Y.t⟩

theorem mul_comm_isomorphic {X Y : TensorObj K d} :
    Restrict (X * Y) (Y * X) ∧ Restrict (Y * X) (X * Y) :=
  ⟨mul_comm_restrict X Y, mul_comm_restrict Y X⟩

/-- Associativity of the interchange map. -/
@[simp]
theorem interchange_assoc {ι : Type*} [Fintype ι] [DecidableEq ι]
    {V W U : ι → Type*}
    [∀ i, AddCommGroup (V i)] [∀ i, Module K (V i)]
    [∀ i, AddCommGroup (W i)] [∀ i, Module K (W i)]
    [∀ i, AddCommGroup (U i)] [∀ i, Module K (U i)]
    (tV : PiTensorProduct K V) (tW : PiTensorProduct K W) (tU : PiTensorProduct K U) :
    liftMap (fun i => (TensorProduct.assoc K (V i) (W i) (U i)).toLinearMap)
      (interchange (interchange tV tW) tU) =
    interchange tV (interchange tW tU) := by
  induction tV using PiTensorProduct.induction_on with
  | smul_tprod c v =>
    simp only [map_smul, interchange, liftMap]
    induction tW using PiTensorProduct.induction_on with
    | smul_tprod c' w =>
      simp only [map_smul]
      induction tU using PiTensorProduct.induction_on with
      | smul_tprod c'' u =>
        dsimp [interchange, liftMap, interchangeMap, interchangeAux]
        simp only [map_smul, LinearMap.smul_apply, smul_smul]
        simp only [PiTensorProduct.lift.tprod, MultilinearMap.coe_mk, MultilinearMap.compLinearMap_apply]
        simp only [_root_.mul_comm, _root_.mul_left_comm]
        congr
      | add x y ih1 ih2 =>
        simp only [map_add, ih1, ih2]
    | add x y ih1 ih2 =>
      simp only [map_add, LinearMap.add_apply, ih1, ih2]
  | add x y ih1 ih2 =>
    simp only [map_add, LinearMap.add_apply, ih1, ih2]

-- Restrict (X*Y*Z) (X*(Y*Z)): f : (X*(Y*Z)).V i →ₗ (X*Y*Z).V i = assoc.symm
private theorem mul_assoc_restrict_fwd (X Y Z : TensorObj K d) : Restrict (X * Y * Z) (X * (Y * Z)) :=
  ⟨fun i => (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).symm.toLinearMap, by
    change liftMap (fun i => (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).symm.toLinearMap)
      (interchange X.t (interchange Y.t Z.t)) = interchange (interchange X.t Y.t) Z.t
    rw [← interchange_assoc, liftMap_comp]
    have : (fun i => (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).symm.toLinearMap ∘ₗ
            (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).toLinearMap) =
           fun _ => LinearMap.id := by
      funext i; apply LinearMap.ext; intro v; simp
    rw [this]; exact liftMap_id (X * Y * Z)⟩

-- Restrict (X*(Y*Z)) (X*Y*Z): f : (X*Y*Z).V i →ₗ (X*(Y*Z)).V i = assoc
theorem mul_assoc_isomorphic {X Y Z : TensorObj K d} :
    Restrict (X * Y * Z) (X * (Y * Z)) ∧ Restrict (X * (Y * Z)) (X * Y * Z) :=
  ⟨mul_assoc_restrict_fwd X Y Z,
   ⟨fun i => (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).toLinearMap, by
    change liftMap (fun i => (TensorProduct.assoc K (X.V i) (Y.V i) (Z.V i)).toLinearMap)
      (interchange (interchange X.t Y.t) Z.t) = interchange X.t (interchange Y.t Z.t)
    rw [interchange_assoc]⟩⟩

-- Restrict (X + zeroObj) X: embed X into X+zeroObj via inl
-- f : X.V i →ₗ (X+zeroObj).V i, need liftMap inl X.t = (X+zeroObj).t
private theorem add_zero_restrict_fwd (X : TensorObj K d) : Restrict (X + zeroObj) X :=
  ⟨fun i => LinearMap.inl K (X.V i) PUnit, by
    simp only [add_t, zeroObj, map_zero, add_zero]; rfl⟩

-- Restrict X (X + zeroObj): project X+zeroObj to X via fst
-- f : (X+zeroObj).V i →ₗ X.V i, need liftMap fst (X+zeroObj).t = X.t
theorem add_zero_isomorphic {X : TensorObj K d} :
    Restrict (X + zeroObj) X ∧ Restrict X (X + zeroObj) :=
  ⟨add_zero_restrict_fwd X,
   ⟨fun i => LinearMap.fst K (X.V i) PUnit, by
    have : liftMap (fun i => LinearMap.fst K (X.V i) PUnit) ((X + zeroObj).t) = X.t := by
      simp only [add_t, zeroObj, map_zero, add_zero]
      have heq : (fun i => (LinearMap.fst K (X.V i) PUnit).comp (LinearMap.inl K (X.V i) PUnit)) =
                 fun _ => LinearMap.id := by funext i; ext; rfl
      rw [liftMap_comp, heq, liftMap_id]
    exact this⟩⟩

-- Restrict (zeroObj + X) X: embed X into zeroObj+X via inr
-- f : X.V i →ₗ (zeroObj+X).V i, need liftMap inr X.t = (zeroObj+X).t
private theorem zero_add_restrict_fwd (X : TensorObj K d) : Restrict (zeroObj + X) X :=
  ⟨fun i => LinearMap.inr K PUnit (X.V i), by
    simp only [add_t, zeroObj, map_zero, zero_add]; rfl⟩

-- Restrict X (zeroObj + X): project via snd
-- f : (zeroObj+X).V i →ₗ X.V i, need liftMap snd (zeroObj+X).t = X.t
theorem zero_add_isomorphic {X : TensorObj K d} :
    Restrict (zeroObj + X) X ∧ Restrict X (zeroObj + X) :=
  ⟨zero_add_restrict_fwd X,
   ⟨fun i => LinearMap.snd K PUnit (X.V i), by
    have : liftMap (fun i => LinearMap.snd K PUnit (X.V i)) ((zeroObj + X).t) = X.t := by
      simp only [add_t, zeroObj, map_zero, zero_add]
      have heq : (fun i => (LinearMap.snd K PUnit (X.V i)).comp (LinearMap.inr K PUnit (X.V i))) =
                 fun _ => LinearMap.id := by funext i; ext; rfl
      rw [liftMap_comp, heq, liftMap_id]
    exact this⟩⟩

-- one_mul: oneObj * X ~ X
-- Restrict (oneObj * X) X: f : X.V i →ₗ (oneObj*X).V i = ULift K ⊗ X.V i
-- We prove this by showing liftMap tensorOne.symm X.t = interchange oneObj.t X.t
-- via: liftMap tensorOne (liftMap tensorOne.symm X.t) = X.t = liftMap tensorOne (interchange ...)
private theorem one_mul_restrict_fwd (X : TensorObj K d) : Restrict (oneObj * X) X := by
  refine ⟨fun i => tensorOne.symm.toLinearMap, ?_⟩
  simp only [mul_t]
  have hTO_symm_TO : (fun i => (tensorOne (K := K) (V := X.V i)).symm.toLinearMap ∘ₗ
      (tensorOne (K := K) (V := X.V i)).toLinearMap) = fun _ => LinearMap.id := by
    funext i; apply LinearMap.ext; intro v; simp [tensorOne]
  have hTO_TO_symm : (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap ∘ₗ
      (tensorOne (K := K) (V := X.V i)).symm.toLinearMap) = fun _ => LinearMap.id := by
    funext i; apply LinearMap.ext; intro v; simp [tensorOne]
  -- suffices: liftMap tensorOne (liftMap tensorOne.symm X.t) = liftMap tensorOne (interchange oneObj.t X.t)
  -- since liftMap tensorOne is injective (has left inverse liftMap tensorOne.symm)
  suffices liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap)
      (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).symm.toLinearMap) X.t) =
      liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap) (interchange oneObj.t X.t) by
    have hinj : Function.Injective
        (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap)) := by
      intro a b hab
      have := congr_arg (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).symm.toLinearMap)) hab
      rw [liftMap_comp, liftMap_comp] at this
      rw [hTO_symm_TO] at this
      unfold liftMap at this; simp at this
      exact this
    exact hinj this
  rw [liftMap_comp, hTO_TO_symm, liftMap_id]
  -- Now prove X.t = liftMap tensorOne (interchange oneObj.t X.t)
  symm
  induction X.t using PiTensorProduct.induction_on with
  | smul_tprod c v =>
    change liftMap (fun i => tensorOne.toLinearMap)
      (interchange (tprod K (fun _ => ULift.up 1)) (c • tprod K v)) = c • tprod K v
    dsimp only [interchange, liftMap, interchangeMap, interchangeAux]
    simp only [map_smul, PiTensorProduct.lift.tprod, MultilinearMap.coe_mk,
               MultilinearMap.compLinearMap_apply]
    congr 1; apply congr_arg; funext i
    simp [tensorOne, uliftEquiv, TensorProduct.lid, TensorProduct.congr]
  | add t1 t2 ih1 ih2 =>
    simp only [(interchange oneObj.t).map_add]
    rw [show (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap)
                ((interchange oneObj.t) t1 + (interchange oneObj.t) t2)) =
            (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap)) ((interchange oneObj.t) t1) +
            (liftMap (fun i => (tensorOne (K := K) (V := X.V i)).toLinearMap)) ((interchange oneObj.t) t2)
          from (liftMap _).map_add _ _]
    rw [ih1, ih2]

-- Restrict X (oneObj * X): f : (oneObj*X).V i = ULift K ⊗ X.V i →ₗ X.V i, i.e. tensorOne
-- liftMap tensorOne (interchange oneObj.t X.t) = X.t
theorem one_mul_isomorphic {X : TensorObj K d} :
    Restrict (oneObj * X) X ∧ Restrict X (oneObj * X) :=
  ⟨one_mul_restrict_fwd X,
   ⟨fun i => tensorOne.toLinearMap, by
    change liftMap (fun i => tensorOne.toLinearMap) (interchange oneObj.t X.t) = X.t
    induction X.t using PiTensorProduct.induction_on with
    | smul_tprod c v =>
      change liftMap (fun i => tensorOne.toLinearMap) (interchange (tprod K (fun _ => ULift.up 1)) (c • tprod K v)) = c • tprod K v
      dsimp only [interchange, liftMap, interchangeMap, interchangeAux]
      simp only [map_smul, PiTensorProduct.lift.tprod, MultilinearMap.coe_mk, MultilinearMap.compLinearMap_apply]
      congr 1; apply congr_arg; funext i
      simp [tensorOne, uliftEquiv, TensorProduct.lid, TensorProduct.congr]
    | add t1 t2 ih1 ih2 =>
      dsimp only [interchange, liftMap] at ih1 ih2 ⊢
      simp only [map_add, ih1, ih2]⟩⟩

theorem mul_one_isomorphic {X : TensorObj K d} :
    Restrict (X * oneObj) X ∧ Restrict X (X * oneObj) :=
  let hc := mul_comm_isomorphic (X := X) (Y := oneObj)
  let h1 := one_mul_isomorphic (X := X)
  ⟨restrict_trans hc.1 h1.1, restrict_trans h1.2 hc.2⟩

-- Restrict (X*(Y+Z)) (X*Y+X*Z): f : (X*Y+X*Z).V i →ₗ (X*(Y+Z)).V i = distribLeft.symm
-- liftMap distribLeft.symm (X*Y+X*Z).t = (X*(Y+Z)).t

end TensorObj