-- QFD/Cosmology/AxisExtraction.lean
import Mathlib.Analysis.InnerProductSpace.PiL2
import Mathlib.Analysis.InnerProductSpace.Basic
import Mathlib.Analysis.InnerProductSpace.Projection.FiniteDimensional
import Mathlib.Algebra.Ring.Parity
import Mathlib.Tactic.Ring
import Mathlib.Tactic.Linarith

noncomputable section

namespace QFD.Cosmology

open scoped BigOperators
open Real

/--
Use PiLp to get a canonical inner product / norm on functions `Fin 3 → ℝ`.
This avoids some `EuclideanSpace` alias/coercion quirks across mathlib commits.
-/
abbrev R3 := PiLp 2 (fun _ : Fin 3 => ℝ)

/-- Unit predicate (kept minimal; avoids importing `Metric.sphere`). -/
def IsUnit (x : R3) : Prop := ‖x‖ = 1

/-- Legendre P₂. -/
def P2 (t : ℝ) : ℝ := (3 * t^2 - 1) / 2

/-
COMMIT-ROBUST: avoid ⟪n,x⟫_ℝ notation completely.
If `inner` is not in scope for any reason, use `Inner.inner`.
-/
def ip (n x : R3) : ℝ := inner ℝ n x

/-- Axis score: square of inner product (captures ± symmetry). -/
def score (n x : R3) : ℝ := (ip n x)^2

/-- Quadrupole pattern about axis `n` evaluated at direction `x`. -/
def quadPattern (n x : R3) : ℝ := P2 (ip n x)

/-- "Argmax set" on the unit sphere. -/
def AxisSet (f : R3 → ℝ) : Set R3 :=
  {x | IsUnit x ∧ ∀ y, IsUnit y → f y ≤ f x}

lemma quadPattern_eq_affine_score (n x : R3) :
    quadPattern n x = (3/2) * score n x - 1/2 := by
  simp [quadPattern, P2, score, ip]
  ring

/-- Cauchy–Schwarz bound specialized to unit vectors: score ≤ 1. -/
lemma score_le_one_of_unit (n x : R3) (hn : IsUnit n) (hx : IsUnit x) :
    score n x ≤ 1 := by
  have habs : |inner ℝ n x| ≤ ‖n‖ * ‖x‖ := abs_real_inner_le_norm n x
  have habs' : |inner ℝ n x| ≤ 1 := by
    rw [hn, hx] at habs
    simpa using habs
  -- Square both sides: |a|^2 = a^2 and |a| ≤ 1 → a^2 ≤ 1
  have hsq : (inner ℝ n x)^2 ≤ 1 := by
    calc (inner ℝ n x)^2
        = |inner ℝ n x|^2 := by rw [sq_abs]
      _ ≤ 1^2 := by
          apply sq_le_sq'
          · linarith [abs_nonneg (inner ℝ n x)]
          · exact habs'
      _ = 1 := by norm_num
  simpa [score, ip] using hsq

/-- `n` is a maximizer of its own quadrupole pattern on the unit sphere. -/
theorem n_mem_AxisSet_quadPattern (n : R3) (hn : IsUnit n) :
    n ∈ AxisSet (quadPattern n) := by
  refine ⟨hn, ?_⟩
  intro y hy
  have h1 : score n y ≤ 1 := score_le_one_of_unit n y hn hy
  have h2 : score n n = 1 := by
    -- ⟪n,n⟫ = ‖n‖^2, so (⟪n,n⟫)^2 = 1 for unit n
    -- This lemma name is stable in most builds:
    have h_inner : inner ℝ n n = ‖n‖^2 := real_inner_self_eq_norm_sq n
    unfold score ip
    rw [h_inner, hn]
    norm_num
  -- Compare via the affine form (monotone in score since 3/2 > 0).
  -- Goal is quadPattern n y ≤ quadPattern n n.
  -- Use the affine identity and linarith.
  have : (3/2:ℝ) * score n y - 1/2 ≤ (3/2:ℝ) * score n n - 1/2 := by
    have h_nonneg : (0:ℝ) ≤ 3/2 := by norm_num
    linarith [mul_le_mul_of_nonneg_left h1 h_nonneg]
  rw [quadPattern_eq_affine_score, quadPattern_eq_affine_score]
  exact this

/--
**Micro-improvement A: -n is also a maximizer**

The quadrupole pattern is even: score n (-n) = score n n, so both n and -n
are in the argmax set. This establishes "axis defined up to sign."
-/
lemma score_neg (n x : R3) : score n (-x) = score n x := by
  unfold score ip
  simp [inner_neg_right]

theorem neg_n_mem_AxisSet_quadPattern (n : R3) (hn : IsUnit n) :
    -n ∈ AxisSet (quadPattern n) := by
  have h_neg_unit : IsUnit (-n) := by
    unfold IsUnit at *
    rw [norm_neg]
    exact hn
  refine ⟨h_neg_unit, ?_⟩
  intro y hy
  have h1 : score n y ≤ 1 := score_le_one_of_unit n y hn hy
  have h2 : score n (-n) = 1 := by
    rw [score_neg]
    have h_inner : inner ℝ n n = ‖n‖^2 := real_inner_self_eq_norm_sq n
    unfold score ip
    rw [h_inner, hn]
    norm_num
  have : (3/2:ℝ) * score n y - 1/2 ≤ (3/2:ℝ) * score n (-n) - 1/2 := by
    have h_nonneg : (0:ℝ) ≤ 3/2 := by norm_num
    linarith [mul_le_mul_of_nonneg_left h1 h_nonneg]
  rw [quadPattern_eq_affine_score, quadPattern_eq_affine_score]
  exact this

/--
**Micro-improvement B: Argmax invariance under positive affine transforms**

The argmax set is invariant under positive affine transformations.
This connects cos²θ → P₂ → score without re-proving monotonicity each time.
-/
lemma AxisSet_affine (f : R3 → ℝ) (a b : ℝ) (ha : 0 < a) :
    AxisSet (fun x => a * f x + b) = AxisSet f := by
  ext x
  unfold AxisSet IsUnit
  constructor
  · intro ⟨hx_unit, hx_max⟩
    refine ⟨hx_unit, ?_⟩
    intro y hy
    have h : a * f y + b ≤ a * f x + b := hx_max y hy
    -- Cancel: a * f y + b ≤ a * f x + b → a * f y ≤ a * f x → f y ≤ f x
    have h1 : a * f y ≤ a * f x := by linarith
    exact le_of_mul_le_mul_left h1 ha
  · intro ⟨hx_unit, hx_max⟩
    refine ⟨hx_unit, ?_⟩
    intro y hy
    have h : f y ≤ f x := hx_max y hy
    linarith [mul_le_mul_of_nonneg_left h (le_of_lt ha)]

/--
**Monotone transform invariance**

If `g : ℝ → ℝ` is strictly monotone increasing, then applying `g` to a
function `f` preserves its argmax set on the unit sphere.

This generalizes `AxisSet_affine` and makes the framework more robust:
any strictly increasing transformation preserves which directions maximize.
-/
lemma AxisSet_monotone (f : R3 → ℝ) (g : ℝ → ℝ)
    (hg : StrictMono g) :
    AxisSet (g ∘ f) = AxisSet f := by
  ext x
  unfold AxisSet
  constructor
  · intro ⟨hx_unit, hx_max⟩
    refine ⟨hx_unit, ?_⟩
    intro y hy
    have h : g (f y) ≤ g (f x) := hx_max y hy
    exact hg.le_iff_le.mp h
  · intro ⟨hx_unit, hx_max⟩
    refine ⟨hx_unit, ?_⟩
    intro y hy
    have h : f y ≤ f x := hx_max y hy
    exact hg.monotone h

/-! ## Phase 2: Full Uniqueness (AxisSet = {n, -n}) -/

/-- Helper: ip = 1 iff x = n (for unit vectors) -/
lemma ip_eq_one_iff_eq (n x : R3) (hn : IsUnit n) (hx : IsUnit x) :
    ip n x = 1 ↔ n = x := by
  unfold ip IsUnit at *
  exact inner_eq_one_iff_of_norm_eq_one hn hx

/-- Helper: ip = -1 iff x = -n (for unit vectors) -/
lemma ip_eq_neg_one_iff_eq_neg (n x : R3) (hn : IsUnit n) (hx : IsUnit x) :
    ip n x = -1 ↔ n = -x := by
  unfold ip IsUnit at *
  exact inner_eq_neg_one_iff_of_norm_eq_one hn hx

/-- Upper bound: quadPattern ≤ 1 on the unit sphere -/
lemma quadPattern_le_one_of_unit (n x : R3) (hn : IsUnit n) (hx : IsUnit x) :
    quadPattern n x ≤ 1 := by
  have hscore : score n x ≤ 1 := score_le_one_of_unit n x hn hx
  have h_nonneg : (0 : ℝ) ≤ (3/2 : ℝ) := by norm_num
  have hmul : (3/2 : ℝ) * score n x ≤ (3/2 : ℝ) * 1 :=
    mul_le_mul_of_nonneg_left hscore h_nonneg
  have hsub : (3/2 : ℝ) * score n x - 1/2 ≤ (3/2 : ℝ) * 1 - 1/2 :=
    sub_le_sub_right hmul (1/2 : ℝ)
  calc quadPattern n x
      = (3/2 : ℝ) * score n x - 1/2 := quadPattern_eq_affine_score n x
    _ ≤ (3/2 : ℝ) * 1 - 1/2 := hsub
    _ = 1 := by norm_num

/--
**Phase 2: Full Uniqueness**

The argmax set of the quadrupole pattern on the unit sphere is exactly {n, -n}.
This completes the characterization: the "Axis of Evil" (if it follows the QFD
predicted pattern) is the motion vector, defined up to sign.
-/
theorem AxisSet_quadPattern_eq_pm (n : R3) (hn : IsUnit n) :
    AxisSet (quadPattern n) = {x | x = n ∨ x = -n} := by
  ext x
  unfold AxisSet IsUnit at *
  constructor
  · intro ⟨hx_unit, hx_max⟩
    have h_ge : quadPattern n n ≤ quadPattern n x := hx_max n hn
    have h_le : quadPattern n x ≤ 1 := quadPattern_le_one_of_unit n x hn hx_unit
    have hnn : quadPattern n n = 1 := by
      have h_inner : inner ℝ n n = ‖n‖^2 := real_inner_self_eq_norm_sq n
      unfold quadPattern P2 ip
      rw [h_inner, hn]
      norm_num
    have hxEq1 : quadPattern n x = 1 := le_antisymm h_le (by simpa [hnn] using h_ge)
    have hscoreEq1 : score n x = 1 := by
      have : (3/2 : ℝ) * score n x - 1/2 = 1 := by
        rw [← quadPattern_eq_affine_score]
        exact hxEq1
      linarith
    have hip_sq : (ip n x)^2 = 1 := by
      unfold score at hscoreEq1
      exact hscoreEq1
    have hip : ip n x = 1 ∨ ip n x = -1 := by
      have := sq_eq_one_iff.mp hip_sq
      exact this
    cases hip with
    | inl hip1 =>
        have : n = x := (ip_eq_one_iff_eq n x hn hx_unit).1 hip1
        exact Or.inl this.symm
    | inr hipm1 =>
        have hneg : n = -x := (ip_eq_neg_one_iff_eq_neg n x hn hx_unit).1 hipm1
        have : x = -n := by
          have := neg_eq_iff_eq_neg.mpr hneg
          exact this.symm
        exact Or.inr this
  · intro hx
    cases hx with
    | inl hxn =>
        rw [hxn]
        exact n_mem_AxisSet_quadPattern n hn
    | inr hxneg =>
        rw [hxneg]
        exact neg_n_mem_AxisSet_quadPattern n hn

/-! ## Model-to-Data Bridge: Temperature Pattern -/

/--
Temperature-like sky pattern: an affine transform of the quadrupole pattern.

This matches the model-to-data bridge used in the cosmology narrative:
  T(x) = A · P₂(⟨n,x⟩) + B

Here we represent it as:
  T(x) = A * quadPattern n x + B

where:
- A is the quadrupole amplitude (positive for detection)
- B is the monopole offset (isotropic background)
- n is the observer's motion vector

**Physical interpretation**: If the CMB temperature anisotropy (after subtracting
isotropic background) has this functional form, then the extracted axis is
deterministically ±n, independent of the values of A and B (as long as A > 0).
-/
def tempPattern (n : R3) (A B : ℝ) (x : R3) : ℝ :=
  A * quadPattern n x + B

/--
**Bridge Theorem (Model-to-Data, still purely geometric)**

For A > 0, the argmax set on the unit sphere is unchanged by the affine transform,
so the extracted axis is still exactly {n, -n}.

**Physical interpretation**: If the observational fit returns a dominant P₂(⟨n,x⟩)
component with positive amplitude A, the extracted axis is **forced** to be ±n.
There is no freedom for the axis to point elsewhere.

**Connection to README assumptions**: This theorem makes formal the claim that
"IF the CMB has form T(x) = A·P₂(⟨n,x⟩) + B with A > 0, THEN the axis is ±n."
The condition A > 0 is necessary (A < 0 would flip the maximizer to the equator).

**Falsifiability**: If the fitted amplitude A is negative or if the axis extracted
from Planck data differs from the dipole direction, the model is falsified.
-/
theorem AxisSet_tempPattern_eq_pm (n : R3) (hn : IsUnit n) (A B : ℝ) (hA : 0 < A) :
    AxisSet (tempPattern n A B) = {x | x = n ∨ x = -n} := by
  -- Affine invariance of AxisSet
  have h_aff :
      AxisSet (fun x => A * quadPattern n x + B) = AxisSet (quadPattern n) := by
    exact AxisSet_affine (quadPattern n) A B hA
  -- Finish using the proven uniqueness theorem
  calc AxisSet (tempPattern n A B)
      = AxisSet (fun x => A * quadPattern n x + B) := by rfl
    _ = AxisSet (quadPattern n) := h_aff
    _ = {x | x = n ∨ x = -n} := AxisSet_quadPattern_eq_pm n hn

/-! ## Negative Amplitude: Equatorial Maximizers -/

/--
**Definition: Equator of a unit vector n**

The equator perpendicular to n is the set of all unit vectors orthogonal to n.
Equivalently: {x | ⟨n,x⟩ = 0 and ‖x‖ = 1}
-/
def Equator (n : R3) : Set R3 :=
  {x | IsUnit x ∧ ip n x = 0}

/--
**Helper: P₂ value at equator is -1/2**

For any unit vector x orthogonal to n, P₂(⟨n,x⟩) = P₂(0) = -1/2.
This is the minimum value of P₂ on [-1,1].
-/
lemma quadPattern_at_equator (n x : R3) (_hn : IsUnit n) (hx : x ∈ Equator n) :
    quadPattern n x = -1/2 := by
  unfold Equator at hx
  have h_ip_zero : ip n x = 0 := hx.2
  unfold quadPattern P2
  rw [h_ip_zero]
  norm_num

/--
**Helper: P₂ minimum is at the equator**

P₂(t) = (3t² - 1)/2 is minimized at t = 0 with value -1/2.
For unit vectors, this means quadPattern is minimized exactly on the equator.
-/
lemma quadPattern_ge_at_equator (n x : R3) (hn : IsUnit n) (hx : IsUnit x) :
    quadPattern n x ≥ -1/2 := by
  unfold quadPattern P2
  have h_bound : (ip n x)^2 ≤ 1 := by
    have h := abs_real_inner_le_norm n x
    unfold ip IsUnit at *
    rw [hn, hx] at h
    -- |ip n x| ≤ 1, so (ip n x)² ≤ 1
    calc (ip n x)^2
        = |ip n x|^2 := by rw [sq_abs]
      _ ≤ 1^2 := by
        apply sq_le_sq'
        · linarith [abs_nonneg (ip n x)]
        · simpa using h
      _ = 1 := by norm_num
  have : 3 * (ip n x)^2 - 1 ≥ 3 * 0 - 1 := by
    have : (0:ℝ) ≤ 3/2 := by norm_num
    linarith [mul_nonneg this (sq_nonneg (ip n x))]
  linarith

/--
**Helper: If ip ≠ 0, then quadPattern > -1/2 (strict inequality)**

Points off the equator have strictly higher quadPattern values than -1/2.
-/
lemma quadPattern_gt_at_non_equator (n x : R3) (_hn : IsUnit n) (_hx : IsUnit x)
    (h_off : ip n x ≠ 0) :
    quadPattern n x > -1/2 := by
  unfold quadPattern P2
  have : (ip n x)^2 > 0 := by
    exact sq_pos_iff.mpr h_off
  linarith

/--
**Theorem: Equator is non-empty**

For any unit vector n in R³, there exists a unit vector orthogonal to n.

**Proof strategy**: Use Mathlib's finite-dimensional inner product space theory:
1. The orthogonal complement of span{n} has dimension 2 (ℝ³ is 3-dimensional)
2. A subspace of positive dimension contains nonzero vectors
3. Normalize any nonzero orthogonal vector to get a unit vector in the Equator

**Mathlib theorems used**:
- `Submodule.finrank_orthogonal_span_singleton`: dim(span{v}ᗮ) = dim(E) - 1
- `Submodule.exists_mem_ne_zero_of_ne_bot`: positive dimension → nonzero element
-/
theorem equator_nonempty (n : R3) (hn : IsUnit n) : ∃ x, x ∈ Equator n := by
  -- Step 1: n ≠ 0 (from IsUnit n : ‖n‖ = 1)
  have hn_ne : n ≠ 0 := by
    intro h_zero
    unfold IsUnit at hn
    rw [h_zero, norm_zero] at hn
    norm_num at hn

  -- Step 2: The orthogonal complement (ℝ ∙ n)ᗮ is nontrivial
  -- Use Mathlib's finrank_orthogonal_span_singleton theorem
  have h_nontrivial : (Submodule.span ℝ {n})ᗮ ≠ ⊥ := by
    intro h_eq_bot

    -- If (span{n})ᗮ = ⊥, then its dimension is 0
    have h_dim_zero : Module.finrank ℝ (Submodule.span ℝ {n})ᗮ = 0 := by
      rw [h_eq_bot]
      simp

    -- But by Submodule.finrank_orthogonal_span_singleton,
    -- dim((span{v})ᗮ) = dim(E) - 1 for nonzero v in finite-dimensional E
    -- For R3 with dim = 3, we get dim((span{n})ᗮ) = 2

    -- Use the dimension formula for orthogonal complements
    have h_dim_two : Module.finrank ℝ (Submodule.span ℝ {n})ᗮ = 2 := by
      -- Apply Submodule.finrank_orthogonal_span_singleton
      -- This theorem states: for E with finrank = n+1, span{v}ᗮ has finrank = n
      -- For R3 with finrank = 3 = 2+1, we get finrank (span{n})ᗮ = 2

      -- Need to provide Fact instance that finrank R3 = 3
      haveI : Fact (Module.finrank ℝ R3 = 2 + 1) := by
        constructor
        -- R3 = PiLp 2 (Fin 3 → ℝ) has dimension 3
        -- Try to let Lean infer this from PiLp instances
        unfold R3
        norm_num

      -- Now apply the theorem
      -- finrank_orthogonal_span_singleton says: finrank (𝕜 ∙ v)ᗮ = n
      -- where 𝕜 ∙ v is notation for span 𝕜 {v}
      exact Submodule.finrank_orthogonal_span_singleton hn_ne

    -- Now we have: 0 = 2, contradiction
    omega

  -- Step 3: Get a nonzero element from the nontrivial orthogonal complement
  obtain ⟨v, hv_mem, hv_ne⟩ := Submodule.exists_mem_ne_zero_of_ne_bot h_nontrivial

  -- Step 4: v is orthogonal to n (from hv_mem : v ∈ (span ℝ {n})ᗮ)
  have hv_orth : inner ℝ n v = 0 := by
    rw [Submodule.mem_orthogonal] at hv_mem
    exact hv_mem n (Submodule.mem_span_singleton_self n)

  -- Step 5: Normalize v to get unit vector
  let v_norm := ‖v‖
  have hv_norm_pos : 0 < v_norm := norm_pos_iff.mpr hv_ne
  let x := v_norm⁻¹ • v

  use x
  constructor
  · -- Prove IsUnit x: ‖x‖ = 1
    unfold IsUnit
    rw [norm_smul, norm_inv, norm_norm]
    field_simp [ne_of_gt hv_norm_pos]
  · -- Prove ip n x = 0
    unfold ip
    rw [inner_smul_right, hv_orth, mul_zero]

/--
**Negative-Amplitude Companion Theorem**

When the amplitude A < 0, the maximizers of T(x) = A·P₂(⟨n,x⟩) + B are
**not** at the poles ±n, but instead at the **equator** (orthogonal to n).

**Physical interpretation**: This proves that the sign of A matters geometrically.
- A > 0: poles aligned with motion vector (bridge theorem)
- A < 0: equator orthogonal to motion vector (this theorem)

**Falsifiability**: If Planck data require A < 0 to fit observations, the
axis-alignment prediction is falsified. The sign is not a free parameter.

**Why this matters**: Reviewers might ask "couldn't you just flip the sign?"
This theorem shows the answer is NO - flipping the sign changes the geometric
prediction from poles to equator, which is observationally distinguishable.
-/
theorem AxisSet_tempPattern_eq_equator (n : R3) (hn : IsUnit n) (A B : ℝ) (hA : A < 0) :
    AxisSet (tempPattern n A B) = Equator n := by
  ext x
  unfold AxisSet tempPattern Equator
  constructor
  · intro ⟨hx_unit, hx_max⟩
    constructor
    · exact hx_unit
    · -- Must show ip n x = 0
      -- Proof by contradiction: if ip n x ≠ 0, x is not a maximizer

      by_contra h_not_zero

      -- If ip n x ≠ 0, then quadPattern n x > -1/2
      have h_quad_gt : quadPattern n x > -1/2 :=
        quadPattern_gt_at_non_equator n x hn hx_unit h_not_zero

      -- Get a point on the equator (exists by axiom)
      obtain ⟨e, he⟩ := equator_nonempty n hn
      have he_unit : IsUnit e := he.1
      have he_ip : ip n e = 0 := he.2

      -- At the equator, quadPattern = -1/2
      have h_eq_val : quadPattern n e = -1/2 := quadPattern_at_equator n e hn he

      -- Since A < 0 and quadPattern n x > quadPattern n e:
      -- A * quadPattern n x < A * quadPattern n e
      have : A * quadPattern n x < A * quadPattern n e := by
        rw [h_eq_val]
        exact mul_lt_mul_of_neg_left h_quad_gt hA

      -- Therefore tempPattern n A B x < tempPattern n A B e
      have : tempPattern n A B x < tempPattern n A B e := by
        unfold tempPattern
        linarith

      -- But this contradicts x being a maximizer
      have : tempPattern n A B e ≤ tempPattern n A B x :=
        hx_max e he_unit
      linarith
  · intro ⟨hx_unit, h_ip_zero⟩
    constructor
    · exact hx_unit
    · intro y hy
      -- Show tempPattern n A B y ≤ tempPattern n A B x when x is on equator
      -- Since A < 0, we want to show:
      -- A * quadPattern n y + B ≤ A * quadPattern n x + B
      -- which means: A * quadPattern n y ≤ A * quadPattern n x
      -- Since A < 0, this means: quadPattern n y ≥ quadPattern n x

      have h_x_min : quadPattern n x = -1/2 := by
        unfold quadPattern P2
        rw [h_ip_zero]
        norm_num

      have h_y_ge : quadPattern n y ≥ -1/2 :=
        quadPattern_ge_at_equator n y hn hy

      calc tempPattern n A B y
          = A * quadPattern n y + B := rfl
        _ ≤ A * (-1/2) + B := by
            have : A * quadPattern n y ≤ A * (-1/2) := by
              have hA_neg : A ≤ 0 := le_of_lt hA
              exact mul_le_mul_of_nonpos_left h_y_ge hA_neg
            linarith
        _ = A * quadPattern n x + B := by rw [h_x_min]
        _ = tempPattern n A B x := rfl

end QFD.Cosmology