import json, time
import numpy as np


def huber(r, delta):
    r = np.asarray(r)
    a = np.abs(r)
    return np.where(a <= delta, 0.5*r*r, delta*(a - 0.5*delta))


def huber_barycenter(Q, weights=None, delta=1.0, iterations=12, sort_output=True):
    """Safeguarded pointwise displacement-Huber barycenter for [objects, quantiles]."""
    Q = np.asarray(Q, dtype=float)
    n, m = Q.shape
    if weights is None:
        weights = np.ones(n) / n
    weights = np.asarray(weights, dtype=float)
    weights = weights / weights.sum()
    q = np.sum(weights[:, None] * Q, axis=0)
    lo, hi = Q.min(axis=0), Q.max(axis=0)
    # Newton is used where local curvature exists; otherwise bisection is safe.
    for _ in range(iterations):
        r = q[None, :] - Q
        score = np.sum(weights[:, None] * np.clip(r, -delta, delta), axis=0)
        active = (np.abs(r) < delta).astype(float)
        curvature = np.sum(weights[:, None] * active, axis=0)
        # Maintain a score bracket: score is monotone increasing in q.
        hi = np.where(score > 0, q, hi)
        lo = np.where(score < 0, q, lo)
        newton = q - score / np.maximum(curvature, 1e-12)
        bisect = 0.5 * (lo + hi)
        # Also reject Newton steps outside the valid bracket.
        q = np.where((curvature > 1e-10) & (newton >= lo) & (newton <= hi), newton, bisect)
    return np.sort(q) if sort_output else q


def objective(q, Q, weights, delta):
    return np.sum(weights[:, None] * huber(q[None, :] - Q, delta))


def mechanism_checks():
    delta = 1.0
    # Prediction 1: for [0,0,T]/3, q=T/3 until T=1.5, then q=delta/2.
    Ts = np.linspace(0, 8, 33)
    observed, predicted = [], []
    for T in Ts:
        Q = np.array([[0.0], [0.0], [T]])
        observed.append(huber_barycenter(Q, delta=delta, iterations=20, sort_output=False)[0])
        predicted.append(min(T/3.0, delta/2.0))
    observed, predicted = np.array(observed), np.array(predicted)
    transition_T = Ts[np.where(np.abs(observed - delta/2) < 2e-4)[0][0]]
    max_err = float(np.max(np.abs(observed-predicted)))

    # Prediction 2: inliers within delta remain exactly quadratic/arithmetic.
    shifts = np.linspace(-0.9, 0.9, 19)
    small_err = max(abs(huber_barycenter(np.array([[-s], [0.0], [s]]), delta=delta,
                                         iterations=12, sort_output=False)[0]) for s in shifts)

    # Prediction 3: one remote token's score contribution is capped at delta/n.
    scores = []
    for T in [2, 10, 100, 1000]:
        Q = np.array([[0.0], [0.0], [T]])
        q = huber_barycenter(Q, delta=delta, iterations=20, sort_output=False)[0]
        scores.append(abs(np.clip(q-T, -delta, delta)/3.0))
    return {
        'predicted_transition_T': 1.5, 'observed_transition_T_grid': float(transition_T),
        'transition_grid_step': 0.25, 'max_saturation_curve_error': max_err,
        'predicted_small_shift_error': 0.0, 'observed_small_shift_max_error': float(small_err),
        'predicted_outlier_score_limit': 1/3, 'observed_outlier_score_max': float(max(scores)),
        'checks_pass': bool(max_err < 0.03 and small_err < 1e-6 and max(scores) <= 1/3+1e-8)
    }


def classification_experiment(seed=7):
    rng = np.random.default_rng(seed)
    m, nobj, ntrain, ntest = 32, 12, 1200, 800
    shape = np.linspace(-0.7, 0.7, m)
    def make(count, corruption_rate):
        y = rng.integers(0, 2, count)
        center = (2*y-1)*0.85
        Q = center[:, None, None] + shape[None, None, :] + rng.normal(0, .16, (count,nobj,m))
        bad = rng.random((count,nobj)) < corruption_rate
        Q += bad[:, :, None] * rng.choice([-1, 1], (count,nobj,1)) * 5.0
        return Q, y
    trainQ, trainy = make(ntrain, 0.0)
    test_clean, testy = make(ntest, 0.0)
    test_bad, _ = make(ntest, 0.5)
    def pool(A, mode):
        out=[]; t=time.perf_counter()
        for x in A:
            out.append(x.mean(axis=0) if mode == 'mean' else huber_barycenter(x, delta=1.0, iterations=8))
        return np.asarray(out), time.perf_counter()-t
    results={}
    for mode in ['mean','huber']:
        Xtr,ttr=pool(trainQ,mode); Xc,tc=pool(test_clean,mode); Xb,tb=pool(test_bad,mode)
        cent=np.stack([Xtr[trainy==k].mean(0) for k in [0,1]])
        def acc(X,y): return float(np.mean(np.argmin(((X[:,None,:]-cent[None,:,:])**2).mean(2),1)==y))
        ac,ab=acc(Xc,testy),acc(Xb,testy)
        results[mode]={'clean_accuracy':ac,'corrupt_accuracy':ab,'degradation':ac-ab,
                       'pool_ms_per_sample_clean':1000*tc/ntest,'pool_ms_per_sample_train':1000*ttr/ntrain}
    results['corruption_rate']=0.5
    results['huber_over_mean_cost_ratio']=results['huber']['pool_ms_per_sample_clean']/results['mean']['pool_ms_per_sample_clean']
    return results


if __name__ == '__main__':
    report={'mechanism_checks':mechanism_checks(),'classification':classification_experiment()}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))