#!/usr/bin/env python3
"""
Generate reference values for SVD building block unit tests.
Run this to verify/implement the C++ SVD implementation against scipy/numpy.

Usage: python3 svd-reference-values.py
"""

import numpy as np
from scipy.linalg import svd, qr as scipy_qr
import json

def compute_householder(x):
    """Compute Householder reflector: H*x = [alpha, 0, 0, ...]^T.
    
    Returns (v_normalized, alpha) where v is the normalized Householder vector.
    H = I - 2*v*v^T / (v^T*v)
    """
    x = np.array(x, dtype=np.float64)
    norm_x = np.linalg.norm(x)
    
    if norm_x < 1e-30:
        return x.copy(), 0.0
    
    alpha = -np.sign(x[0]) * norm_x if x[0] != 0 else -norm_x
    v = x.copy()
    v[0] -= alpha
    v_norm = np.linalg.norm(v)
    
    if v_norm < 1e-30:
        return np.zeros_like(x), alpha
    
    v /= v_norm
    return v, alpha

def apply_householder_left(A, v, start_row):
    """Apply Householder reflection from the left: A = (I - 2vv^T) @ A.
    
    v is the normalized Householder vector operating on rows [start_row:].
    The length of v must match the number of rows affected.
    """
    A = A.copy()
    k = len(v)
    
    for col in range(A.shape[1]):
        dot = np.dot(v, A[start_row:start_row+k, col])
        A[start_row:start_row+k, col] -= 2.0 * dot * v
    
    return A

def apply_householder_right(A, v, start_col):
    """Apply Householder reflection from the right: A = A @ (I - 2vv^T).
    
    v is the normalized Householder vector operating on columns [start_col:].
    The length of v must match the number of columns affected.
    """
    A = A.copy()
    k = len(v)
    
    for row in range(A.shape[0]):
        dot = np.dot(A[row, start_col:start_col+k], v)
        A[row, start_col:start_col+k] -= 2.0 * dot * v
    
    return A

def compute_givens(x, y):
    """Compute Givens rotation that zeros out y.
    
    Returns (c, s) such that [c s; -s c] @ [x; y] = [r; 0].
    """
    r = np.sqrt(x*x + y*y)
    if r < 1e-30:
        return 1.0, 0.0
    c = x / r
    s = y / r
    return c, s

def apply_givens_left(A, i, j, c, s):
    """Apply Givens rotation from the left to rows i and j of A.
    
    [c  s]   [row_i]
    [-s  c] @ [row_j] = [new_row_i]
                          [new_row_j]
    """
    A = A.copy()
    new_i = c * A[i] + s * A[j]
    new_j = -s * A[i] + c * A[j]
    A[i] = new_i
    A[j] = new_j
    return A

def apply_givens_right(A, i, j, c, s):
    """Apply Givens rotation from the right to columns i and j of A.
    
    [col_i  col_j] @ [c -s] = [new_col_i  new_col_j]
                       [s  c]
    """
    A = A.copy()
    new_i = c * A[:, i] + s * A[:, j]
    new_j = -s * A[:, i] + c * A[:, j]
    A[:, i] = new_i
    A[:, j] = new_j
    return A

def householder_bidiagonalization(A):
    """Full Householder bidiagonalization: A = Q_L @ B @ Q_R^T.
    
    Returns (B, Q_L, Q_R) where B is upper bidiagonal.
    """
    m, n = A.shape
    p = min(m, n)
    
    QL = np.eye(m, dtype=np.float64)
    QR = np.eye(n, dtype=np.float64)
    W = A.copy()
    
    for k in range(p):
        # Left HH: zero out W[k+1:, k]
        if k < m - 1:
            x = W[k+1:, k].copy()
            v, alpha = compute_householder(x)
            if np.linalg.norm(v) > 1e-30:
                W = apply_householder_left(W, v, k + 1)
                QL = apply_householder_right(QL, v, k + 1)
        
        # Right HH: zero out W[k, k+2:] (superdiagonal)
        if k < p - 1 and k + 2 <= n:
            x = W[k, k+2:].copy()
            v, alpha = compute_householder(x)
            if np.linalg.norm(v) > 1e-30:
                W = apply_householder_right(W, v, k + 2)
                QR = apply_householder_right(QR, v, k + 2)
    
    return W, QL, QR

def implicit_qr_iteration(B, QR_acc):
    """Implicit QR iteration on a bidiagonal matrix.
    
    Returns (Sigma, QR_acc) where Sigma is diagonal with singular values
    and QR_acc contains the accumulated right transformations.
    """
    m, n = B.shape
    p = min(m, n)
    W = B.copy()
    
    max_iter = 1000
    tol = 1e-10
    
    for iteration in range(max_iter):
        # Deflate negligible subdiagonal elements
        for i in range(p - 1, 0, -1):
            if abs(W[i, i-1]) < tol * (abs(W[i-1, i-1]) + abs(W[i, i])):
                W[i, i-1] = 0.0
        
        # Find smallest unreduced block [start, end]
        start = 0
        for i in range(p - 1):
            if abs(W[i+1, i]) >= tol * (abs(W[i, i]) + abs(W[i+1, i+1])):
                start = i + 1
        
        end = p - 1
        for i in range(p - 2, -1, -1):
            if abs(W[i+1, i]) >= tol * (abs(W[i, i]) + abs(W[i+1, i+1])):
                end = i
                break
        
        if start >= end:
            continue
        
        # Wilkinson shift from bottom 2x2 corner
        a, b = W[end-1, end-1], W[end-1, end]
        c_val, d = W[end, end-1], W[end, end]
        trace = a + d
        det = a * d - b * c_val
        disc = trace**2 - 4 * det
        
        if disc >= 0:
            sqrt_disc = np.sqrt(disc)
            e1, e2 = (trace + sqrt_disc) / 2, (trace - sqrt_disc) / 2
            shift = e1 if abs(e1 - d) < abs(e2 - d) else e2
        else:
            shift = d
        
        # Implicit QR step using Givens rotations
        # Process from top to bottom within the block
        x = W[start, start] - shift
        y = W[start + 1, start]
        
        for i in range(start, end):
            r = np.sqrt(x*x + y*y)
            if r < 1e-30:
                x = W[i + 1, i]
                y = W[i + 1, i + 1] if i + 2 <= end else 0.0
                continue
            
            c_rot = x / r
            s_rot = y / r
            
            # Apply from left to rows i, i+1 (columns i..n-1)
            for j in range(i, n):
                t1, t2 = W[i, j], W[i + 1, j]
                W[i, j] = c_rot * t1 + s_rot * t2
                W[i + 1, j] = -s_rot * t1 + c_rot * t2
            
            # Apply from right to columns i, i+1 (rows 0..i)
            if i > start:
                for j in range(i + 1):
                    t1, t2 = W[j, i], W[j, i + 1]
                    W[j, i] = c_rot * t1 + s_rot * t2
                    W[j, i + 1] = -s_rot * t1 + c_rot * t2
            
            # Accumulate into QR_acc
            for j in range(QR_acc.shape[0]):
                t1, t2 = QR_acc[j, i], QR_acc[j, i + 1]
                QR_acc[j, i] = c_rot * t1 + s_rot * t2
                QR_acc[j, i + 1] = -s_rot * t1 + c_rot * t2
            
            # Prepare for next rotation
            x = W[i + 1, i]
            y = W[i + 1, i + 1] if i + 2 <= end else 0.0
    
    return W, QR_acc


def main():
    print("=" * 70)
    print("SVB BUILDING BLOCK REFERENCE VALUES")
    print("Generated with scipy/numpy for C++ unit test verification")
    print("=" * 70)
    
    # ------------------------------------------------------------------
    # Test 1: Householder Vector Computation
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 1: computeHouseholderVector")
    print("=" * 70)
    
    test_vectors = [
        ("2D [1,3]", [1.0, 3.0]),
        ("2D [3,4] (norm=5)", [3.0, 4.0]),
        ("3D [1,2,3]", [1.0, 2.0, 3.0]),
        ("3D [0,0,1]", [0.0, 0.0, 1.0]),
        ("4D [5,-3,2,1]", [5.0, -3.0, 2.0, 1.0]),
    ]
    
    for name, vec in test_vectors:
        v, alpha = compute_householder(vec)
        x = np.array(vec)
        Hx = x - 2 * np.dot(v, x) * v
        
        print(f"\n{name}:")
        print(f"  Input:            {list(x)}")
        print(f"  ||x||:            {np.linalg.norm(x):.15f}")
        print(f"  alpha:            {alpha:.15f}")
        print(f"  v (normalized):   {[round(float(vi), 12) for vi in v]}")
        print(f"  H*x = [alpha,0..]: {[round(float(xi), 12) for xi in Hx]}")
        print(f"  Off-diagonal ~0:  {np.allclose(Hx[1:], 0, atol=1e-12)}")
    
    # ------------------------------------------------------------------
    # Test 2: Householder Apply Left
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 2: applyHouseholderLeft")
    print("=" * 70)
    
    A_test = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float64)
    x_col = A_test[1:, 0].copy()
    v_left, _ = compute_householder(x_col)
    
    print(f"\nInput matrix:\n{A_test}")
    print(f"Householder vector (rows 1:3): {[round(float(vi), 12) for vi in v_left]}")
    
    A_result = apply_householder_left(A_test, v_left, 1)
    print(f"\nAfter applyHouseholderLeft:\n{A_result}")
    print(f"  A[1,0] = {A_result[1,0]:.2e}, A[2,0] = {A_result[2,0]:.2e} (should be ~0)")
    
    # ------------------------------------------------------------------
    # Test 3: Householder Apply Right
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 3: applyHouseholderRight")
    print("=" * 70)
    
    A_test = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=np.float64)
    x_row = A_test[0, 1:].copy()
    v_right, _ = compute_householder(x_row)
    
    print(f"\nInput matrix:\n{A_test}")
    print(f"Householder vector (cols 1:3): {[round(float(vi), 12) for vi in v_right]}")
    
    A_result = apply_householder_right(A_test, v_right, 1)
    print(f"\nAfter applyHouseholderRight:\n{A_result}")
    print(f"  A[0,1] = {A_result[0,1]:.2e}, A[0,2] = {A_result[0,2]:.2e} (should be ~0)")
    
    # ------------------------------------------------------------------
    # Test 4: Givens Rotation Computation
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 4: computeGivens")
    print("=" * 70)
    
    givens_tests = [
        ("3-4-5 triangle", 3.0, 4.0),
        ("y already zero", 1.0, 0.0),
        ("x is zero", 0.0, 5.0),
        ("Both negative", -3.0, -4.0),
        ("45 degree case", 1.0, -1.0),
    ]
    
    for name, x, y in givens_tests:
        c, s = compute_givens(x, y)
        result_x = c * x + s * y
        result_y = -s * x + c * y
        
        print(f"\n{name}: x={x}, y={y}")
        print(f"  r = {np.sqrt(x*x+y*y):.12f}")
        print(f"  c = {c:.12f}, s = {s:.12f}")
        print(f"  [c s; -s c] @ [x;y] = [{result_x:.2e}, {result_y:.2e}]")
    
    # ------------------------------------------------------------------
    # Test 5: Apply Givens Left/Right
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 5: applyGivensLeft / applyGivensRight")
    print("=" * 70)
    
    A_test = np.array([[3.0, 4.0], [1.0, 2.0]], dtype=np.float64)
    c, s = compute_givens(3.0, 1.0)
    
    print(f"\nInput matrix:\n{A_test}")
    print(f"Givens rotation (rows 0,1): c={c:.12f}, s={s:.12f}")
    
    A_left = apply_givens_left(A_test, 0, 1, c, s)
    print(f"\nAfter applyGivensLeft:\n{A_left}")
    print(f"  A[1,0] = {A_left[1,0]:.2e} (should be ~0)")
    
    A_test = np.array([[3.0, 1.0], [4.0, 2.0]], dtype=np.float64)
    c, s = compute_givens(3.0, 4.0)
    
    print(f"\nInput matrix:\n{A_test}")
    print(f"Givens rotation (cols 0,1): c={c:.12f}, s={s:.12f}")
    
    A_right = apply_givens_right(A_test, 0, 1, c, s)
    print(f"\nAfter applyGivensRight:\n{A_right}")
    print(f"  A[0,1] = {A_right[0,1]:.2e} (should be ~0)")
    
    # ------------------------------------------------------------------
    # Test 6: Full Bidiagonalization
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 6: householderBidiagonalization")
    print("=" * 70)
    
    bidiag_tests = [
        ("2x2 [[1,2],[3,4]]", np.array([[1.0, 2.0], [3.0, 4.0]])),
        ("3x3 SPD [[5,3],[3,5]]", np.array([[5.0, 3.0], [3.0, 5.0]])),
        ("3x3 diag [[10,0,0],[0,5,0],[0,0,2]]", 
         np.array([[10.0, 0, 0], [0, 5.0, 0], [0, 0, 2.0]])),
        ("3x3 full [[1,2,3],[4,5,6],[7,8,10]]", 
         np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 10.0]])),
        ("Tall 4x3", np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], dtype=np.float64)),
    ]
    
    for name, A in bidiag_tests:
        B, QL, QR = householder_bidiagonalization(A)
        m, n = A.shape
        p = min(m, n)
        
        print(f"\n{name}:")
        print(f"  Original:\n{A}")
        print(f"\n  Bidiagonal B:\n{B}")
        print(f"  Diagonal: {[round(float(B[i,i]), 10) for i in range(p)]}")
        print(f"  Superdiag: {[round(float(B[i,i+1]), 10) for i in range(min(p-1, n-1))]}")
        
        recon = QL @ B @ QR.T
        err = np.linalg.norm(recon - A, 'fro')
        print(f"  ||QL @ B @ QR^T - A||_F = {err:.2e}")
    
    # ------------------------------------------------------------------
    # Test 7: Full SVD Reference Values
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 7: Full SVD Reference Values (scipy.linalg.svd)")
    print("=" * 70)
    
    test_matrices = [
        ("Simple 2x2", np.array([[1,2],[3,4]], dtype=np.float64)),
        ("SPD 2x2", np.array([[5,3],[3,5]], dtype=np.float64)),
        ("Full-rank 3x3", np.array([[1,2,3],[4,5,6],[7,8,10]], dtype=np.float64)),
        ("Rank-deficient 3x3", np.array([[1,2,3],[4,5,6],[7,8,9]], dtype=np.float64)),
        ("Diagonal 3x3", np.array([[10,0,0],[0,5,0],[0,0,2]], dtype=np.float64)),
        ("Tall 4x3", np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], dtype=np.float64)),
        ("Wide 3x5", np.array([[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15]], dtype=np.float64)),
        ("Symmetric tri 5x5", np.array([[2,-1,0,0,0],[-1,2,-1,0,0],[0,-1,2,-1,0],[0,0,-1,2,-1],[0,0,0,-1,2]], dtype=np.float64)),
        ("Neg values 2x3", np.array([[0.5,-0.3,0.8],[-0.2,0.7,0.1]], dtype=np.float64)),
        ("Near-singular 2x2", np.array([[1,0],[0,1e-6]], dtype=np.float64)),
        ("Orthogonal 3x3", np.array([[np.cos(np.pi/4), -np.sin(np.pi/4), 0],
                                      [np.sin(np.pi/4), np.cos(np.pi/4), 0],
                                      [0, 0, 1]], dtype=np.float64)),
        ("Identity 3x3", np.eye(3)),
        ("Zero 3x3", np.zeros((3,3))),
        ("Col vector 2x1", np.array([[3],[4]], dtype=np.float64)),
        ("Row vector 1x2", np.array([[3,4]], dtype=np.float64)),
    ]
    
    for name, A in test_matrices:
        U, s, Vt = svd(A, full_matrices=False)
        
        print(f"\n{name}: shape={A.shape}")
        print(f"  Singular values: {[round(float(x), 12) for x in s]}")
        print(f"  U:\n{np.array2string(U, precision=6, floatmode='maxprec_equal')}")
        print(f"  Vt:\n{np.array2string(Vt, precision=6, floatmode='maxprec_equal')}")
        recon_err = np.linalg.norm(A - U @ np.diag(s) @ Vt, 'fro')
        print(f"  Reconstruction error: {recon_err:.2e}")
    
    # ------------------------------------------------------------------
    # Test 8: Implicit QR Iteration on Bidiagonal
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("TEST 8: implicitQRIteration")
    print("=" * 70)
    
    qr_tests = [
        ("2x2 [[1,2],[3,4]]", np.array([[1.0, 2.0], [3.0, 4.0]])),
        ("3x3 diag", np.array([[10.0, 0, 0], [0, 5.0, 0], [0, 0, 2.0]])),
    ]
    
    for name, A in qr_tests:
        B, QL, QR = householder_bidiagonalization(A)
        Sigma, QR_final = implicit_qr_iteration(B.copy(), QR.copy())
        
        print(f"\n{name}:")
        print(f"  Bidiagonal B:\n{B}")
        print(f"  After QR iteration (Sigma):\n{Sigma}")
        print(f"  Diagonal entries: {[round(float(Sigma[i,i]), 10) for i in range(min(Sigma.shape))]}")
        
        # Verify: QL @ Sigma @ QR_final^T ≈ A
        recon = QL @ Sigma @ QR_final.T
        err = np.linalg.norm(recon - A, 'fro')
        print(f"  ||QL @ Sigma @ QR^T - A||_F = {err:.2e}")
    
    # ------------------------------------------------------------------
    # JSON output for easy import into C++ tests
    # ------------------------------------------------------------------
    print("\n" + "=" * 70)
    print("JSON OUTPUT (for easy C++ integration)")
    print("=" * 70)
    
    json_data = {}
    
    # Householder test vectors
    hh_tests = {}
    for name, vec in test_vectors:
        v, alpha = compute_householder(vec)
        x = np.array(vec)
        Hx = x - 2 * np.dot(v, x) * v
        hh_tests[name] = {
            "input": [float(xi) for xi in x],
            "norm": float(np.linalg.norm(x)),
            "alpha": float(alpha),
            "v_normalized": [round(float(vi), 12) for vi in v],
            "Hx": [round(float(xi), 12) for xi in Hx],
        }
    json_data["householder_vectors"] = hh_tests
    
    # Full SVD reference values
    svd_tests = {}
    for name, A in test_matrices:
        U, s, Vt = svd(A, full_matrices=False)
        svd_tests[name] = {
            "shape": list(A.shape),
            "singular_values": [round(float(x), 12) for x in s],
            "U": [[round(float(U[i,j]), 8) for j in range(U.shape[1])] for i in range(U.shape[0])],
            "Vt": [[round(float(Vt[i,j]), 8) for j in range(Vt.shape[1])] for i in range(Vt.shape[0])],
        }
    json_data["svd_reference"] = svd_tests
    
    print(json.dumps(json_data, indent=2))


if __name__ == "__main__":
    main()