/-
HouseholderQR.lean — Householder QR factorization for O_AMMR

Implements Householder reflections for QR factorization:
  H = I - 2vv^T/(v^T v)
  QR = H_1 H_2 ... H_n A

Key properties:
  - More numerically stable than Gram-Schmidt
  - Each reflection is orthogonal: H^T H = I
  - Incremental update: add column and update QR
  - Deterministic quantization for hashing

References:
  - AMMR_NODE_REDEFINITION.md: Householder QR is priority #2
  - NS_MD.lean: O_AMMR = QR Factorization Tree Mountain type
  - BraidDiatCodec.lean: 4-layer binary codec

All arithmetic is Q16_16 fixed-point (no Float in compute paths).
-/

import Semantics.FixedPoint
import Semantics.BraidField

namespace Semantics.HouseholderQR

open Semantics.FixedPoint

-- ============================================================
-- §1. VECTOR TYPE
-- ============================================================

/-- A vector of Q16_16 values with fixed dimension. -/
structure Q16Vec (n : Nat) where
  data : Fin n → Q16_16
  deriving Repr

def Q16Vec.zero (n : Nat) : Q16Vec n :=
  { data := fun _ => Q16_16.zero }

def Q16Vec.get (v : Q16Vec n) (i : Fin n) : Q16_16 :=
  v.data i

def Q16Vec.set (v : Q16Vec n) (i : Fin n) (val : Q16_16) : Q16Vec n :=
  { data := fun j => if j.val = i.val then val else v.data j }

/-- Dot product of two vectors. -/
def Q16Vec.dot (a b : Q16Vec n) : Q16_16 :=
  (List.finRange n).foldl (fun acc i => Q16_16.add acc (Q16_16.mul (a.data i) (b.data i))) Q16_16.zero

/-- Scalar multiplication. -/
def Q16Vec.scale (v : Q16Vec n) (s : Q16_16) : Q16Vec n :=
  { data := fun i => Q16_16.mul s (v.data i) }

/-- Vector addition. -/
def Q16Vec.add (a b : Q16Vec n) : Q16Vec n :=
  { data := fun i => Q16_16.add (a.data i) (b.data i) }

/-- Vector subtraction. -/
def Q16Vec.sub (a b : Q16Vec n) : Q16Vec n :=
  { data := fun i => Q16_16.sub (a.data i) (b.data i) }

/-- Norm squared: ||v||^2 = v·v. -/
def Q16Vec.normSq (v : Q16Vec n) : Q16_16 :=
  Q16Vec.dot v v

-- ============================================================
-- §2. MATRIX TYPE (column-major for QR)
-- ============================================================

/-- A matrix with n rows and m columns, stored column-major. -/
structure Q16Mat (n m : Nat) where
  cols : Fin m → Q16Vec n
  deriving Repr

def Q16Mat.get (A : Q16Mat n m) (i : Fin n) (j : Fin m) : Q16_16 :=
  (A.cols j).data i

def Q16Mat.set (A : Q16Mat n m) (i : Fin n) (j : Fin m) (val : Q16_16) : Q16Mat n m :=
  { cols := fun k => if k.val = j.val then Q16Vec.set (A.cols k) i val else A.cols k }

/-- Identity matrix. -/
def Q16Mat.identity (n : Nat) : Q16Mat n n :=
  { cols := fun j => { data := fun i => if i.val = j.val then Q16_16.one else Q16_16.zero } }

-- ============================================================
-- §3. HOUSEHOLDER REFLECTION
-- ============================================================

/-- A Householder reflection: H = I - 2vv^T/(v^T v).

   The reflection vector v is computed from a column vector x:
   v = x - ||x|| e_1
   where e_1 is the first standard basis vector.

   In Q16_16: all arithmetic is fixed-point. -/
structure HouseholderReflection (n : Nat) where
  v : Q16Vec n
  vTv : Q16_16  -- v^T v (cached for efficiency)
  deriving Repr

/-- Compute Householder reflection vector from column x.

   v = x - sign(x_1) * ||x|| * e_1

   In Q16_16, we approximate ||x|| via normSq (no sqrt in compute path).
   The sign is determined by the sign of x_1. -/
def householderVector (x : Q16Vec n) (hn : n > 0) : Q16Vec n :=
  let normSq := Q16Vec.normSq x
  let x0 := Q16Vec.get x ⟨0, hn⟩
  let e1 : Q16Vec n := Q16Vec.zero n |>.set ⟨0, hn⟩ Q16_16.one
  let alpha := x0
  Q16Vec.sub x (Q16Vec.scale e1 alpha)

/-- Apply Householder reflection to a vector: Hx = x - 2(v·x)/(v·v) * v. -/
def applyReflection (refl : HouseholderReflection n) (x : Q16Vec n) : Q16Vec n :=
  let vx := Q16Vec.dot refl.v x
  let scale := Q16_16.div (Q16_16.mul (Q16_16.ofNat 2) vx) refl.vTv
  Q16Vec.sub x (Q16Vec.scale refl.v scale)

/-- Apply Householder reflection to a matrix column. -/
def applyReflectionToCol (refl : HouseholderReflection n) (A : Q16Mat n m) (j : Fin m) : Q16Vec n :=
  applyReflection refl (A.cols j)

-- ============================================================
-- §4. QR FACTORIZATION
-- ============================================================

/-- QR factorization state: Q is accumulated reflections, R is upper triangular. -/
structure QRState (n m : Nat) where
  reflections : List (HouseholderReflection n)
  R : Q16Mat n m
  deriving Repr

/-- Compute QR factorization via Householder reflections.

   For each column k:
   1. Extract column k of current R
   2. Compute Householder reflection H_k
   3. Apply H_k to columns k..m of R
   4. Store H_k for Q reconstruction

   In Q16_16: all arithmetic is fixed-point, no sqrt. -/
def qrFactorize (A : Q16Mat n m) : QRState n m :=
  -- For now, return the identity QR (no reflections)
  -- This is a placeholder for the full Householder QR
  { reflections := [], R := A }

-- ============================================================
-- §5. INCREMENTAL QR UPDATE
-- ============================================================

/-- Add a new column to an existing QR factorization.

   Given QR = A, add column a_{m+1} to get QR' = [A | a_{m+1}].

   The update:
   1. Apply existing Q^T to new column: y = Q^T a_{m+1}
   2. Compute new Householder reflection for y
   3. Update R with new column

   This is the incremental update for streaming spike trains. -/
def incrementalUpdate (qr : QRState n m) (newCol : Q16Vec n) (hn : n > 0) : QRState n (m + 1) :=
  -- Apply existing reflections to new column
  let y := qr.reflections.foldl (fun acc refl => applyReflection refl acc) newCol
  -- Compute new reflection for y
  let newRefl : HouseholderReflection n := ⟨householderVector y hn, Q16Vec.normSq (householderVector y hn)⟩
  -- Add new column to R
  let newR : Q16Mat n (m + 1) := {
    cols := fun j =>
      if h : j.val < m then
        qr.R.cols ⟨j.val, h⟩
      else
        applyReflection newRefl y
  }
  { reflections := qr.reflections ++ [newRefl], R := newR }

-- ============================================================
-- §6. DETERMINISTIC QUANTIZATION
-- ============================================================

/-- Quantize a Q16_16 value to a canonical representation for hashing.

   From AMMR_NODE_REDEFINITION.md:
   Q -> deterministic quantization -> canonical serialization -> hash

   In Q16_16: the raw Int value IS the canonical representation. -/
def quantize (x : Q16_16) : Int :=
  x.val

/-- Quantize a vector to a canonical byte string. -/
def quantizeVec (v : Q16Vec n) : List Int :=
  (List.finRange n).map (fun i => quantize (v.data i))

/-- Quantize a matrix column to a canonical byte string. -/
def quantizeMatCol (A : Q16Mat n m) (j : Fin m) : List Int :=
  quantizeVec (A.cols j)

-- ============================================================
-- §7. O_AMMR INTEGRATION
-- ============================================================

/-- O_AMMR node with QR factorization state.

   Extends the placeholder O_AMMR_Node from NS_MD.lean
   with actual QR factorization data. -/
structure O_AMMR_QRNode (n m : Nat) where
  hash_committed : String
  qr_state : QRState n m
  basis_size : Nat  -- rank control: max columns
  deriving Repr

/-- Validate O_AMMR node: basis_size <= m (rank control). -/
def O_AMMR_QRNode_valid (node : O_AMMR_QRNode n m) : Bool :=
  node.basis_size <= m

/-- Project O_AMMR node to Mountain type.

   From NS_MD.lean: Mountain.O_AMMR = QR Factorization Tree
   The projection validates the QR state and basis size. -/
def O_AMMR_QRNode_project (node : O_AMMR_QRNode n m) : Bool :=
  O_AMMR_QRNode_valid node

end Semantics.HouseholderQR