import json
import numpy as np
from scipy.optimize import minimize

SEED = 1071
K = 3
EPS = 1e-9


def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def pushforward(response_prob, phi, n_states):
    out = np.zeros(n_states)
    for r, j in enumerate(phi):
        out[j] += response_prob[r]
    return out


def tv(p, q):
    return 0.5 * np.abs(p - q).sum(axis=-1)


def make_data(n, rng):
    # A range of exact Bayesian posteriors, represented by logits.
    z = rng.normal(0, 1.25, size=(n, K))
    q = softmax(z)
    y = np.array([rng.choice(K, p=row) for row in q])
    return q, y


def raw_from_q(q, gamma):
    return q ** gamma / (q ** gamma).sum(axis=1, keepdims=True)


def fit_calibrator(p, y, ridge=1e-3):
    # log pi = a_j log(p_j) + b_j, fitted by held-out-style train labels.
    lp = np.log(np.clip(p, EPS, 1))
    def unpack(x):
        return x[:K], x[K:]
    def objective(x):
        a, b = unpack(x)
        logits = lp * a[None, :] + b[None, :]
        logz = np.logaddexp.reduce(logits, axis=1)
        nll = -(logits[np.arange(len(y)), y] - logz).mean()
        return nll + ridge * (a @ a + b @ b)
    res = minimize(objective, np.r_[np.ones(K), np.zeros(K)], method='BFGS',
                   options={'maxiter': 500, 'gtol': 1e-8})
    a, b = unpack(res.x)
    return a, b, bool(res.success)


def apply_calibrator(p, a, b):
    return softmax(np.log(np.clip(p, EPS, 1)) * a[None, :] + b[None, :])


def metrics(pred, q, y):
    n = len(y)
    nll = -np.log(np.clip(pred[np.arange(n), y], EPS, 1)).mean()
    brier = ((pred - np.eye(K)[y]) ** 2).sum(axis=1).mean()
    tv_mean = tv(pred, q).mean()
    conf = pred.max(axis=1)
    acc = (pred.argmax(axis=1) == y).astype(float)
    bins = np.linspace(0, 1, 11)
    ece = 0.
    for lo, hi in zip(bins[:-1], bins[1:]):
        ix = (conf >= lo) & ((conf < hi) if hi < 1 else (conf <= hi))
        if ix.any():
            ece += ix.mean() * abs(acc[ix].mean() - conf[ix].mean())
    # smallest top-probability set with nominal 90% mass, empirical coverage
    order = np.argsort(-pred, axis=1)
    mass = np.take_along_axis(pred, order, axis=1).cumsum(1)
    size = (mass < .9).sum(1) + 1
    covered = np.array([y[i] in order[i, :size[i]] for i in range(n)]).mean()
    return {'nll': float(nll), 'brier': float(brier), 'tv_to_exact_q': float(tv_mean),
            'ece': float(ece), 'set90_coverage': float(covered),
            'mean_set_size': float(size.mean())}


def run_gamma(gamma, rng):
    qtr, ytr = make_data(5000, rng)
    qte, yte = make_data(5000, rng)
    ptr, pte = raw_from_q(qtr, gamma), raw_from_q(qte, gamma)
    a, b, ok = fit_calibrator(ptr, ytr)
    cal = apply_calibrator(pte, a, b)
    return {
        'gamma': gamma,
        'predicted_a': 1.0 / gamma,
        'fitted_a_mean': float(a.mean()),
        'gamma_times_fitted_a': float(gamma * a.mean()),
        'raw': metrics(pte, qte, yte),
        'calibrated': metrics(cal, qte, yte),
        'optimizer_success': ok,
    }


def main():
    rng = np.random.default_rng(SEED)
    # Core pushforward sanity check: response probability is conserved by phi.
    # Each state has several synonymous responses; split mass arbitrarily.
    phi = np.array([0, 0, 0, 1, 1, 1, 2, 2, 3])  # state 3 = abstain/malformed
    response = rng.dirichlet(np.ones(len(phi)))
    pushed = pushforward(response, phi, 4)
    direct = np.zeros(4)
    for j in range(4): direct[j] = response[phi == j].sum()
    pushforward_error = float(np.max(np.abs(pushed - direct)))
    # Invariance test: redistribute each state's mass among additional synonyms.
    state_mass = pushed.copy()
    split = np.concatenate([rng.dirichlet(np.ones(5)) * state_mass[j] for j in range(4)])
    phi_split = np.repeat(np.arange(4), 5)
    invariance_error = float(np.max(np.abs(pushforward(split, phi_split, 4) - state_mass)))

    rows = [run_gamma(g, np.random.default_rng(SEED + int(g * 100)))
            for g in [0.5, 0.75, 1.0, 1.5, 2.0]]
    result = {
        'seed': SEED,
        'predictions': {
            'P1_pushforward_exact': 'aggregation error is 0 up to floating point, and synonym splitting leaves state mass unchanged',
            'P2_power_distortion': 'raw log-odds scale with gamma; fitted calibration slope a should satisfy gamma*a ~= 1',
            'P3_calibration_effect': 'calibrated TV-to-exact-posterior should be near 0 for every gamma, while raw TV is 0 at gamma=1 and grows away from 1'
        },
        'sanity': {'pushforward_max_abs_error': pushforward_error,
                   'synonym_split_max_abs_error': invariance_error,
                   'response_mass': float(response.sum()),
                   'pushed_mass': float(pushed.sum())},
        'sweep': rows,
        'comparison_gamma_2': next(r for r in rows if r['gamma'] == 2.0)
    }
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()