import PL1.ALFragments
import PL1.BalancedFragmentTree
import PL1.DyadicL1Coordinates
import PL1.Definitions
import Mathlib.Analysis.Normed.Operator.Completeness
import Mathlib.Logic.Equiv.Fin.Basic

/-! The interval L1 copy generated by an atomless balanced fragment tree. -/
noncomputable section
set_option backward.isDefEq.respectTransparency false
open scoped BigOperators Topology
open Filter
namespace PL1.ALDyadicEmbedding
open PL1.BalancedFragmentTree PL1.FiniteSumming PL1.DyadicTrace PL1.ALOrder

/-- Binary word ordering matching the interval children 2i and 2i+1. -/
def wordEquiv : (n : ℕ) → Fin (2 ^ n) ≃ Word n
  | 0 => {
      toFun := fun _ i => Fin.elim0 i
      invFun := fun _ => 0
      left_inv := fun i => by apply Fin.ext; have hi := i.isLt; norm_num at hi; omega
      right_inv := fun _ => Subsingleton.elim _ _ }
  | n + 1 => (finCongr (pow_succ 2 n)).trans
      ((finProdFinEquiv.symm.trans (Equiv.prodCongr (wordEquiv n) finTwoEquiv)).trans
        (snocEquiv n))

def childIndex (n : ℕ) (i : Fin (2 ^ n)) (b : Bool) : Fin (2 ^ (n + 1)) :=
  (finCongr (pow_succ 2 n)).symm (finProdFinEquiv (i, finTwoEquiv.symm b))

@[simp] theorem wordEquiv_childIndex (n : ℕ) (i : Fin (2 ^ n)) (b : Bool) :
    wordEquiv (n + 1) (childIndex n i b) = Fin.snoc (wordEquiv n i) b := by
  simp only [wordEquiv, childIndex, Equiv.trans_apply, Equiv.apply_symm_apply,
    Equiv.symm_apply_apply, Equiv.prodCongr_apply]
  simp [snocEquiv]

@[simp] theorem childIndex_false_val (n : ℕ) (i : Fin (2 ^ n)) :
    (childIndex n i false).val = 2 * i.val := by
  simp only [childIndex, finCongr_symm_apply_coe]
  simp [finProdFinEquiv, finTwoEquiv]

@[simp] theorem childIndex_true_val (n : ℕ) (i : Fin (2 ^ n)) :
    (childIndex n i true).val = 2 * i.val + 1 := by
  simp only [childIndex, finCongr_symm_apply_coe]
  simp [finProdFinEquiv, finTwoEquiv, Nat.add_comm]

variable {E : Type*} [NormedAddCommGroup E] [Lattice E]
  [IsOrderedAddMonoid E] [NormedVectorLattice E]

theorem exists_balancedSplit_of_atomless [CompleteSpace E] (hAL : ALRenorming.IsALNorm E)
    (u : E) (hna : ∀ a : E, IsLatticeAtom a → ¬ a ≤ u) :
    ∃ split : E → E, IsBalancedSplit (normFunctional hAL) u split := by
  apply exists_balancedSplit
  intro x hx hxu
  exact ALFragments.exists_half_fragment hAL (normFunctional hAL)
    (normFunctional_strictlyPositive hAL) hx
    (fun a ha hax => hna a ha (hax.trans hxu))

def node (u : E) (split : E → E) (n : ℕ) (i : Fin (2 ^ n)) : E :=
  BalancedFragmentTree.cell u split n (wordEquiv n i)

omit [Lattice E] [IsOrderedAddMonoid E] [NormedVectorLattice E] in
theorem node_children_add (u : E) (split : E → E) (n : ℕ) (i : Fin (2 ^ n)) :
    node u split (n + 1) (childIndex n i false) +
      node u split (n + 1) (childIndex n i true) = node u split n i := by
  simp only [node, wordEquiv_childIndex, cell_children_add]

theorem node_nonneg {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split)
    (n : ℕ) (i : Fin (2 ^ n)) : 0 ≤ node u split n i :=
  (cell_mem_Icc hu hs n _).1

theorem norm_node {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split)
    (n : ℕ) (i : Fin (2 ^ n)) : ‖node u split n i‖ = ‖u‖ / (2 : ℝ)^n := by
  rw [← normFunctional_nonneg hAL (node_nonneg hu hs n i)]
  exact (cell_phi hu hs n _).trans (by rw [normFunctional_nonneg hAL hu])

theorem node_disjoint {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (n : ℕ) :
    Pairwise (fun i j : Fin (2 ^ n) => IsVLDisjoint (node u split n i) (node u split n j)) :=
  fun _ _ hij => cell_pairwise_disjoint hu hs n (fun h => hij ((wordEquiv n).injective h))

def treeSynthesis (u : E) (split : E → E) (n : ℕ) : EllOne (2 ^ n) →L[ℝ] E :=
  ∑ i : Fin (2 ^ n), (PiLp.proj 1 (fun _ : Fin (2 ^ n) => ℝ) i).smulRight
    ((2 : ℝ)^n • node u split n i)

theorem treeSynthesis_apply (u : E) (split : E → E) (n : ℕ) (x : EllOne (2 ^ n)) :
    treeSynthesis u split n x = ∑ i, ((2 : ℝ)^n * x i) • node u split n i := by
  simp [treeSynthesis, smul_smul, mul_comm]

theorem norm_treeSynthesis {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split)
    (n : ℕ) (x : EllOne (2 ^ n)) : ‖treeSynthesis u split n x‖ = ‖u‖ * ‖x‖ := by
  rw [treeSynthesis_apply, ← norm_abs_eq_norm]
  rw [← normFunctional_nonneg hAL (abs_nonneg _),
    abs_sum_smul_disjoint _ _ (node_disjoint hu hs n), map_sum]
  simp only [map_smul, smul_eq_mul, abs_of_nonneg (node_nonneg hu hs n _)]
  simp_rw [normFunctional_nonneg hAL (node_nonneg hu hs n _), norm_node hu hs]
  rw [norm_ellOne, Finset.mul_sum]
  apply Finset.sum_congr rfl
  intro i _
  rw [abs_mul, abs_of_nonneg (pow_nonneg (by norm_num : (0 : ℝ) ≤ 2) n)]
  field_simp

 theorem norm_treeSynthesis_le {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (n : ℕ) :
    ‖treeSynthesis u split n‖ ≤ ‖u‖ := by
  apply ContinuousLinearMap.opNorm_le_bound _ (norm_nonneg u)
  intro x
  exact (norm_treeSynthesis hu hs n x).le






open PL1.DyadicL1Coordinates

theorem synthesis_single (n : ℕ) (i : Fin (2 ^ n)) :
    synthesis n (PiLp.single 1 i 1) = normalizedVector n i := by
  rw [synthesis_apply]
  simp [PiLp.single_apply]

theorem treeSynthesis_single (u : E) (split : E → E) (n : ℕ) (i : Fin (2 ^ n)) :
    treeSynthesis u split n (PiLp.single 1 i 1) = (2 : ℝ)^n • node u split n i := by
  rw [treeSynthesis_apply]
  simp [PiLp.single_apply]

def treeApprox (u : E) (split : E → E) (n : ℕ) : IntervalL1 ℝ →L[ℝ] E :=
  (treeSynthesis u split n).comp (analysis n)

theorem treeApprox_normalizedVector (u : E) (split : E → E) (n : ℕ) (i : Fin (2 ^ n)) :
    treeApprox u split n (normalizedVector n i) = (2 : ℝ)^n • node u split n i := by
  rw [← synthesis_single n i]
  change treeSynthesis u split n (analysis n (synthesis n _)) = _
  rw [analysis_synthesis, treeSynthesis_single]

theorem treeApprox_vector (u : E) (split : E → E) (n : ℕ) (i : Fin (2 ^ n)) :
    treeApprox u split n (dyadicVector n i) = node u split n i := by
  have hp : (2 : ℝ)^n ≠ 0 := pow_ne_zero _ (by norm_num)
  have he : dyadicVector (𝕜 := ℝ) n i = ((2 : ℝ)^n)⁻¹ • normalizedVector n i := by
    simp [normalizedVector, smul_smul, hp]
  rw [he, map_smul, treeApprox_normalizedVector, smul_smul, inv_mul_cancel₀ hp, one_smul]

theorem norm_treeApprox {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split)
    (n : ℕ) (f : IntervalL1 ℝ) : ‖treeApprox u split n f‖ = ‖u‖ * ‖dyadicAverage n f‖ := by
  change ‖treeSynthesis u split n (analysis n f)‖ = _
  rw [norm_treeSynthesis hu hs, ← norm_synthesis n (analysis n f), synthesis_analysis]

theorem norm_treeApprox_le {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (n : ℕ) :
    ‖treeApprox u split n‖ ≤ ‖u‖ := by
  calc
    _ ≤ ‖treeSynthesis u split n‖ * ‖analysis n‖ := ContinuousLinearMap.opNorm_comp_le _ _
    _ ≤ ‖u‖ * ‖analysis n‖ := mul_le_mul_of_nonneg_right (norm_treeSynthesis_le hu hs n) (norm_nonneg _)
    _ ≤ ‖u‖ * 1 := mul_le_mul_of_nonneg_left (norm_analysis_le_one n) (norm_nonneg _)
    _ = _ := mul_one _

theorem treeApprox_stable_vector (u : E) (split : E → E) {n m : ℕ} (hnm : n ≤ m)
    (i : Fin (2 ^ n)) : treeApprox u split m (dyadicVector n i) = node u split n i := by
  induction hnm using Nat.decreasingInduction with
  | self => exact treeApprox_vector u split m i
  | of_succ n hn ih =>
    have h0 := ih (childIndex n i false)
    have h1 := ih (childIndex n i true)
    calc
      _ = treeApprox u split m
          (dyadicVector (n + 1) (childIndex n i false) +
            dyadicVector (n + 1) (childIndex n i true)) := by
        rw [childIndex_false_val, childIndex_true_val, ← dyadicVector_refine]
      _ = node u split (n + 1) (childIndex n i false) +
          node u split (n + 1) (childIndex n i true) := by rw [map_add, h0, h1]
      _ = _ := node_children_add u split n i

theorem treeApprox_average_self (u : E) (split : E → E) (n : ℕ) (f : IntervalL1 ℝ) :
    treeApprox u split n (dyadicAverage n f) = treeApprox u split n f := by
  rw [← synthesis_analysis]
  change treeSynthesis u split n (analysis n (synthesis n (analysis n f))) = _
  rw [analysis_synthesis]
  rfl

theorem treeApprox_stable_average (u : E) (split : E → E) {n m : ℕ} (hnm : n ≤ m)
    (f : IntervalL1 ℝ) : treeApprox u split m (dyadicAverage n f) = treeApprox u split n f := by
  rw [← treeApprox_average_self u split n f]
  rw [dyadicAverage, PartitionAverages.average_apply, map_sum, map_sum]
  simp_rw [map_smul]
  apply Finset.sum_congr rfl
  intro i hi
  have hi' : i < 2 ^ n := Finset.mem_range.mp hi
  rw [treeApprox_stable_vector u split hnm ⟨i, hi'⟩,
    treeApprox_vector u split n ⟨i, hi'⟩]





theorem treeApprox_cauchy {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (f : IntervalL1 ℝ) :
    CauchySeq (fun n => treeApprox u split n f) := by
  apply Metric.cauchySeq_iff'.mpr
  intro ε hε
  have ht : Tendsto (fun n => ‖u‖ * ‖f - dyadicAverage n f‖) atTop (𝓝 0) := by
    simpa only [sub_self, norm_zero, mul_zero] using
      (((tendsto_const_nhds (x := f)).sub (dyadicAverage_tendsto f)).norm).const_mul ‖u‖
  obtain ⟨N, hN⟩ := (ht.eventually_lt_const hε).exists
  refine ⟨N, fun n hn => ?_⟩
  calc
    _ = ‖treeApprox u split n (f - dyadicAverage N f)‖ := by
      rw [dist_eq_norm, map_sub, treeApprox_stable_average u split hn f]
    _ ≤ ‖treeApprox u split n‖ * ‖f - dyadicAverage N f‖ := ContinuousLinearMap.le_opNorm _ _
    _ ≤ ‖u‖ * ‖f - dyadicAverage N f‖ :=
      mul_le_mul_of_nonneg_right (norm_treeApprox_le hu hs n) (norm_nonneg _)
    _ < ε := hN

variable [CompleteSpace E]

def treeLimit {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (f : IntervalL1 ℝ) : E :=
  (cauchySeq_tendsto_of_complete (treeApprox_cauchy hu hs f)).choose

theorem treeLimit_tendsto {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (f : IntervalL1 ℝ) :
    Tendsto (fun n => treeApprox u split n f) atTop (𝓝 (treeLimit hu hs f)) :=
  (cauchySeq_tendsto_of_complete (treeApprox_cauchy hu hs f)).choose_spec

def treeEmbedding {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) : IntervalL1 ℝ →L[ℝ] E :=
  ContinuousLinearMap.ofTendstoOfBoundedRange (treeLimit hu hs) (treeApprox u split)
    (tendsto_pi_nhds.mpr (treeLimit_tendsto hu hs))
    (isBounded_iff_forall_norm_le.mpr ⟨‖u‖, by
      rintro _ ⟨n, rfl⟩
      exact norm_treeApprox_le hu hs n⟩)

theorem treeEmbedding_tendsto {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (f : IntervalL1 ℝ) :
    Tendsto (fun n => treeApprox u split n f) atTop (𝓝 (treeEmbedding hu hs f)) :=
  treeLimit_tendsto hu hs f

theorem norm_treeEmbedding {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 ≤ u) (hs : IsBalancedSplit (normFunctional hAL) u split) (f : IntervalL1 ℝ) :
    ‖treeEmbedding hu hs f‖ = ‖u‖ * ‖f‖ := by
  apply tendsto_nhds_unique (treeEmbedding_tendsto hu hs f).norm
  have ht := ((dyadicAverage_tendsto f).norm).const_mul ‖u‖
  simpa only [norm_treeApprox hu hs] using ht

/-- The fragment tree yields a genuine isometry of the actual interval L1 norm. -/
def treeIsometry {hAL : ALRenorming.IsALNorm E} {u : E} {split : E → E}
    (hu : 0 < u) (hs : IsBalancedSplit (normFunctional hAL) u split) : IntervalL1 ℝ →ₗᵢ[ℝ] E where
  toLinearMap := ((‖u‖⁻¹ : ℝ) • treeEmbedding hu.le hs).toLinearMap
  norm_map' f := by
    have hn : ‖u‖ ≠ 0 := norm_ne_zero_iff.mpr (ne_of_gt hu)
    change ‖(‖u‖⁻¹ : ℝ) • treeEmbedding hu.le hs f‖ = ‖f‖
    rw [norm_smul, Real.norm_eq_abs, abs_of_nonneg (inv_nonneg.mpr (norm_nonneg u)),
      norm_treeEmbedding, ← mul_assoc, inv_mul_cancel₀ hn, one_mul]

theorem containsIsomorphicCopy_of_atomless_vector (hAL : ALRenorming.IsALNorm E)
    (u : E) (hu : 0 < u) (hna : ∀ a : E, IsLatticeAtom a → ¬ a ≤ u) :
    ContainsIsomorphicCopy ℝ (IntervalL1 ℝ) E := by
  obtain ⟨split, hs⟩ := exists_balancedSplit_of_atomless hAL u hna
  let J := treeIsometry hu hs
  have hc : IsClosed (J.toContinuousLinearMap.range : Set E) := J.isometry.isClosedEmbedding.isClosed_range
  exact ⟨J.toContinuousLinearMap.range, hc,
    ⟨J.toContinuousLinearMap.equivRange J.injective hc⟩⟩

theorem containsIsomorphicCopy_of_not_atomic (hAL : ALRenorming.IsALNorm E)
    (hna : ¬ IsAtomicLattice E) : ContainsIsomorphicCopy ℝ (IntervalL1 ℝ) E := by
  rw [IsAtomicLattice] at hna
  push Not at hna
  obtain ⟨u, hu, hna⟩ := hna
  exact containsIsomorphicCopy_of_atomless_vector hAL u hu (fun a ha hau => hna a ha hau)

#print axioms containsIsomorphicCopy_of_atomless_vector
#print axioms containsIsomorphicCopy_of_not_atomic
end PL1.ALDyadicEmbedding