"""Toy verification of conformal early rejection for synthetic diffusion/NAS trials.

The simulator deliberately keeps final evaluation cheap (a latent table lookup), while
counting it as the expensive operation.  The monitor is calibrated only on a split of
completed trajectories and is never used as the authoritative final label.
"""
import json
from pathlib import Path
import numpy as np

SEED = 2762
CHECKPOINTS = (0.75, 0.50, 0.25)
ALPHAS = (0.05, 0.10)


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -40, 40)))


def draw_trajectories(n, rng, beta=5.0, shift=0.0, noise=0.45):
    # x is architecture quality; failures are the hard external evaluation label.
    x = rng.normal(loc=shift, scale=1.0, size=n)
    failure = x < 0.0
    scores = {}
    for j, t in enumerate(CHECKPOINTS):
        # Later denoising checkpoints have more informative hidden states.
        # The same latent x is observed with independent checkpoint noise.
        s = sigmoid(beta * (-x + rng.normal(0, noise * (1.15 - 0.15*j), n)))
        scores[t] = s
    # authoritative accuracy/utility, only available after full evaluation
    accuracy = 0.70 + 0.20 * sigmoid(1.6*x) + rng.normal(0, .004, n)
    feasible = (~failure) & (accuracy >= 0.80)
    return x, failure, feasible, accuracy, scores


def conformal_threshold(values, alpha):
    # Split-conformal order statistic: k=ceil((n+1)(1-alpha)), one-indexed.
    n = len(values)
    k = int(np.ceil((n + 1) * (1-alpha)))
    k = max(1, min(n, k))
    return np.sort(values)[k-1]


def thresholds(cal_scores, alpha):
    return {t: conformal_threshold(cal_scores[t], alpha) for t in CHECKPOINTS}


def run_controller(scores, failure, feasible, accuracy, tau=None, mode='conformal'):
    n = len(failure)
    if mode == 'none':
        rejected = np.zeros(n, dtype=bool)
    elif mode == 'uncalibrated':
        rejected = np.zeros(n, dtype=bool)
        for t in CHECKPOINTS:
            rejected |= scores[t] > 0.5
    else:
        rejected = np.zeros(n, dtype=bool)
        for t in CHECKPOINTS:
            rejected |= scores[t] > tau[t]
    evaluated = ~rejected
    accepted_failure = float(failure[evaluated].mean()) if evaluated.any() else 1.0
    # best authoritative result among evaluated candidates, with infeasible candidates
    # excluded from the utility metric.
    good = evaluated & feasible
    best = float(accuracy[good].max()) if good.any() else float('nan')
    return dict(rejected=int(rejected.sum()), evaluated=int(evaluated.sum()),
                reject_rate=float(rejected.mean()), accepted_failure=accepted_failure,
                best_accuracy=best)


def iid_trial(ncal, ntest, beta, alpha, seed, shift_cal=0., shift_test=0.):
    rng = np.random.default_rng(seed)
    _, fc, _, _, sc = draw_trajectories(ncal, rng, beta=beta, shift=shift_cal)
    _, ft, fe, acc, st = draw_trajectories(ntest, rng, beta=beta, shift=shift_test)
    tau = thresholds(sc, alpha)
    # marginal calibration property: probability that a fresh score exceeds its
    # split-conformal threshold, checked checkpoint-wise.
    exceed = {str(t): float((st[t] > tau[t]).mean()) for t in CHECKPOINTS}
    return tau, exceed, run_controller(st, ft, fe, acc, tau, 'conformal'), run_controller(st, ft, fe, acc, mode='uncalibrated'), run_controller(st, ft, fe, acc, mode='none')


def aggregate(rows):
    keys = ['rejected','evaluated','reject_rate','accepted_failure','best_accuracy']
    return {k: float(np.nanmean([r[k] for r in rows])) for k in keys}


def main():
    out = Path('results.json')
    # Math sanity check: order-statistic exceedance and its finite-sample prediction.
    coverage = []
    for ncal in (20, 50, 100, 500):
        vals = []
        for rep in range(300):
            _, ex, *_ = iid_trial(ncal, 2000, 5.0, .10, 10000+ncal*1000+rep)
            vals.append(np.mean(list(ex.values())))
        observed = float(np.mean(vals))
        # For continuous iid scores, P(fresh > kth order statistic) is approximately
        # (n+1-k)/(n+1), hence <= alpha with this conservative order statistic.
        k = int(np.ceil((ncal+1)*.90))
        predicted = float((ncal-k+1)/(ncal+1))
        coverage.append({'ncal': ncal, 'predicted_exceedance': predicted, 'observed_exceedance': observed})

    # Prediction 1/2: stronger score separability gives lower accepted failure and
    # more useful rejection; prediction 3: exchangeability shift breaks calibration.
    separability = []
    for beta in (0.5, 1., 2., 5., 10.):
        rows = []
        for rep in range(80):
            _, _, c, u, no = iid_trial(400, 1000, beta, .10, 20000+rep+int(beta*100))
            rows.append(c)
        separability.append({'beta': beta, **aggregate(rows)})

    shift = []
    for d in (0., -0.25, -0.5, -1.0):
        rows = []
        for rep in range(100):
            _, _, c, u, no = iid_trial(400, 1000, 5., .10, 30000+rep+int(abs(d)*100), shift_test=d)
            rows.append(c)
        a = aggregate(rows)
        shift.append({'test_shift': d, **a})

    # Fixed proposal stream comparison at two alpha values.
    controllers = {}
    for alpha in ALPHAS:
        allc, allu, alln = [], [], []
        for rep in range(100):
            _, _, c, u, no = iid_trial(400, 1000, 5., alpha, 40000+rep)
            allc.append(c); allu.append(u); alln.append(no)
        controllers[str(alpha)] = {'conformal': aggregate(allc), 'uncalibrated': aggregate(allu), 'none': aggregate(alln)}

    result = {'seed': SEED, 'checkpoints': CHECKPOINTS, 'coverage_sweep': coverage,
              'separability_sweep': separability, 'distribution_shift_sweep': shift,
              'controller_comparison': controllers}
    out.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()