import json, math
import numpy as np


def householder_qr_factors(A):
    """Unblocked Householder QR, returning reflector factors and R.
    The factors (v,tau) are an implicit representation of Q; no Q is stored.
    """
    A = np.array(A, dtype=np.float64, copy=True)
    m, n = A.shape
    p = min(m, n)
    factors = []
    for j in range(p):
        x = A[j:, j]
        nx = np.linalg.norm(x)
        if nx == 0:
            factors.append((np.zeros_like(x), 0.0)); continue
        # Sign choice avoids cancellation.
        alpha = -math.copysign(nx, x[0])
        v = x.copy(); v[0] -= alpha
        vv = float(v @ v)
        tau = 0.0 if vv == 0 else 2.0 / vv
        A[j:, j:] -= np.outer(v, tau * (v @ A[j:, j:]))
        A[j, j] = alpha
        A[j+1:, j] = 0.0
        factors.append((v, tau))
    return factors, A


def apply_factors_left(factors, X, transpose=False):
    """Apply product of stored reflectors to X (reflectors are symmetric)."""
    Y = np.array(X, dtype=np.float64, copy=True)
    seq = factors if transpose else factors[::-1]
    for j, (v, tau) in enumerate(seq):
        row = j if transpose else len(factors)-1-j
        # In reverse traversal, the factor's original row is row.
        if transpose:
            row = j
        Y[row:] -= np.outer(v, tau * (v @ Y[row:]))
    return Y


def householder_block_basis(Y):
    fac, R = householder_qr_factors(Y)
    # Q(:,1:b) is obtained by applying Q to first b coordinate vectors.
    b = Y.shape[1]
    E = np.eye(Y.shape[0], b)
    Q = apply_factors_left(fac, E, transpose=False)
    return Q, fac, R


def gs_block_basis(Y):
    """One-pass classical Gram-Schmidt, intentionally no reorthogonalization."""
    m, b = Y.shape; Q = np.zeros((m, b)); r = 0
    for j in range(b):
        v = Y[:, j].copy()
        if r: v -= Q[:, :r] @ (Q[:, :r].T @ v)
        nv = np.linalg.norm(v)
        if nv > 1e-14:
            Q[:, r] = v / nv; r += 1
    return Q[:, :r]


def adaptive(A, sigma, block, method="householder", kmax=None, seed=0):
    rng = np.random.default_rng(seed)
    m, n = A.shape; kmax = min(m, n) if kmax is None else min(kmax, m)
    R = A.copy(); initial = float(np.sum(R*R)); E = initial
    blocks = []; cols = 0
    while E > sigma*sigma and cols < kmax:
        b = min(block, kmax-cols)
        # Fresh Gaussian range sketch of the current residual.
        Y = R @ rng.standard_normal((n, b))
        if method == "householder": Q, _, _ = householder_block_basis(Y)
        else: Q = gs_block_basis(Y)
        if Q.shape[1] == 0: break
        # Q is already orthogonal to previous blocks because it is formed from R.
        R -= Q @ (Q.T @ R)
        blocks.append(Q); cols += Q.shape[1]; E = float(np.sum(R*R))
    Qall = np.concatenate(blocks, axis=1) if blocks else np.zeros((m,0))
    return {"Q": Qall, "residual_energy": E, "initial_energy": initial,
            "rank": Qall.shape[1], "orth_error": float(np.linalg.norm(Qall.T@Qall-np.eye(Qall.shape[1]))) }


def matrix_with_spectrum(m, n, decay):
    rng = np.random.default_rng(1234 + int(decay*1000))
    U,_=np.linalg.qr(rng.standard_normal((m,m))); V,_=np.linalg.qr(rng.standard_normal((n,n)))
    s=np.exp(-decay*np.arange(min(m,n)))
    return U[:, :min(m,n)] @ np.diag(s) @ V[:min(m,n), :]


def run():
    # Prediction 1: exact Householder orthogonality stays near epsilon as blocks grow;
    # one-pass GS worsens strongly for ill-conditioned panels.
    rng=np.random.default_rng(7); m=96; b=4
    cond_rows=[]
    for cond in [1e2,1e6,1e10,1e14]:
        x=np.linspace(0,1,m); Y=np.column_stack([x**j + 1e-14*rng.standard_normal(m) for j in range(b)])
        # Make the conditioning control explicit through singular values.
        U,_=np.linalg.qr(rng.standard_normal((m,b))); W,_=np.linalg.qr(rng.standard_normal((b,b)))
        Y=U@np.diag(np.geomspace(1,1/cond,b))@W.T
        qh,_,_=householder_block_basis(Y); qg=gs_block_basis(Y)
        cond_rows.append([cond,float(np.linalg.norm(qh.T@qh-np.eye(b))),float(np.linalg.norm(qg.T@qg-np.eye(qg.shape[1])))])
    # Prediction 2: for geometric spectrum, rank grows monotonically as sigma tightens.
    A=matrix_with_spectrum(72,48,0.16); normA=np.linalg.norm(A,'fro')
    tol_rows=[]
    for rel in [0.5,0.25,0.12,0.06,0.03]:
        out=adaptive(A, rel*normA, 4, "householder", seed=11)
        tol_rows.append([rel,out['rank'],math.sqrt(out['residual_energy'])/normA, out['orth_error']])
    # Prediction 3: captured residual is below prescribed sigma (up to sampling/roundoff).
    checks=[]
    for rel in [0.4,0.2,0.1]:
        out=adaptive(A, rel*normA, 4, "householder", seed=21)
        checks.append([rel, math.sqrt(out['residual_energy'])/normA, out['rank']])
    # Secondary baseline comparison at same max rank.
    base=adaptive(A, .1*normA, 4, "gs", seed=21); idea=adaptive(A,.1*normA,4,"householder",seed=21)
    result={"predicted": {"orthogonality": "Householder O(eps), GS grows with cond", "rank": "tighter sigma -> nondecreasing rank", "residual": "relative residual <= sigma/||A||"}, "orthogonality_sweep":cond_rows,"tolerance_sweep":tol_rows,"residual_checks":checks,"comparison":{"gs":{k:base[k] for k in ['rank','residual_energy','orth_error']},"householder":{k:idea[k] for k in ['rank','residual_energy','orth_error']}}}
    print(json.dumps(result, indent=2))

if __name__ == '__main__': run()