import RNLAPhaseB.B3Base
import RNLAPhaseB.Complexity

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

namespace RNLAPhaseB.MT2020.RandHBS

open RNLAPhaseB.Complexity

noncomputable section

/-- Perfect binary cluster tree of Sec. 20.4: node (l, j) with l ≤ depth and j < 2^l
carries the index vector `box l j`. The root holds the full index set. Each parent is
the disjoint union of its two children. The leaves sit at level `depth`.
The source stops splitting once an index vector is short, and calls the meaning of
short application-dependent. That rule is not formalized here. Instead, the HBS
convention's stated consequences (O(k) indices per leaf, O(n/k) nodes) are imposed as
hypotheses of the theorem. -/
structure ClusterTree (n : ℕ) where
  depth : ℕ
  box : ℕ → ℕ → Finset (Fin n)
  root : box 0 0 = Finset.univ
  split : ∀ l, l < depth → ∀ j, j < 2 ^ l →
    box l j = box (l + 1) (2 * j) ∪ box (l + 1) (2 * j + 1) ∧
    Disjoint (box (l + 1) (2 * j)) (box (l + 1) (2 * j + 1))

/-- Sec. 20.6 convention: each leaf index vector holds O(k) indices, i.e. at most `b`. -/
def ClusterTree.LeafSizeAtMost {n : ℕ} (T : ClusterTree n) (b : ℕ) : Prop :=
  ∀ j, j < 2 ^ T.depth → (T.box T.depth j).card ≤ b

/-- Number of nodes of the perfect binary tree: levels 0..depth, with 2^l nodes at level l. -/
def ClusterTree.numNodes {n : ℕ} (T : ClusterTree n) : ℕ :=
  2 ^ (T.depth + 1) - 1

/-- Data of an HBS representation of rank k (Sec. 20.4-20.6). It consists of:
* the leaf diagonal blocks D_tau;
* the leaf (long) bases U_tau and V_tau, with rows indexed by the box;
* the short 2k x k nested bases U_tau and V_tau of the parents;
* the k x k sibling matrices Ktilde_{alpha,beta} (false) and Ktilde_{beta,alpha} (true). -/
structure HBSRep (n k : ℕ) where
  diag : ℕ → Matrix (Fin n) (Fin n) ℂ
  leafU : ℕ → Matrix (Fin n) (Fin k) ℂ
  leafV : ℕ → Matrix (Fin n) (Fin k) ℂ
  shortU : ℕ → ℕ → Matrix (Fin k ⊕ Fin k) (Fin k) ℂ
  shortV : ℕ → ℕ → Matrix (Fin k ⊕ Fin k) (Fin k) ℂ
  sib : ℕ → ℕ → Bool → Matrix (Fin k) (Fin k) ℂ

/-- Long basis of the node at distance m above the leaves, built by the nested-basis
relation (eq:nestedbasis): U_tau^long = [U_alpha^long 0; 0 U_beta^long] U_tau.
Rows outside the node's box are zero. -/
def nestedLong {n k : ℕ} (depth : ℕ) (box : ℕ → ℕ → Finset (Fin n))
    (leaf : ℕ → Matrix (Fin n) (Fin k) ℂ)
    (short : ℕ → ℕ → Matrix (Fin k ⊕ Fin k) (Fin k) ℂ) :
    ℕ → ℕ → Matrix (Fin n) (Fin k) ℂ
  | 0, j => fun i c => if i ∈ box depth j then leaf j i c else 0
  | m + 1, j => fun i c =>
      if i ∈ box (depth - m) (2 * j) then
        ∑ c', nestedLong depth box leaf short m (2 * j) i c' *
          short (depth - (m + 1)) j (Sum.inl c') c
      else if i ∈ box (depth - m) (2 * j + 1) then
        ∑ c', nestedLong depth box leaf short m (2 * j + 1) i c' *
          short (depth - (m + 1)) j (Sum.inr c') c
      else 0

def HBSRep.longU {n k : ℕ} (R : HBSRep n k) (T : ClusterTree n) (l j : ℕ) :
    Matrix (Fin n) (Fin k) ℂ :=
  nestedLong T.depth T.box R.leafU R.shortU (T.depth - l) j

def HBSRep.longV {n k : ℕ} (R : HBSRep n k) (T : ClusterTree n) (l j : ℕ) :
    Matrix (Fin n) (Fin k) ℂ :=
  nestedLong T.depth T.box R.leafV R.shortV (T.depth - l) j

/-- `R` represents `K` exactly on the tessellation (eq:siblingboxes, eq:diagbox).
The leaf diagonal blocks equal D_tau. Every sibling block K(I_alpha, I_beta) equals
U_alpha^long Ktilde_{alpha,beta} (V_beta^long)^*, and symmetrically for (beta, alpha). -/
def HBSRep.Represents {n k : ℕ} (R : HBSRep n k) (T : ClusterTree n)
    (K : Matrix (Fin n) (Fin n) ℂ) : Prop :=
  (∀ j, j < 2 ^ T.depth → ∀ i ∈ T.box T.depth j, ∀ i' ∈ T.box T.depth j,
      K i i' = R.diag j i i') ∧
  (∀ l, l < T.depth → ∀ j, j < 2 ^ l →
    (∀ i ∈ T.box (l + 1) (2 * j), ∀ i' ∈ T.box (l + 1) (2 * j + 1),
      K i i' = (R.longU T (l + 1) (2 * j) * R.sib l j false *
        (R.longV T (l + 1) (2 * j + 1))ᴴ) i i') ∧
    (∀ i ∈ T.box (l + 1) (2 * j + 1), ∀ i' ∈ T.box (l + 1) (2 * j),
      K i i' = (R.longU T (l + 1) (2 * j + 1) * R.sib l j true *
        (R.longV T (l + 1) (2 * j))ᴴ) i i'))

/-- K is compressible in the HBS format of Sec. 20.6 with rank k on tree T. -/
def HBSCompressible {n : ℕ} (T : ClusterTree n) (k : ℕ)
    (K : Matrix (Fin n) (Fin n) ℂ) : Prop :=
  ∃ R : HBSRep n k, R.Represents T K

def cpx (z : ℂ) : List ℝ := [z.re, z.im]

def boxList {n : ℕ} (T : ClusterTree n) (l j : ℕ) : List (Fin n) :=
  (List.finRange n).filter (fun i => decide (i ∈ T.box l j))

def twoK (k : ℕ) : List (Fin k ⊕ Fin k) :=
  (List.finRange k).map Sum.inl ++ (List.finRange k).map Sum.inr

/-- Explicit emitted HBS data in level order. Each node contributes the following:
* a leaf: D_tau over box x box, and, when below the root, the leaf U and V over box x k;
* a parent below the root: its short 2k x k U and V;
* a parent: its two k x k sibling matrices.
Complex scalars are emitted as (re, im). -/
def encodeRep {n k : ℕ} (T : ClusterTree n) (R : HBSRep n k) : List ℝ :=
  (List.range (T.depth + 1)).flatMap fun l => (List.range (2 ^ l)).flatMap fun j =>
    if l = T.depth then
      (boxList T l j).flatMap (fun i => (boxList T l j).flatMap fun i' => cpx (R.diag j i i')) ++
      (if 1 ≤ l then
        (boxList T l j).flatMap (fun i => (List.finRange k).flatMap fun c => cpx (R.leafU j i c)) ++
        (boxList T l j).flatMap (fun i => (List.finRange k).flatMap fun c => cpx (R.leafV j i c))
       else [])
    else
      (if 1 ≤ l then
        (twoK k).flatMap (fun a => (List.finRange k).flatMap fun c => cpx (R.shortU l j a c)) ++
        (twoK k).flatMap (fun a => (List.finRange k).flatMap fun c => cpx (R.shortV l j a c))
       else []) ++
      (List.finRange k).flatMap (fun a => (List.finRange k).flatMap fun c => cpx (R.sib l j false a c)) ++
      (List.finRange k).flatMap (fun a => (List.finRange k).flatMap fun c => cpx (R.sib l j true a c))

/-- The output is a full (exact) HBS representation of K on T. -/
def IsFullHBSOutput {n : ℕ} (T : ClusterTree n) (k : ℕ) (K : Matrix (Fin n) (Fin n) ℂ)
    (out : OutputEncoding.Buffer) : Prop :=
  ∃ R : HBSRep n k, R.Represents T K ∧ out = ⟨[], encodeRep T R⟩

def gaussC {n k : ℕ} (G : Matrix (Fin n) (Fin k) ℝ) : Matrix (Fin n) (Fin k) ℂ :=
  G.map (fun x => (x : ℂ))

/-- Supplied input. It contains the tree, the Gaussian G, Y = K G and Z = K^* G. -/
def encodeInput {n k : ℕ} (T : ClusterTree n) (G : Matrix (Fin n) (Fin k) ℝ)
    (Y Z : Matrix (Fin n) (Fin k) ℂ) : RealRAM.Input :=
  ⟨[n, k, T.depth] ++ (List.range (T.depth + 1)).flatMap (fun l =>
      (List.range (2 ^ l)).flatMap fun j =>
        (boxList T l j).length :: (boxList T l j).map Fin.val),
   (List.finRange n).flatMap (fun i => (List.finRange k).map fun c => G i c) ++
   (List.finRange n).flatMap (fun i => (List.finRange k).flatMap fun c => cpx (Y i c)) ++
   (List.finRange n).flatMap (fun i => (List.finRange k).flatMap fun c => cpx (Z i c))⟩

/-- The input for the draw G: the tree, G, Y = K G and Z = K^* G (eq:randomhudsonsample1-2). -/
def sampleInput {n k : ℕ} (T : ClusterTree n) (K : Matrix (Fin n) (Fin n) ℂ)
    (G : Matrix (Fin n) (Fin k) ℝ) : RealRAM.Input :=
  encodeInput T G (K * gaussC G) (Kᴴ * gaussC G)

/-- Entry access: request 0 with indices [i, j] returns (re, im) of K(i, j). -/
def entryOracle {n : ℕ} (K : Matrix (Fin n) (Fin n) ℂ) : Oracle.Interface where
  answer := fun request input =>
    if request = 0 then
      match input.indices with
      | [i, j] =>
          if hi : i < n then
            if hj : j < n then some ⟨[], cpx (K ⟨i, hi⟩ ⟨j, hj⟩)⟩ else none
          else none
      | _ => none
    else none
  work := fun _ _ => 1
  setupCost := 0

/-- Placeholder sampling interface. It is unreachable because admissible programs
contain no `sample` instruction; the only randomness is the supplied G. -/
def unusedSampling : Sampling.Model where
  law := fun _ => Measure.dirac 0
  probability := fun _ => inferInstance
  setupCost := 0
  drawCost := fun _ => 1
  positiveCost := fun _ => Nat.one_pos

def entryAccess {n : ℕ} (K : Matrix (Fin n) (Fin n) ℂ) : ChargedRAM.Access :=
  ⟨unusedSampling, entryOracle K, .unboundedUnit⟩

/-- Programs have no sampling, no uncharged subroutine calls and no random bits.
All computation therefore occurs in charged base steps. -/
def Admissible (p : ChargedRAM.Program) : Prop :=
  ∀ ins ∈ p.code, match ins with
    | .sample _ _ => False
    | .call _ _ _ => False
    | .base (.randomBit _) => False
    | _ => True

/-- One entry evaluation per oracle query. -/
def entryCharge (p : ChargedRAM.Program) (s : ChargedRAM.State) : ℕ :=
  match p.code[s.ram.pc]? with
  | some (.query _ _ _ _) => 1
  | _ => 0

/-- One floating point operation per scalar arithmetic, square root or scalar comparison. -/
def flopCharge (p : ChargedRAM.Program) (s : ChargedRAM.State) : ℕ :=
  match p.code[s.ram.pc]? with
  | some (.base (.scalarBinary _ _ _ _)) => 1
  | some (.base (.squareRoot _ _)) => 1
  | some (.base (.branchScalar _ _ _ _)) => 1
  | _ => 0

/-- An actual ChargedRAM execution, counting entry evaluations and flops along its steps. -/
inductive CountedRun (p : ChargedRAM.Program) (a : ChargedRAM.Access) (input : RealRAM.Input) :
    ChargedRAM.State → ℕ → ℕ → ChargedRAM.State → Prop
  | done (s : ChargedRAM.State) : CountedRun p a input s 0 0 s
  | next {s u t : ChargedRAM.State} {c e f : ℕ} :
      ChargedRAM.Step p a input (fun _ => false) (fun _ => 0) s c u →
      CountedRun p a input u e f t →
      CountedRun p a input s (entryCharge p s + e) (flopCharge p s + f) t

/-- mt2020 Theorem thm:randHBS.

One fixed program receives the tree, the Gaussian G, Y = K G and Z = K^* G. Producing
these samples is what costs T_apply. The program has entry access to K at T_entry per
entry and charges T_flop per floating point operation.

The tree domain follows the Sec. 20.6 HBS convention: each leaf index vector holds O(k)
indices, which means there are overall O(n/k) nodes. Both are imposed with one constant c:
leaf size at most c k, and k times the number of nodes at most c n.

The claim has two parts:
* for every draw, the run halts with total cost at most
  T_apply + T_entry * O(nk) + T_flop * O(nk^2);
* almost surely in G, the output of that halted run is a full HBS representation of K.
The O-constants depend only on c. -/
theorem randHBS :
    ∃ prog : ChargedRAM.Program, Admissible prog ∧
    ∀ c : ℕ, ∃ C : ℝ, 0 < C ∧ ∃ N : ℕ,
    ∀ (n k : ℕ) (T : ClusterTree n), T.LeafSizeAtMost (c * k) →
      k * T.numNodes ≤ c * n →
    ∀ K : Matrix (Fin n) (Fin n) ℂ, HBSCompressible T k K →
    ∀ (Ω : Type) [MeasurableSpace Ω] (P : Measure Ω) (G : Ω → Matrix (Fin n) (Fin k) ℝ),
      RNLAPhaseB.StandardGaussianTestMatrix P G →
      (∀ ω, ∃ (final : ChargedRAM.State) (e f : ℕ),
        CountedRun prog (entryAccess K) (sampleInput T K (G ω)) {} e f final ∧
        prog.halted final ∧
        (N ≤ n → ∀ Tapply Tentry Tflop : ℝ, 0 ≤ Tapply → 0 ≤ Tentry → 0 ≤ Tflop →
          Tapply + Tentry * (e : ℝ) + Tflop * (f : ℝ) ≤
            Tapply + Tentry * (C * ((n : ℝ) * k)) + Tflop * (C * ((n : ℝ) * (k : ℝ) ^ 2)))) ∧
      (∀ᵐ ω ∂P, ∀ (final : ChargedRAM.State) (e f : ℕ),
        CountedRun prog (entryAccess K) (sampleInput T K (G ω)) {} e f final →
        prog.halted final → IsFullHBSOutput T k K final.output) := by
  sorry

end

end RNLAPhaseB.MT2020.RandHBS