import json
import numpy as np
from scipy.special import ndtr

SEED = 2816
BETA = 0.99


def cvar_rockafellar(y, beta=BETA):
    y = np.asarray(y)
    eta = np.quantile(y, beta, method='linear')
    return float(eta + np.mean(np.maximum(y - eta, 0.0)) / (1-beta))


def ridge_fit(P, z, lam):
    P = np.asarray(P); z = np.asarray(z)
    return np.linalg.solve(P.T @ P + lam*np.eye(P.shape[1]), P.T @ z)


def phi(x):
    x = np.asarray(x)
    bump = np.exp(-0.5*((x-0.90)/0.025)**2)
    return np.column_stack([np.ones_like(x), x, bump, np.sin(6*np.pi*x)])


def low_members(x):
    # Five cheap models: close globally, but uncertain near the tail boundary.
    x = np.asarray(x)
    base = 0.15 + 0.35*x + 0.08*np.sin(4*np.pi*x)
    bump = np.exp(-0.5*((x-0.90)/0.025)**2)
    offsets = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
    return np.array([base + 0.025*o*(0.25 + 2.0*bump) for o in offsets])


def high_oracle(x):
    x = np.asarray(x)
    # Narrow, high-impact failure region missing from the cheap model.
    return 0.15 + 0.35*x + 0.08*np.sin(4*np.pi*x) + 1.25*np.exp(-0.5*((x-0.90)/0.025)**2) + 0.04*np.sin(18*np.pi*x)


def residual_uncertainty(x, Xh, zh, lam):
    if len(Xh) < 3: return np.ones(len(x))
    P = phi(Xh); w = ridge_fit(P, zh, lam)
    pred = P @ w
    rng = np.random.default_rng(123)
    ws = []
    for _ in range(20):
        idx = rng.integers(0, len(Xh), len(Xh))
        ws.append(ridge_fit(P[idx], zh[idx], lam))
    return np.std(np.asarray(ws) @ phi(x).T, axis=0)


def acquire(policy, cand, Xh, zh, n, eps=1.0, rho=0.5, lam=1e-3):
    members = low_members(cand); mu = members.mean(0); s = members.std(0, ddof=1)
    q = np.quantile(mu, BETA)
    region = np.abs(mu-q) <= eps*np.maximum(s, 1e-5)
    if policy == 'random':
        score = np.random.default_rng(999+len(Xh)).random(len(cand))
        return cand[np.argsort(score)[-n:]]
    if policy == 'global_uncertainty':
        score = s
    else:
        # stage 1: half from tail boundary by uncertainty; stage 2: residual/correction impact
        ids = np.where(region)[0]
        if len(ids) == 0: ids = np.arange(len(cand))
        n1 = n//2
        take1 = ids[np.argsort(s[ids])[-min(n1,len(ids)):]]
        sr = residual_uncertainty(cand, Xh, zh, lam)
        if len(Xh): w = ridge_fit(phi(Xh), zh, lam); impact = np.abs(phi(cand) @ w)
        else: impact = np.zeros(len(cand))
        score = (sr/(sr.std()+1e-9) + rho*impact/(impact.std()+1e-9))
        available = np.ones(len(cand), dtype=bool); available[take1] = False
        ids2 = np.where(available)[0]
        take2 = ids2[np.argsort(score[ids2])[-(n-n1):]]
        return np.concatenate([cand[take1], cand[take2]])
    return cand[np.argsort(score)[-n:]]


def mechanism_checks():
    # Prediction 1: Rockafellar minimizer is the beta quantile.
    y = np.sort(np.random.default_rng(SEED).normal(size=10001))
    etas = np.linspace(y[int(.985*len(y))], y[int(.995*len(y))], 401)
    obj = [e + np.mean(np.maximum(y-e,0))/(1-BETA) for e in etas]
    eta_obs = etas[np.argmin(obj)]; q = np.quantile(y, BETA)
    cvar_direct = y[int(np.ceil(BETA*len(y))):].mean()
    # Prediction 2: in a single orthogonal direction, ridge correction ratio is 1/(1+lambda).
    z = np.array([2.0]); P = np.array([[1.0]])
    lams = np.array([0., .01, .1, 1., 10.])
    ratios = np.array([ridge_fit(P,z,l)[0]/2.0 for l in lams])
    pred_ratios = 1/(1+lams)
    # Prediction 3: tail-boundary region probability for standard normal standardized distance is 2Phi(eps)-1.
    # Generate standardized distances directly, matching the definition |mu-q|/s.
    d = np.random.default_rng(SEED+1).normal(size=200000)
    eps = np.array([0.5, 1., 2.])
    observed = np.array([(np.abs(d)<=e).mean() for e in eps])
    predicted = 2*ndtr(eps)-1
    return {
        'cvar_quantile_check': {'predicted_eta_quantile':float(q), 'observed_eta':float(eta_obs), 'abs_error':float(abs(q-eta_obs)), 'direct_cvar':float(cvar_direct), 'rockafellar_at_eta':float(min(obj))},
        'ridge_shrinkage_check': {'lambda':lams.tolist(), 'predicted_ratio':pred_ratios.tolist(), 'observed_ratio':ratios.tolist(), 'max_abs_error':float(np.max(abs(ratios-pred_ratios)))},
        'tail_region_check': {'epsilon':eps.tolist(), 'predicted_fraction':predicted.tolist(), 'observed_fraction':observed.tolist(), 'max_abs_error':float(np.max(abs(observed-predicted)))}
    }


def mini_experiment():
    rng = np.random.default_rng(SEED)
    pool = np.linspace(0,1,5000,endpoint=False)+0.5/5000
    truth = high_oracle(pool); true_cvar = cvar_rockafellar(truth)
    results = {}
    for policy in ['random','global_uncertainty','tail_residual']:
        # same initial design across policies
        init = rng.choice(len(pool), 24, replace=False)
        Xh = pool[init].copy(); zh = high_oracle(Xh)-low_members(Xh).mean(0)
        rows=[]
        for rnd in range(5):
            # corrected predictor uses low mean plus fitted residual
            w = ridge_fit(phi(Xh), zh, 1e-3)
            pred = low_members(pool).mean(0) + phi(pool)@w
            rows.append({'calls':int(len(Xh)), 'cvar_abs_error':float(abs(cvar_rockafellar(pred)-true_cvar)), 'global_mse':float(np.mean((pred-truth)**2)), 'tail_mean_error':float(abs(np.mean(np.sort(pred)[-50:])-np.mean(np.sort(truth)[-50:])))})
            cand = pool[rng.choice(len(pool), 2500, replace=False)]
            new = acquire(policy, cand, Xh, zh, 16)
            Xh = np.concatenate([Xh,new]); zh = np.concatenate([zh, high_oracle(new)-low_members(new).mean(0)])
        results[policy] = rows
    return {'true_cvar':true_cvar, 'runs':results}


if __name__ == '__main__':
    out = {'seed':SEED, 'beta':BETA, 'mechanism_checks':mechanism_checks(), 'mini_experiment':mini_experiment()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out, indent=2))