#!/usr/bin/env python3
"""Schur-Riesz Greedy Adapter Expansion numerical MVP."""
import json
import numpy as np

SEED = 2145


def weighted_projector(B, y, rtol=1e-11):
    y = np.asarray(y, float)
    n = len(y)
    if B.size == 0:
        return np.zeros((n, n)), np.diag(np.sqrt(y))
    Y = np.diag(y)
    gram = B.T @ Y @ B
    P = B @ np.linalg.pinv(gram, rcond=rtol) @ B.T @ Y
    return P, np.diag(np.sqrt(y))


def residual_response(B, C, y):
    P, Ys = weighted_projector(B, y)
    Q = C - P @ C
    WQ = Ys @ Q
    sv = np.linalg.svd(WQ, compute_uv=False)
    scale = max(1.0, float(sv[0]) if len(sv) else 0.0)
    positive = sv[sv > 1e-9 * scale]
    return Q, float(np.sum(WQ * WQ)), positive


def fit_loss(X, target, y):
    sw = np.sqrt(y)
    Xw, tw = sw[:, None] * X, sw * target
    residual = tw - Xw @ np.linalg.pinv(Xw) @ tw
    return 0.5 * float(residual @ residual)


def exact_increment(X, C, target, y):
    before = fit_loss(X, target, y)
    after = fit_loss(np.column_stack([X, C]), target, y)
    Q, _, _ = residual_response(X, C, y)
    sw = np.sqrt(y)
    Qw, rw = sw[:, None] * Q, sw * target
    predicted = 0.5 * float(rw @ Qw @ np.linalg.pinv(Qw.T @ Qw) @ Qw.T @ rw)
    return before - after, predicted


def make_problem(seed=SEED, n=180, d=3, candidates=18):
    r = np.random.default_rng(seed)
    y = np.exp(r.normal(0, .35, n))
    base = r.normal(size=(n, d))
    # Several candidates are redundant, several are useful, and several are noise.
    blocks = []
    labels = []
    for j in range(candidates):
        if j < 5:
            C = base[:, [j % d]] + .015 * r.normal(size=(n, 1))
            labels.append("redundant")
        elif j in (5, 7, 11):
            C = r.normal(size=(n, 1))
            labels.append("useful")
        else:
            C = .07 * r.normal(size=(n, 1))
            labels.append("weak")
        blocks.append(C)
    target = 1.8 * blocks[5][:, 0] - 1.2 * blocks[7][:, 0] + .25 * base[:, 0] + .35 * r.normal(size=n)
    return y, base, blocks, labels, target


def math_checks():
    r = np.random.default_rng(SEED)
    n = 80
    y = np.exp(r.normal(size=n))
    B = r.normal(size=(n, 3))
    C = r.normal(size=(n, 2))
    P, Ys = weighted_projector(B, y)
    Q, gain, sv = residual_response(B, C, y)
    orth = np.linalg.norm(B.T @ (y[:, None] * Q))
    idem = np.linalg.norm(P @ P - P)

    # Prediction 1: residual gain vanishes as candidate approaches incumbent.
    scales = [0.0, .25, .5, 1.0, 2.0]
    # candidate = incumbent direction + scale * novel direction
    novel = r.normal(size=(n, 1))
    gain_scale = []
    for a in scales:
        _, g, _ = residual_response(B, B[:, [0]] + a * novel, y)
        gain_scale.append(g)
    # Prediction 2: gain scales quadratically with candidate amplitude.
    amps = [.25, .5, 1., 2.]
    amp_gain = [residual_response(B, a * novel, y)[1] for a in amps]
    ratios = [amp_gain[i] / amp_gain[2] for i in range(4)]
    # Prediction 3: exact least-squares improvement equals projected residual energy.
    target = r.normal(size=n)
    observed, predicted = exact_increment(B, C, target, y)
    return {
        "projection_orthogonality_abs": orth,
        "projector_idempotence_abs": idem,
        "scale_sweep": {str(a): g for a, g in zip(scales, gain_scale)},
        "amplitude_sweep_gain": {str(a): g for a, g in zip(amps, amp_gain)},
        "amplitude_ratios_vs_a1": {str(a): ratios[i] for i, a in enumerate(amps)},
        "expected_amplitude_ratios": {str(a): a*a for a in amps},
        "exact_loss_drop": observed,
        "predicted_loss_drop": predicted,
        "relative_prediction_error": abs(observed-predicted) / max(1e-12, abs(observed)),
        "positive_singular_values": sv.tolist(),
    }


def greedy_experiment():
    y, base, blocks, labels, target = make_problem()
    # Baseline: fixed order; idea: largest projected energy, with a sensible
    # lower bound that rejects nearly-null blocks. Both use one scalar block.
    def run(order):
        X = base.copy()
        losses = [fit_loss(X, target, y)]
        selected = []
        for j in order:
            Q, g, sv = residual_response(X, blocks[j], y)
            if len(sv) and sv[0] >= .15 and sv[0] <= 30:
                selected.append(j)
                X = np.column_stack([X, blocks[j]])
                losses.append(fit_loss(X, target, y))
            if len(selected) == 3: break
        return selected, losses
    fixed, fixed_losses = run(list(range(len(blocks))))
    # Greedy recomputes residual gains each round.
    remaining = set(range(len(blocks)))
    X = base.copy(); greedy = [fit_loss(X, target, y)]; accepted = []; gains = []
    for _ in range(3):
        scored = []
        for j in remaining:
            _, g, sv = residual_response(X, blocks[j], y)
            if len(sv) and .15 <= sv[0] <= 30:
                scored.append((g, j, float(sv[0])))
        if not scored: break
        g, j, s = max(scored)
        accepted.append(j); gains.append(g)
        X = np.column_stack([X, blocks[j]]); remaining.remove(j)
        greedy.append(fit_loss(X, target, y))
    return {
        "fixed_order_selected": fixed,
        "fixed_order_labels": [labels[j] for j in fixed],
        "fixed_losses": fixed_losses,
        "greedy_selected": accepted,
        "greedy_labels": [labels[j] for j in accepted],
        "greedy_projected_gains": gains,
        "greedy_losses": greedy,
        "final_loss_ratio_greedy_over_fixed": greedy[-1] / fixed_losses[-1],
    }


def main():
    out = {"seed": SEED, "math_checks": math_checks(), "toy_experiment": greedy_experiment()}
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()