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

SEED = 3133
rng = np.random.default_rng(SEED)
D = 8
N = 64
K = 120

# Eight separated wells in an 8-D neural-energy-like landscape.
centers = np.zeros((8, D))
for j in range(8):
    centers[j, j] = 2.4
    centers[j, (j + 1) % D] = -1.4
sigma = 0.72
log2pi = math.log(2 * math.pi)

def logp0(x):
    return -0.5 * np.sum(x*x, axis=-1) - D*0.5*log2pi

def score(x):
    # log likelihood of a mixture of narrow wells; max/logsumexp is stable.
    z = -0.5*np.sum((x[:, None, :] - centers[None, :, :])**2, axis=-1)/sigma**2
    z -= D*math.log(sigma) + D*0.5*log2pi + math.log(8.)
    m = z.max(axis=1)
    return m + np.log(np.exp(z-m[:,None]).sum(axis=1))

def log_gaussian(x, mean, cov):
    L = np.linalg.cholesky(cov)
    y = np.linalg.solve(L, (x-mean).T).T
    return -0.5*np.sum(y*y, axis=1) - np.log(np.diag(L)).sum() - D*0.5*log2pi

def systematic_resample(w, rr):
    w = w / w.sum()
    c = np.cumsum(w)
    return int(np.searchsorted(c, rr))

def importance_sanity():
    # Deliberately biased proposal. Weighted moments should recover constrained-prior moments.
    M = 180000
    qmean = np.ones(D) * 1.3
    qcov = np.eye(D) * 2.0
    x = rng.multivariate_normal(qmean, qcov, size=M)
    s = score(x)
    lam = np.quantile(s, .88)
    lpq = log_gaussian(x, qmean, qcov)
    lp = logp0(x)
    ok = s > lam
    lw = lp[ok] - lpq[ok]
    lw -= lw.max()
    w = np.exp(lw)
    ess = w.sum()**2 / np.sum(w*w)
    # Independent prior sample gives a reference constrained mean (not used in weighting).
    y = rng.normal(size=(500000, D)); sy = score(y); yy = y[sy > lam]
    weighted_mean = (w[:,None]*x[ok]).sum(0)/w.sum()
    ref_mean = yy.mean(0)
    mean_err = float(np.linalg.norm(weighted_mean-ref_mean)/math.sqrt(D))
    return dict(threshold=float(lam), ess=float(ess), ess_fraction=float(ess/max(1,ok.sum())),
                weighted_reference_rmse=mean_err, accepted=int(ok.sum()))

def beta_sanity():
    # Exact claimed shrinkage law, with a sizeable independent replication.
    u = rng.random((200000, N))
    logs = np.log(u.max(axis=1))
    return dict(empirical_mean=float(logs.mean()), predicted=-1/N,
                relative_error=float(abs(logs.mean()+1/N)/(1/N)),
                empirical_sd=float(logs.std()), predicted_sd=float(1/N))

def fit_flow(live):
    # Conditional affine flow: ML Gaussian location/scale fitted on recent live set.
    # Shrinkage/regularization prevents singular densities in a tiny online window.
    mu = live.mean(0)
    var = live.var(0) + 0.20**2
    return mu, var

def run_nested(method):
    local = np.random.default_rng(SEED + (0 if method=='baseline' else 1))
    live = local.normal(size=(N,D)); ls = score(live)
    evals = 0; accepted = 0; ess_values=[]; acc_values=[]; trajectories=[]
    logX = 0.0
    for it in range(K):
        worst = int(np.argmin(ls)); lam = ls[worst]
        # For baseline, exact prior rejection. For idea, batch proposals from fitted affine flow.
        if method == 'baseline':
            ntry = 0
            while True:
                x = local.normal(size=(1,D)); sx = score(x)[0]; ntry += 1
                if sx > lam: break
            evals += ntry; acc_values.append(1.0/ntry); ess_values.append(1.0)
        else:
            mu, var = fit_flow(live)
            # A modest batch keeps the experiment small while allowing importance resampling.
            M = 128
            x = mu + local.normal(size=(M,D))*np.sqrt(var)
            sx = score(x); evals += M
            ok = sx > lam
            if not np.any(ok):
                # Robust fallback: exact prior rejection for this replacement.
                ntry=0
                while True:
                    xx=local.normal(size=(1,D)); ss=score(xx)[0]; ntry+=1
                    if ss>lam: x=xx; sx=np.array([ss]); break
                evals += ntry; accepted += 1; live[worst]=x[0]; ls[worst]=sx[0]
                ess_values.append(0.0); acc_values.append(0.0); logX += -1/N; trajectories.append(logX); continue
            lpq = log_gaussian(x[ok], mu, np.diag(var))
            lp = logp0(x[ok])
            lw = lp-lpq; lw -= lw.max(); w=np.exp(lw)
            ess_values.append(float(w.sum()**2/np.sum(w*w)))
            acc_values.append(float(ok.mean()))
            j = systematic_resample(w, local.random())
            xnew=x[ok][j]; snew=sx[ok][j]
        live[worst]=xnew if method!='baseline' else x[0]
        ls[worst]=snew if method!='baseline' else sx
        accepted += 1
        logX += -1/N # ideal nested trajectory comparator
        trajectories.append(logX)
    return dict(score_evals=evals, accepted=accepted,
                evals_per_accept=evals/accepted, mean_candidate_acceptance=float(np.mean(acc_values)),
                mean_ess=float(np.mean(ess_values)), final_logX=float(logX),
                trajectory=trajectories)

def main():
    t=time.time(); out={'seed':SEED, 'dimension':D, 'live_points':N, 'iterations':K,
        'math_checks':{'beta_shrinkage':beta_sanity(), 'importance_sampling':importance_sanity()}}
    out['baseline']=run_nested('baseline'); out['idea']=run_nested('idea')
    # Avoid bloating the report while retaining a trajectory snapshot for inspection.
    out['baseline']['trajectory_sample']=out['baseline'].pop('trajectory')[::20]
    out['idea']['trajectory_sample']=out['idea'].pop('trajectory')[::20]
    out['runtime_sec']=time.time()-t
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__=='__main__': main()