import RNLAPhaseB.B3Base
import RNLAPhaseB.Complexity

open MeasureTheory
open scoped Matrix
open scoped Matrix.Norms.L2Operator
open scoped BigOperators
open scoped ENNReal

namespace RNLAPhaseB.MT2020RrfSrtt

/-- Unitary DFT, entries n^{-1/2} exp(-2πi jk/n) (complex field only). -/
noncomputable def dftMatrix (𝕜 : Type*) [RCLike 𝕜] (n : ℕ) : Matrix (Fin n) (Fin n) 𝕜 :=
  Matrix.of fun j k =>
    ((Real.cos (2 * Real.pi * ((j : ℕ) : ℝ) * ((k : ℕ) : ℝ) / n) : 𝕜)
      - RCLike.I * (Real.sin (2 * Real.pi * ((j : ℕ) : ℝ) * ((k : ℕ) : ℝ) / n) : 𝕜))
      / (Real.sqrt n : 𝕜)

/-- Orthonormal DCT-II. -/
noncomputable def dct2Matrix (𝕜 : Type*) [RCLike 𝕜] (n : ℕ) : Matrix (Fin n) (Fin n) 𝕜 :=
  Matrix.of fun j k =>
    ((Real.sqrt (2 / n) * (if (j : ℕ) = 0 then 1 / Real.sqrt 2 else 1)
      * Real.cos (Real.pi * (2 * ((k : ℕ) : ℝ) + 1) * ((j : ℕ) : ℝ) / (2 * n)) : ℝ) : 𝕜)

/-- Unitary discrete Hartley transform, entries n^{-1/2} cas(2π jk/n). -/
noncomputable def dhtMatrix (𝕜 : Type*) [RCLike 𝕜] (n : ℕ) : Matrix (Fin n) (Fin n) 𝕜 :=
  Matrix.of fun j k =>
    (((Real.cos (2 * Real.pi * ((j : ℕ) : ℝ) * ((k : ℕ) : ℝ) / n)
      + Real.sin (2 * Real.pi * ((j : ℕ) : ℝ) * ((k : ℕ) : ℝ) / n)) / Real.sqrt n : ℝ) : 𝕜)

/-- Unitary Walsh–Hadamard transform for n = 2^p: entries (-1)^{⟨j,k⟩}/√n,
with ⟨j,k⟩ the number of common one-bits among the p binary digits. -/
noncomputable def whtMatrix (𝕜 : Type*) [RCLike 𝕜] (n p : ℕ) : Matrix (Fin n) (Fin n) 𝕜 :=
  Matrix.of fun j k =>
    ((((-1 : ℝ) ^ ((Finset.range p).filter
        (fun z => ((j : ℕ).testBit z && (k : ℕ).testBit z) = true)).card) / Real.sqrt n : ℝ) : 𝕜)

/-- The source's unitary trigonometric transforms F (DFT in the complex case;
DCT2, DHT; WHT when n is a power of two). -/
def IsUnitaryTrigTransform {𝕜 : Type*} [RCLike 𝕜] {n : ℕ}
    (F : Matrix (Fin n) (Fin n) 𝕜) : Prop :=
  ((RCLike.I : 𝕜) ≠ 0 ∧ F = dftMatrix 𝕜 n) ∨ F = dct2Matrix 𝕜 n ∨ F = dhtMatrix 𝕜 n ∨
    ∃ p : ℕ, n = 2 ^ p ∧ F = whtMatrix 𝕜 n p

/-- Random restriction R ∈ F^{ℓ×n}: for ℓ ≤ n row i selects coordinate ι i
(ι a uniformly random injection, see `SRTTLaw`); for ℓ > n, the source's
extension R_{ij} = 1 iff i = j. -/
def restrictionMatrix {𝕜 : Type*} [RCLike 𝕜] {ℓ n : ℕ} (ι : Fin ℓ → Fin n) :
    Matrix (Fin ℓ) (Fin n) 𝕜 :=
  Matrix.of fun i j =>
    if ℓ ≤ n then (if j = ι i then 1 else 0) else (if (i : ℕ) = (j : ℕ) then 1 else 0)

/-- Permutation matrix of σ. -/
def permutationMatrix {𝕜 : Type*} [RCLike 𝕜] {n : ℕ} (σ : Equiv.Perm (Fin n)) :
    Matrix (Fin n) (Fin n) 𝕜 :=
  Matrix.of fun i j => if σ i = j then 1 else 0

/-- SRTT embedding S = sqrt(n/ℓ) R F E Π ∈ F^{ℓ×n}. -/
noncomputable def srttEmbedding {𝕜 : Type*} [RCLike 𝕜] {ℓ n : ℕ}
    (F : Matrix (Fin n) (Fin n) 𝕜) (ι : Fin ℓ → Fin n) (σ : Equiv.Perm (Fin n))
    (e : Fin n → 𝕜) : Matrix (Fin ℓ) (Fin n) 𝕜 :=
  let R : Matrix (Fin ℓ) (Fin n) 𝕜 := restrictionMatrix (𝕜 := 𝕜) ι
  let RF : Matrix (Fin ℓ) (Fin n) 𝕜 := R * F
  let RFE : Matrix (Fin ℓ) (Fin n) 𝕜 := RF * Matrix.diagonal e
  let RFEP : Matrix (Fin ℓ) (Fin n) 𝕜 := RFE * permutationMatrix (𝕜 := 𝕜) σ
  (Real.sqrt ((n : ℝ) / ℓ) : 𝕜) • RFEP

/-- SRTT test matrix Ω = S^* ∈ F^{n×ℓ}. -/
noncomputable def srttTestMatrix {𝕜 : Type*} [RCLike 𝕜] {ℓ n : ℕ}
    (F : Matrix (Fin n) (Fin n) 𝕜) (ι : Fin ℓ → Fin n) (σ : Equiv.Perm (Fin n))
    (e : Fin n → 𝕜) : Matrix (Fin n) (Fin ℓ) 𝕜 :=
  (srttEmbedding F ι σ e)ᴴ

open Classical in
/-- Law of the SRTT randomness: E is a random sign flip (iid entries uniform on
{z ∈ F : |z| = 1}, expressed as unit modulus plus invariance under every unit
multiplier), Π a uniform permutation, the restriction a uniform injection
(when ℓ ≤ n), and the three factors independent. -/
structure SRTTLaw {𝕜 : Type*} [RCLike 𝕜] [MeasurableSpace 𝕜] {Ω : Type*} [MeasurableSpace Ω]
    {ℓ n : ℕ} (P : Measure Ω) (ι : Ω → Fin ℓ → Fin n) (π : Ω → Equiv.Perm (Fin n))
    (e : Ω → Fin n → 𝕜) : Prop where
  isProbabilityMeasure : IsProbabilityMeasure P
  sign_measurable : ∀ i, Measurable (fun ω => e ω i)
  sign_unit : ∀ i, ∀ᵐ ω ∂P, ‖e ω i‖ = 1
  sign_uniform : ∀ i (u : 𝕜), ‖u‖ = 1 →
    P.map (fun ω => u * e ω i) = P.map (fun ω => e ω i)
  sign_iIndep : ProbabilityTheory.iIndepFun (fun i ω => e ω i) P
  joint : ∀ (f : Fin ℓ → Fin n) (σ : Equiv.Perm (Fin n)) (A : Set (Fin n → 𝕜)),
    MeasurableSet A →
    P {ω | (ℓ ≤ n → ι ω = f) ∧ π ω = σ ∧ e ω ∈ A} =
      (if ℓ ≤ n then (if Function.Injective f then ((n.descFactorial ℓ : ℝ≥0∞))⁻¹ else 0)
        else 1) * (((n.factorial : ℝ≥0∞))⁻¹ * P {ω | e ω ∈ A})

/-- σ_j(B) (1-based): the j-th largest square root of the eigenvalues of B^*B,
and 0 for j > n (sSup ∅ = 0). -/
noncomputable def singularValue {𝕜 : Type*} [RCLike 𝕜] {m n : ℕ}
    (B : Matrix (Fin m) (Fin n) 𝕜) (j : ℕ) : ℝ :=
  sSup {t : ℝ | j ≤ (Finset.univ.filter (fun i : Fin n =>
    t ≤ Real.sqrt ((Matrix.isHermitian_conjTranspose_mul_self B).eigenvalues i))).card}

/-- mt2020 thm:rrf-srtt. For B ∈ F^{m×n}, natural k, SRTT test matrix
Ω ∈ F^{n×ℓ} with ℓ ≥ 8(k + 8 log(kn)) log k and Y = BΩ:
‖(I - P_Y)B‖ ≤ (1 + 3√(n/ℓ)) σ_{k+1}, failure probability O(k^{-1}) with one
constant C uniform over all inputs (k = 0 makes the bound ⊤). log = natural log. -/
theorem rrf_srtt {𝕜 : Type*} [RCLike 𝕜] [MeasurableSpace 𝕜] [BorelSpace 𝕜] :
    ∃ C : ℝ, 0 < C ∧
      ∀ (m n k ℓ : ℕ) (B : Matrix (Fin m) (Fin n) 𝕜) (F : Matrix (Fin n) (Fin n) 𝕜),
        IsUnitaryTrigTransform F →
        (ℓ : ℝ) ≥ 8 * ((k : ℝ) + 8 * Real.log ((k : ℝ) * (n : ℝ))) * Real.log (k : ℝ) →
        ∀ {Ω : Type} [MeasurableSpace Ω] (P : Measure Ω) (ι : Ω → Fin ℓ → Fin n)
          (π : Ω → Equiv.Perm (Fin n)) (e : Ω → Fin n → 𝕜),
          SRTTLaw P ι π e →
          P {ω | (1 + 3 * Real.sqrt ((n : ℝ) / ℓ)) * singularValue B (k + 1) <
              ‖(1 - orthogonalProjector (B * srttTestMatrix F (ι ω) (π ω) (e ω))) * B‖}
            ≤ ENNReal.ofReal C * ((k : ℝ≥0∞))⁻¹ := by
  sorry

end RNLAPhaseB.MT2020RrfSrtt