/-
Copyright (c) 2026 SilverSight Contributors. All rights reserved.
Released under Apache 2.0 license.

Generic n×n matrix operations for spectral analysis.
All functions are dimension-polymorphic via an explicit (n : Nat) parameter.
-/

import SilverSight.FixedPoint

namespace SilverSight.PIST.MatrixN

open SilverSight.FixedPoint
open SilverSight.FixedPoint.Q16_16

/-- Safely get an entry from a possibly ragged 2D array. -/
@[inline]
def getEntry (mat : Array (Array Int)) (i j : Nat) : Int :=
  mat.getD i #[] |>.getD j 0

/-- Row sum of row i in an n×n matrix. -/
def rowSum (mat : Array (Array Int)) (i n : Nat) : Int :=
  (List.range n).foldl (fun acc j => acc + getEntry mat i j) 0

/-- Symmetrize an n×n matrix: (A + Aᵀ)/2. -/
def symmetrize (mat : Array (Array Int)) (n : Nat) : Array (Array Int) :=
  Array.ofFn (n := n) fun i =>
    Array.ofFn (n := n) fun j =>
      (getEntry mat i.val j.val + getEntry mat j.val i.val) / 2

/-- Build the Laplacian L = D - A from a symmetric matrix. -/
def buildLaplacian (sym : Array (Array Int)) (n : Nat) : Array (Array Int) :=
  Array.ofFn (n := n) fun i =>
    let deg := rowSum sym i.val n
    Array.ofFn (n := n) fun j =>
      if i.val = j.val then deg else -(getEntry sym i.val j.val)

/-- Build AᵀA (Gram matrix) from an n×n matrix. -/
def buildATA (mat : Array (Array Int)) (n : Nat) : Array (Array Int) :=
  Array.ofFn (n := n) fun i =>
    Array.ofFn (n := n) fun j =>
      (List.range n).foldl (fun acc k =>
        acc + getEntry mat k i.val * getEntry mat k j.val) 0

/-- Squared norm of a Q16_16 vector (sum of squares of raw Ints). -/
def normSqRaw (v : Array Q16_16) : Int :=
  v.foldl (fun acc x => acc + x.toInt * x.toInt) 0

/- Matrix-vector multiplication: mat[n×n] × v[n]. -/
def matVecMul (mat : Array (Array Int)) (n : Nat) (v : Array Q16_16) : Array Q16_16 :=
  Array.ofFn (n := n) fun i =>
    let s : Int := (List.range n).foldl (fun acc j =>
      acc + getEntry mat i.val j * (v.getD j zero).toInt) 0
    ofRawInt s

/-- Integer square root via Newton's method. -/
def isqrt (n : Int) : Int :=
  if n ≤ 0 then 0
  else
    let rec loop (x : Int) (fuel : Nat) : Int :=
      match fuel with
      | 0     => x
      | f + 1 =>
        let x' := (x + n / x) / 2
        if x' ≥ x then x else loop x' f
    loop (n / 2 + 1) 64

/-- Dominant eigenvalue of an n×n Int matrix via power iteration.

    ⚠️ KNOWN LIMITATION: Power iteration fails to converge on matrices
    with peripheral spectrum (eigenvalues on the unit circle at non-trivial
    angles). In the 250-equation corpus, 20 matrices have exact ρ = 1.0
    but power iteration produces values in [0.67, 0.95].

    For exact eigenvalues, use the characteristic polynomial approach.
    This function is retained for backward compatibility with existing
    proofs that depend on its specific behavior.
    -/
def powerIteration (mat : Array (Array Int)) (n : Nat) (maxIter : Nat := 100) : Q16_16 :=
  if n = 0 then zero
  else
    let initVal := ofRawInt ((q16Scale : Int) / (n : Int))
    let v₀ : Array Q16_16 := Array.replicate n initVal
    let rec iterate (v : Array Q16_16) (fuel : Nat) : Array Q16_16 :=
      match fuel with
      | 0     => v
      | f + 1 =>
        let mv  := matVecMul mat n v
        let nm2 := normSqRaw mv
        if nm2 ≤ 0 then v
        else
          let nm := isqrt nm2
          if nm = 0 then v
          else
            let vn := mv.map (fun x => ofRawInt (x.toInt * q16Scale / nm))
            iterate vn f
    let vFinal := iterate v₀ maxIter
    let mv  := matVecMul mat n vFinal
    let num : Int := (List.range n).foldl (fun acc i =>
      acc + (vFinal.getD i zero).toInt * (mv.getD i zero).toInt) 0
    let den : Int := normSqRaw vFinal
    if den ≤ 0 then zero
    else ofRawInt (num * q16Scale / den)

#eval isqrt 9    -- expect 3
#eval isqrt 16   -- expect 4
#eval isqrt 200  -- expect 14

#eval! (powerIteration #[#[1, 0], #[0, 1]] 2).toInt  -- expect 65536

end SilverSight.PIST.MatrixN
