"""Spatial-quantile conformal bands: math check and small operator experiment."""
import json, math, random
from pathlib import Path
import numpy as np


def weighted_quantile(values, weights, mass):
    values = np.asarray(values, float).ravel()
    weights = np.asarray(weights, float).ravel()
    order = np.argsort(values, kind="mergesort")
    v, w = values[order], weights[order]
    c = np.cumsum(w / w.sum())
    return float(v[np.searchsorted(c, mass, side="left")])


def spatial_score(residual, gamma, weights=None):
    """Smallest s whose weighted residual mass is at least 1-gamma."""
    residual = np.asarray(residual)
    if weights is None:
        weights = np.ones(residual.shape[-1]) / residual.shape[-1]
    return np.array([weighted_quantile(x, weights, 1.0 - gamma) for x in residual])


def conformal_q(scores, alpha):
    scores = np.sort(np.asarray(scores, float))
    n = len(scores)
    k = min(n, int(math.ceil((n + 1) * (1.0 - alpha))))
    return float(scores[k - 1])


def bands(center, scale, q):
    return center, q * scale


def math_sanity():
    # Directly checks weighted quantile definition, monotonicity in gamma,
    # and the finite-sample order-statistic indexing including k > n.
    x = np.array([[1., 2., 10., 20.]])
    w = np.array([.1, .2, .3, .4])
    s01 = spatial_score(x, .1, w)[0]
    s50 = spatial_score(x, .5, w)[0]
    assert s01 >= s50 and s01 == 20.0 and s50 == 10.0
    assert conformal_q(np.arange(1., 11.), .1) == 10.0  # ceil(11*.9)=10
    assert conformal_q(np.arange(1., 4.), .01) == 3.0  # clipped k
    # Every returned score covers at least the requested weighted spatial mass.
    rng = np.random.default_rng(4)
    r = rng.lognormal(size=(20, 37))
    for gamma in (.01, .05, .1, .3):
        ss = spatial_score(r, gamma)
        covered = (r <= ss[:, None]).mean(axis=1)
        assert np.all(covered >= 1 - gamma - 1e-12)
    return {"weighted_gamma_0.1": s01, "weighted_gamma_0.5": s50,
            "score_monotone": True, "order_statistic_checks": True}


def make_data(seed=17, n=1800, m=48, d=6):
    rng = np.random.default_rng(seed)
    x = np.linspace(0, 1, m)
    z = rng.normal(size=(n, d))
    # Smooth latent operator with heteroscedastic, localized spatial noise.
    truth = (z[:, 0, None] * np.sin(2*np.pi*x)[None, :] +
             .7*z[:, 1, None] * np.cos(4*np.pi*x)[None, :] +
             .4*z[:, 2, None] * np.sin(6*np.pi*x)[None, :] +
             .25*z[:, 3, None] * x[None, :])
    truth += .08 * z[:, 4, None] * np.cos(10*np.pi*x)[None, :]
    return z, truth, x


def predictor_fit(xtr, ytr, xall, yall, seed=19):
    # A deliberately small frozen operator surrogate: ridge regression on
    # input coefficients and fixed Fourier features. It is a neural-operator-
    # style field map with shared input representation, without heavy training.
    rng = np.random.default_rng(seed)
    d = xtr.shape[1]
    m = ytr.shape[1]
    phi = np.concatenate([xtr, np.tanh(xtr), np.ones((len(xtr), 1))], axis=1)
    # tiny random feature layer makes this nonlinear while remaining reproducible
    W = rng.normal(0, .55, size=(phi.shape[1], 18))
    H = np.tanh(phi @ W)
    H = np.concatenate([phi, H], axis=1)
    lam = 1e-2
    coef = np.linalg.solve(H.T @ H + lam*np.eye(H.shape[1]), H.T @ ytr)
    def predict(xx):
        p = np.concatenate([xx, np.tanh(xx), np.ones((len(xx), 1))], axis=1)
        hh = np.tanh(p @ W)
        hh = np.concatenate([p, hh], axis=1)
        return hh @ coef
    return predict


def run_experiment():
    z, y, grid = make_data()
    # Fixed disjoint train/calibration/test split.
    pred = predictor_fit(z[:900], y[:900], z, y)
    cal = slice(900, 1100)
    test = slice(1100, 1800)
    yc, yt = y[cal], y[test]
    pc, pt = pred(z[cal]), pred(z[test])
    # Scale is learned only from the training residual envelope, then frozen.
    train_res = np.abs(y[:900] - pred(z[:900]))
    scale = np.maximum(np.quantile(train_res, .75, axis=0), .035)
    rc = np.abs(yc-pc) / scale[None, :]
    rt = np.abs(yt-pt) / scale[None, :]
    out = {"n_cal": len(yc), "n_test": len(yt), "grid": len(grid), "scale_mean": float(scale.mean())}
    for gamma in (.01, .05, .10):
        qscores = spatial_score(rc, gamma)
        q = conformal_q(qscores, .10)
        maxq = conformal_q(rc.max(axis=1), .10)
        frac_q = (rt <= q).mean(axis=1)
        frac_max = (rt <= maxq).mean(axis=1)
        target = 1-gamma
        out[f"gamma_{gamma:.2f}"] = {
            "quantile_q": q, "max_q": maxq,
            "quantile_mean_width": float(2*q*scale.mean()),
            "max_mean_width": float(2*maxq*scale.mean()),
            "quantile_mean_fraction": float(frac_q.mean()),
            "max_mean_fraction": float(frac_max.mean()),
            "quantile_event_rate": float((frac_q >= target).mean()),
            "max_event_rate": float((frac_max >= target).mean()),
            "target_event_rate": .90,
        }
    return out


def main():
    result = {"math_sanity": math_sanity(), "experiment": run_experiment()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()