#!/usr/bin/env python3
"""
Reference values for the QR eigen-decomposition building block tests
(unit-tests/qr-build-blocks-tests.cpp). Run this to verify/implement the
C++ implementation in src/QR.hpp / src/QR.cpp against numpy/scipy.

Conventions (match the C++ exactly):
  * Givens zeroing rotation: G = [[c, s], [-s, c]], c = x/r, s = y/r,
    r = hypot(x, y).  G * (x, y)^T = (r, 0)^T.
  * Similarity transform: A <- G A G^T (ApplyRotationBothSides).
  * Eigenvector accumulation: V <- V G^T (ApplyRotationToVectors).
    Vblock in the 2x2 closed form is [[c, -s], [s, c]] (same shape as G^T).
  * Tridiagonalization: bottom-up Givens (i = N-2 down to k+1 per column k).
  * Shifted QR loop: Wilkinson shift mu from the trailing 2x2, chase on the
    trailing unreduced block [lo, hi], deflate by relative tolerance, peel
    exact-zero subdiagonals, 2x2 closed-form termination.
  * Pipeline: M0 = U * Mtri * U^T and Mtri = V * D * V^T  =>
    eigenvectors of M0 = U * V (columns), eigenvalues = diag(D).

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

import numpy as np
import scipy.linalg as sla

np.set_printoptions(precision=9, linewidth=120)


def givens(x, y):
    """c = x/r, s = y/r with r = hypot(x, y)."""
    r = np.hypot(x, y)
    if r == 0.0:
        return 1.0, 0.0
    return x / r, y / r


def rot(n, i, c, s):
    """G = I with [[c, s], [-s, c]] embedded at (i, i+1)."""
    G = np.eye(n)
    G[i:i + 2, i:i + 2] = np.array([[c, s], [-s, c]])
    return G


def tridiagonalize(M0):
    """Bottom-up Givens tridiagonalization. Returns (Mtri, U) with
    M0 = U Mtri U^T."""
    n = len(M0)
    M = M0.copy()
    U = np.eye(n)
    for k in range(n - 2):
        for i in range(n - 2, k, -1):
            c, s = givens(M[i, k], M[i + 1, k])
            G = rot(n, i, c, s)
            M = G @ M @ G.T
            U = U @ G.T
    return M, U


def wilkinson(a, b, d):
    """Eigenvalue of [[a, b], [b, d]] closest to d."""
    delta = 0.5 * (a - d)
    spread = np.sqrt(delta * delta + b * b)
    return 0.5 * (a + d) - (spread if delta >= 0 else -spread)


def solve2x2(A, lo):
    """Closed form for the block at (lo, lo+1): (lHi, lLo, c, s) with
    vHi = (c, s), vLo = (-s, c)."""
    a = A[lo, lo]
    b = A[lo, lo + 1]
    e = A[lo + 1, lo]
    d = A[lo + 1, lo + 1]
    tr = a + d
    det = a * d - b * e
    disc = max(0.0, tr * tr - 4 * det)
    lhi = 0.5 * (tr + np.sqrt(disc))
    llo = 0.5 * (tr - np.sqrt(disc))
    if b != 0.0:
        v1 = lhi - a
        nn = np.hypot(b, v1)
        c, s = b / nn, v1 / nn
    elif a >= d:
        c, s = 1.0, 0.0
    else:
        c, s = 0.0, 1.0
    return lhi, llo, c, s


def eigenqr(M0, tol=1e-12, max_iter=100000):
    """Full pipeline mirroring QR::EigenQR. Returns (eigs, W) where W has
    the eigenvectors of M0 as columns."""
    n = len(M0)
    if n == 2:
        l1, l2, c, s = solve2x2(M0, 0)
        return np.array([l1, l2]), np.array([[c, -s], [s, c]])
    M, U = tridiagonalize(M0)
    V = np.eye(n)
    hi = n - 1
    for _ in range(max_iter):
        # deflate: zero tiny subdiagonals (relative test)
        for i in range(hi):
            t = M[i + 1, i]
            scale = abs(M[i, i]) + abs(M[i + 1, i + 1])
            if abs(t) <= tol * scale:
                M[i + 1, i] = M[i, i + 1] = 0.0
        # peel exact-zero trailing subdiagonals
        while hi > 0 and M[hi, hi - 1] == 0.0:
            hi -= 1
        if hi == 0:
            break
        # find start of trailing unreduced block
        lo = hi
        for i in range(hi - 1, -1, -1):
            if M[i + 1, i] == 0.0:
                break
            lo = i
        if lo + 1 == hi:
            # closed-form 2x2 termination: set diagonal, fold Vblock in
            l1, l2, c, s = solve2x2(M, lo)
            Vb = np.eye(n)
            Vb[lo:lo + 2, lo:lo + 2] = np.array([[c, -s], [s, c]])
            V = V @ Vb
            M[lo, lo] = l1
            M[lo + 1, lo + 1] = l2
            M[lo + 1, lo] = M[lo, lo + 1] = 0.0
            if lo == 0:
                break
            hi = lo - 1
            continue
        # full shifted step on [lo, hi] (shift applies to the active block)
        mu = wilkinson(M[hi - 1, hi - 1], M[hi, hi - 1], M[hi, hi])
        diag = M.diagonal().copy()
        diag[lo:hi + 1] -= mu
        np.fill_diagonal(M, diag)
        c, s = givens(M[lo, lo], M[lo + 1, lo])
        G = rot(n, lo, c, s)
        M = G @ M @ G.T
        V = V @ G.T
        for i in range(lo + 1, hi):
            c, s = givens(M[i, i], M[i + 1, i])
            G = rot(n, i, c, s)
            M = G @ M @ G.T
            V = V @ G.T
        diag = M.diagonal().copy()
        diag[lo:hi + 1] += mu
        np.fill_diagonal(M, diag)

    eigs = np.diag(M).astype(float)
    order = np.argsort(eigs)[::-1]  # descending, like the C++ test harness
    eigs = eigs[order]
    W = U @ V
    W = W[:, order]
    return eigs, W


def report(name, val, ref=None, tol=1e-6):
    ok = "OK " if ref is None or np.allclose(val, ref, rtol=tol, atol=tol) else "FAIL"
    print(f"[{ok}] {name} = {val}")
    if ref is not None:
        print(f"        scipy/numpy ref = {ref}")


def main():
    print("=== TEST 1: GivensRotation ===")
    c, s = givens(2.0, 1.0)
    print(f"  c = {c}  s = {s}")
    # G * (x, y)^T = (r, 0)^T:  G = [[c, s], [-s, c]]
    assert abs(c * 2 + s * 1 - np.sqrt(5)) < 1e-15
    assert abs(-s * 2 + c * 1) < 1e-15

    print("\n=== TEST 2: ApplyRotationBothSides A <- G A G^T ===")
    A = np.array([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0], [9.0, 10.0, 11.0]])
    G = rot(3, 0, 0.6, 0.8)
    B = G @ A @ G.T
    print(B)

    A = np.array([[5.0, 0.0, 1.0], [0.0, 6.0, 2.0], [1.0, 2.0, 7.0]])
    c, s = givens(6.0, 2.0)
    G = rot(3, 1, c, s)
    B = G @ A @ G.T
    print(B)

    print("\n=== TEST 3: V accumulation V <- V G^T ===")
    V = np.eye(3)
    G = rot(3, 0, 0.894427191, 0.447213595)
    V = V @ G.T
    print(V)
    V2 = V @ rot(3, 1, 0.848874681, 0.528748047).T
    print(V2)

    print("\n=== TEST 4: Solve2x2Eigen ===")
    for A in (np.array([[5.0, 8.0], [8.0, 9.0]]), np.array([[1.0, 2.0], [3.0, 4.0]])):
        l1, l2, c, s = solve2x2(A, 0)
        ref = np.linalg.eigvalsh(A) if np.allclose(A, A.T) else np.linalg.eigvals(A)
        print(f"  A={A.ravel()}  lHi={l1} lLo={l2}  c={c} s={s}  ref={np.sort(ref)[::-1]}")

    print("\n=== TEST 8: WilkinsonShift ===")
    print(f"  W(5, 8, 9) = {wilkinson(5, 8, 9)}")
    print(f"  W(4, 2, 7) = {wilkinson(4, 2, 7)}")
    print(f"  W(9, 2, 5) = {wilkinson(9, 2, 5)}")

    print("\n=== TEST 8b: one full shifted chase step on tridiagonal 3x3 ===")
    T = np.array([[1.0, 2.0, 0.0], [2.0, 5.0, 2.0], [0.0, 2.0, 9.0]])
    mu = wilkinson(5, 2, 9)
    M = T - mu * np.eye(3)
    c, s = givens(M[0, 0], M[1, 0])
    M = rot(3, 0, c, s) @ M @ rot(3, 0, c, s).T
    c, s = givens(M[1, 1], M[2, 1])
    M = rot(3, 1, c, s) @ M @ rot(3, 1, c, s).T
    M = M + mu * np.eye(3)
    print(f"  mu = {mu}")
    print(M)
    print(f"  corners: {M[0, 2]}, {M[2, 0]} (exact-arithmetic zeros)")
    print(f"  trace {M.trace():.15f} (was {T.trace()})")

    print("\n=== TEST 7: Tridiagonalize ===")
    M4 = np.array([[2.0, 1, 0, 1], [1, 3, 1, 0], [0, 1, 4, 1], [1, 0, 1, 5]])
    M, U = tridiagonalize(M4)
    print("  M4 tridiagonalized:\n", M)
    print(f"  reconstruction U M U^T == M4: {np.allclose(U @ M @ U.T, M4, atol=1e-9)}")
    M5 = np.array([[3.0, 1, 0, 0, 1], [1, 4, 1, 0, 0], [0, 1, 5, 1, 0],
                   [0, 0, 1, 6, 1], [1, 0, 0, 1, 7]])
    M, U = tridiagonalize(M5)
    print("  M5 tridiagonalized:\n", M)
    print(f"  reconstruction: {np.allclose(U @ M @ U.T, M5, atol=1e-9)}")

    print("\n=== End-to-end: random symmetric vs scipy.linalg.eigh ===")
    rng = np.random.default_rng(12345)
    worst = 0.0
    for n in range(3, 9):
        M0 = rng.normal(size=(n, n))
        M0 = (M0 + M0.T) / 2
        eigs, W = eigenqr(M0.astype(float))
        ref = sla.eigh(M0)
        e_err = np.max(np.abs(np.sort(eigs) - ref[0]))
        resid = np.linalg.norm(W @ np.diag(eigs) @ W.T - M0)
        ortho = np.linalg.norm(W.T @ W - np.eye(n))
        print(f"  n={n}: eigs_err={e_err:.2e}  resid={resid:.2e}  ortho={ortho:.2e}")
        worst = max(worst, e_err, resid, ortho)
    print(f"\nworst over all n: {worst:.2e}")
    assert worst < 1e-10, "end-to-end reference FAILED"
    print("ALL REFERENCES OK")


if __name__ == "__main__":
    main()