import json
from pathlib import Path
import numpy as np

SEED = 988

def fit_joint(w, y, ridge=1e-6):
    mw, my = w.mean(0), y.mean(0)
    wc, yc = w - mw, y - my
    sww = wc.T @ wc / (len(w) - 1)
    syw = yc.T @ wc / (len(w) - 1)
    syy = yc.T @ yc / (len(w) - 1)
    gain = np.linalg.solve(sww + ridge * np.eye(w.shape[1]), syw.T).T
    cov = syy - gain @ syw.T
    return mw, my, gain, (cov + cov.T) / 2

def condition(model, w):
    mw, my, gain, cov = model
    return my + (w - mw) @ gain.T, cov

def nll(errors, cov, jitter=1e-8):
    c = (cov + cov.T) / 2 + jitter * np.eye(cov.shape[0])
    sign, ld = np.linalg.slogdet(c)
    if sign <= 0:
        return float("inf")
    q = np.mean(np.sum(errors * np.linalg.solve(c, errors.T).T, axis=1))
    return float(.5 * (q + ld + errors.shape[1] * np.log(2 * np.pi)))

def coverage(errors, cov, alpha=.95):
    c = (cov + cov.T) / 2 + 1e-8 * np.eye(len(cov))
    q = np.sum(errors * np.linalg.solve(c, errors.T).T, axis=1)
    # chi-square(10) 95% quantile, avoiding a scipy dependency in the MVP
    threshold = 18.307038053275146 if len(cov) == 10 else np.quantile(q, alpha)
    return float(np.mean(q <= threshold)), float(np.mean(q))

def ar(n, total, rho, sigma, rng):
    x = np.empty((n, total))
    x[:, 0] = rng.normal(size=n) * sigma / np.sqrt(1 - rho*rho)
    for t in range(1, total):
        x[:, t] = rho * x[:, t-1] + sigma * rng.normal(size=n)
    return x

def main():
    rng = np.random.default_rng(SEED)
    out = {}

    # Core algebra: Schur complement of a positive definite joint covariance.
    a = rng.normal(size=(10, 10))
    joint = a @ a.T + .2 * np.eye(10)
    sww, swy = joint[:4, :4], joint[:4, 4:]
    syw, syy = joint[4:, :4], joint[4:, 4:]
    schur = syy - syw @ np.linalg.solve(sww, swy)
    out["math_check"] = {
        "min_joint_eigenvalue": float(np.linalg.eigvalsh(joint).min()),
        "min_schur_eigenvalue": float(np.linalg.eigvalsh(schur).min()),
        "schur_trace": float(np.trace(schur)),
        "psd_verified": bool(np.linalg.eigvalsh(schur).min() > -1e-10),
    }

    # Same process and windows for both predictors.
    rho, sigma, H = .85, 1.0, 10
    ntrain, ntest = 16000, 6000
    train = ar(ntrain, 20 + H, rho, sigma, rng)
    test = ar(ntest, 20 + H, rho, sigma, rng)
    wtr, ytr = train[:, 19:20], train[:, 20:]
    wte, yte = test[:, 19:20], test[:, 20:]

    # Standard rollout: estimate one-step AR coefficient, recursively apply it.
    x0 = wtr[:, 0]
    rho_hat = float((x0 @ ytr[:, 0]) / (x0 @ x0))
    powers = rho_hat ** np.arange(1, H + 1)
    base_mean = wte * powers[None, :]
    # Innovation variance accumulated by recursive rollout; diagonal baseline.
    marginal_var = sigma*sigma * np.array([(1-rho_hat**(2*k))/(1-rho_hat**2) for k in range(1, H+1)])
    base_cov = np.diag(marginal_var)

    # Idea: learn all future coordinates jointly and condition their covariance on w.
    joint_model = fit_joint(wtr, ytr, ridge=1e-6)
    idea_mean, idea_cov = condition(joint_model, wte)
    base_err, idea_err = yte - base_mean, yte - idea_mean
    out["rollout"] = {
        "rho_hat": rho_hat,
        "baseline_recursive": {
            "mse": float(np.mean(base_err**2)),
            "nll": nll(base_err, base_cov),
            "coverage95": coverage(base_err, base_cov)[0],
        },
        "covariance_conditioned": {
            "mse": float(np.mean(idea_err**2)),
            "nll": nll(idea_err, idea_cov),
            "coverage95": coverage(idea_err, idea_cov)[0],
            "conditional_trace": float(np.trace(idea_cov)),
            "min_cov_eigenvalue": float(np.linalg.eigvalsh(idea_cov).min()),
        },
    }

    # Nested-prefix check at a fixed prediction origin. AR theory predicts p=1 helps;
    # older history is redundant, so this explicitly tests the limitation.
    prefix_rows = []
    target = train[:, 20:30]
    target_test = test[:, 20:30]
    for p in [0, 1, 2, 4, 8, 12]:
        if p == 0:
            mean = target.mean(0)
            cov = np.diag(np.var(target, axis=0, ddof=1))
            pred = np.broadcast_to(mean, target_test.shape)
        else:
            model = fit_joint(train[:, 20-p:20], target, ridge=1e-6)
            pred, cov = condition(model, test[:, 20-p:20])
        prefix_rows.append({
            "prefix": p,
            "conditional_trace": float(np.trace(cov)),
            "test_mse": float(np.mean((target_test-pred)**2)),
            "min_eigenvalue": float(np.linalg.eigvalsh(cov).min()),
        })
    out["prefix_sweep"] = prefix_rows
    Path("mvp_results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()