import json
import math
from pathlib import Path
import numpy as np

SEED = 2747


def empirical_cvar_lower(z, gamma):
    """Max over eta of eta - (gamma*K)^-1 sum(max(eta-z,0))."""
    z = np.asarray(z, dtype=float)
    K = len(z)
    if not (0 < gamma <= 1) or K == 0:
        raise ValueError("gamma and samples")
    # The objective is piecewise linear; a maximizer is an order statistic.
    candidates = np.r_[z, z.min() - 1.0, z.max() + 1.0]
    vals = np.array([eta - np.maximum(eta - z, 0).sum() / (gamma * K)
                     for eta in candidates])
    # Return the optimized objective, not eta itself.
    return float(vals.max())


def empirical_bottom_cvar(z, gamma):
    """Equivalent weighted mean of the lower empirical tail (fraction gamma)."""
    z = np.sort(np.asarray(z, dtype=float))
    K = len(z)
    h = gamma * K
    m = int(math.floor(h))
    frac = h - m
    # Integral of the empirical quantile over [0,gamma], divided by gamma.
    total = z[:m].sum()
    if frac > 1e-12 and m < K:
        total += frac * z[m]
    return float(total / h)


def conformal_q(residuals, alpha):
    r = np.sort(np.asarray(residuals, dtype=float))
    k = int(math.ceil((len(r) + 1) * (1 - alpha)))
    if k > len(r):
        return float("inf")
    return float(r[k - 1])


def episodes(rng, n, K=16, gamma=0.25, shift=0.0):
    """Selected-trajectory episodes: latent true clearance and noisy biased samples."""
    # Heteroskedastic scenes make the calibration nontrivial while retaining exchangeability.
    y = rng.normal(0.72 + shift, 0.20, n)
    y = np.clip(y, -0.8, 1.5)
    bias = 0.12 + 0.12 * np.maximum(0, 0.5 - y)  # predictor overstates danger clearance
    # Candidate/stochastic predictions for the already-selected trajectory.
    z = y[:, None] + bias[:, None] + rng.normal(0, 0.13, (n, K))
    # Vectorized empirical lower-tail mean, equal to the eta maximization.
    zs = np.sort(z, axis=1)
    h = gamma * K
    mfull = int(math.floor(h))
    frac = h - mfull
    if mfull == 0:
        m = zs[:, 0].copy()
    else:
        m = zs[:, :mfull].sum(axis=1)
        if frac > 1e-12 and mfull < K:
            m = m + frac * zs[:, mfull]
        m = m / h
    nominal = z.mean(axis=1)
    return y, m, nominal


def one_trial(rng, ncal, ntest, alpha, gamma, shift_test=0.0):
    yc, mc, nc = episodes(rng, ncal, gamma=gamma)
    q = conformal_q(mc - yc, alpha)
    qt = conformal_q(nc - yc, alpha)
    yt, mt, nt = episodes(rng, ntest, gamma=gamma, shift=shift_test)
    cert = mt - q
    cert_nom = nt - qt
    # Uncalibrated is C=M (the common baseline).
    return {
        "q_cvar": q, "q_nominal": qt,
        "coverage_cvar": float(np.mean(yt >= cert)),
        "coverage_nominal_cal": float(np.mean(yt >= cert_nom)),
        "coverage_uncalibrated": float(np.mean(yt >= mt)),
        "nonnegative_cvar": float(np.mean(cert >= 0)),
        "nonnegative_nominal": float(np.mean(cert_nom >= 0)),
    }


def mean_trials(ncal, alpha, gamma, trials=30, ntest=700, shift=0.0):
    out = []
    for t in range(trials):
        out.append(one_trial(np.random.default_rng(SEED + 10000*t + ncal), ncal,
                             ntest, alpha, gamma, shift))
    keys = out[0].keys()
    return {k: float(np.mean([x[k] for x in out])) for k in keys}


def main():
    rng = np.random.default_rng(SEED)
    # Math sanity: CVaR optimizer versus direct lower-tail integral.
    cvar_rows = []
    for K in [8, 16, 64, 127]:
        z = rng.normal(size=K)
        for gamma in [0.1, 0.25, 0.5, 1.0]:
            a = empirical_cvar_lower(z, gamma)
            b = empirical_bottom_cvar(z, gamma)
            cvar_rows.append({"K": K, "gamma": gamma, "optimizer": a,
                              "lower_tail_mean": b, "abs_error": abs(a-b)})

    # Prediction 1: exact finite-sample conformal rank coverage is approximately
    # 1-alpha (and never systematically below it) under exchangeability.
    coverage_sweep = {}
    for alpha in [0.05, 0.10, 0.20, 0.30]:
        coverage_sweep[str(alpha)] = mean_trials(200, alpha, 0.25, trials=80)

    # Prediction 2: increasing calibration n makes the quantile/certificate more
    # stable; the expected rank granularity is 1/(n+1).
    n_sweep = {}
    for n in [20, 50, 100, 200, 500]:
        n_sweep[str(n)] = mean_trials(n, 0.10, 0.25, trials=80)

    # Prediction 3: the lower-tail mechanism is conservative relative to mean
    # prediction when prediction noise is nonzero, increasing the required q.
    gamma_sweep = {}
    for gamma in [0.10, 0.25, 0.50, 1.0]:
        gamma_sweep[str(gamma)] = mean_trials(200, 0.10, gamma, trials=80)

    # Deliberate distribution shift: exchangeability assumption fails.
    shifted = mean_trials(200, 0.10, 0.25, trials=30, shift=-0.45)
    result = {
        "seed": SEED,
        "cvar_math_max_abs_error": max(x["abs_error"] for x in cvar_rows),
        "cvar_math_rows": cvar_rows,
        "coverage_sweep": coverage_sweep,
        "calibration_size_sweep": n_sweep,
        "gamma_sweep": gamma_sweep,
        "shifted_test": shifted,
        "predictions": {
            "exchangeable_coverage": "coverage should be >= 1-alpha in aggregate; finite n gives rank granularity about 1/(n+1)",
            "alpha_monotonicity": "larger alpha lowers q and target coverage toward 1-alpha",
            "lower_tail_conservatism": "smaller gamma lowers M_CVaR and should require a smaller/equal residual correction than gamma=1"
        }
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps({"cvar_max_error": result["cvar_math_max_abs_error"],
                      "coverage": result["coverage_sweep"],
                      "n": result["calibration_size_sweep"],
                      "gamma": result["gamma_sweep"],
                      "shifted": shifted}, indent=2))


if __name__ == "__main__":
    main()