import json
from pathlib import Path
import numpy as np

SEED = 2486

def gale(X):
    """Return rows spanning the left nullspace of X (X is a x b)."""
    u, s, vh = np.linalg.svd(X.T, full_matrices=True)
    a, b = X.shape
    return u[:, a:].T, s

def normalized_penalty(Y, X, D=None, eps=1e-12):
    if D is None:
        D = np.eye(X.shape[1])
    r = Y @ D @ X.T
    return np.linalg.norm(r, 'fro')**2 / (np.linalg.norm(Y, 'fro')**2 * np.linalg.norm(X, 'fro')**2 + eps)

def prediction_sweeps(rng):
    # P1: exact SVD Gale residual is floating-point roundoff for every valid size.
    exact = []
    for a, b in [(2, 5), (3, 7), (5, 9), (8, 13)]:
        X = rng.normal(size=(a, b))
        Y, _ = gale(X)
        exact.append({'a': a, 'b': b, 'residual_norm': float(np.linalg.norm(Y @ X.T)),
                      'relative_penalty': float(normalized_penalty(Y, X)),
                      'rank': int(np.linalg.matrix_rank(X))})

    # P2: perturbing Y by gamma Z predicts residual norm O(gamma), penalty O(gamma^2).
    a, b = 4, 9
    X = rng.normal(size=(a, b)); Y, _ = gale(X)
    Z = rng.normal(size=Y.shape); Z /= np.linalg.norm(Z)
    gammas = np.logspace(-6, -1, 8)
    rows = []
    for g in gammas:
        Yn = Y + g * Z
        rn = np.linalg.norm(Yn @ X.T)
        p = normalized_penalty(Yn, X)
        rows.append({'gamma': float(g), 'residual_norm': float(rn), 'penalty': float(p),
                     'norm_over_gamma': float(rn / g), 'penalty_over_gamma2': float(p / g**2)})
    slope_norm = float(np.polyfit(np.log(gammas), np.log([r['residual_norm'] for r in rows]), 1)[0])
    slope_pen = float(np.polyfit(np.log(gammas), np.log([r['penalty'] for r in rows]), 1)[0])

    # P3: D=cI preserves the null relation; a nonconstant diagonal gauge generally breaks it.
    Drows = []
    for spread in [0., .01, .03, .1, .3, 1.0]:
        logd = rng.normal(0, spread, size=b)
        D = np.diag(np.exp(logd))
        Drows.append({'spread': spread, 'residual_norm': float(np.linalg.norm(Y @ D @ X.T)),
                      'penalty': float(normalized_penalty(Y, X, D))})
    positive = [r for r in Drows if r['spread'] > 0]
    gauge_slope = float(np.polyfit(np.log([r['spread'] for r in positive]), np.log([r['residual_norm'] for r in positive]), 1)[0])
    return exact, rows, slope_norm, slope_pen, Drows, gauge_slope

def fit_ridge(A, y, lam=1e-6):
    # Deterministic least-squares ridge, with intercept.
    A1 = np.column_stack([A, np.ones(len(A))])
    reg = lam * np.eye(A1.shape[1]); reg[-1, -1] = 0
    return np.linalg.solve(A1.T @ A1 + reg, A1.T @ y)

def toy_comparison(rng):
    # A tiny representation probe: primary summary has a dimensions; adding
    # YX^T supplies an orthogonal (b-a)-dimensional feature summary.
    a, b, d, n = 3, 8, 12, 400
    W = rng.normal(size=(a, d)); W /= np.linalg.norm(W, axis=1, keepdims=True)
    primary, combined, targets = [], [], []
    for _ in range(n):
        H = rng.normal(size=(d, b))
        P = W @ H
        Y, _ = gale(P)
        dual = Y @ H.T                         # (b-a) x d
        p = P.mean(axis=1)
        q = dual.mean(axis=1)
        primary.append(p)
        combined.append(np.concatenate([p, q]))
        # Target depends on both channels; this tests whether the complement
        # carries useful information, not whether it changes the null identity.
        targets.append(1.2 * p[0] - .7 * p[1] + .5 * q[0] + .8 * q[1] + .05 * rng.normal())
    primary, combined, targets = map(np.asarray, (primary, combined, targets))
    cut = 280
    wb = fit_ridge(primary[:cut], targets[:cut]); wg = fit_ridge(combined[:cut], targets[:cut])
    pb = np.column_stack([primary[cut:], np.ones(n-cut)]) @ wb
    pg = np.column_stack([combined[cut:], np.ones(n-cut)]) @ wg
    return {'primary_test_mse': float(np.mean((pb-targets[cut:])**2)),
            'gale_combined_test_mse': float(np.mean((pg-targets[cut:])**2)),
            'feature_width_primary': a, 'feature_width_combined': a + (b-a)}

def main():
    rng = np.random.default_rng(SEED)
    exact, rows, sn, sp, drows, gs = prediction_sweeps(rng)
    result = {'seed': SEED, 'predictions': {
        'P1_exact_nullspace': '||YX^T|| remains at floating-point roundoff',
        'P2_noise_scaling': '||YX^T|| ~ gamma^1 and normalized penalty ~ gamma^2',
        'P3_diagonal_gauge': 'D=cI preserves zero; nonconstant D produces residual'},
        'exact_sweep': exact, 'noise_sweep': rows,
        'noise_loglog_slopes': {'residual_norm': sn, 'penalty': sp},
        'gauge_loglog_slope': gs,
        'diagonal_sweep': drows, 'toy_comparison': toy_comparison(rng)}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()