import json
import numpy as np


def fit_jacobian(z0, r0, Z, R, eta=1e-5, h=1.0):
    X = np.asarray(Z) - z0[None, :]
    Y = np.asarray(R) - r0[None, :]
    w = np.exp(-np.sum(X * X, axis=1) / max(h * h, 1e-12))
    H = X.T @ (w[:, None] * X) + eta * np.eye(X.shape[1])
    B = X.T @ (w[:, None] * Y)
    return np.linalg.solve(H, B).T, w


def geometry_steps(J, r0, lam, rho, rng):
    p = J.shape[1]
    A = J.T @ J + lam * np.eye(p)
    dgn = -np.linalg.solve(A, J.T @ r0)
    if np.linalg.norm(dgn) > rho:
        dgn *= rho / np.linalg.norm(dgn)
    ev, Q = np.linalg.eigh(A)
    u = rng.normal(size=p)
    de = Q @ ((Q.T @ u) / np.sqrt(np.maximum(ev, 1e-12)))
    de *= rho / max(np.linalg.norm(de), 1e-12)
    return dgn, de, np.linalg.cond(A)


def verification(seed=7):
    rng = np.random.default_rng(seed)
    p, m, n = 4, 6, 80
    Jtrue = rng.normal(size=(m, p))
    z0, r0 = rng.normal(size=p), rng.normal(size=m)
    X = rng.normal(scale=.18, size=(n, p))
    Jhat, _ = fit_jacobian(z0, r0, z0 + X, r0 + X @ Jtrue.T, eta=1e-10, h=10.)
    rel_err = np.linalg.norm(Jhat - Jtrue) / np.linalg.norm(Jtrue)
    d, _, cond = geometry_steps(Jtrue, r0, .03, 10., rng)
    ev = np.linalg.eigvalsh(Jtrue.T @ Jtrue + .03 * np.eye(p))
    inv_ev = 1.0 / ev
    alignment = bool(np.all(np.diff(ev) >= -1e-10) and np.all(np.diff(inv_ev) <= 1e-10))
    before, after = np.linalg.norm(r0), np.linalg.norm(r0 + Jtrue @ d)
    return {'jacobian_relative_error': float(rel_err),
            'residual_norm_before': float(before), 'residual_norm_after': float(after),
            'residual_ratio': float(after / before), 'metric_condition': float(cond),
            'metric_alignment_eigenvalue_check': alignment,
            'sensitivity_eigenvalues': ev.tolist(), 'inverse_metric_eigenvalues': inv_ev.tolist()}


def residual(z):
    target = np.array([1.2, -1., .7, -.5, .35, -.25, .15, -.1])
    scales = np.array([.12, .15, .25, .35, .5, .7, 1., 1.3])
    q = (z - target) / scales
    r = q.copy()
    r[0] += .22 * z[1] ** 2
    r[2] += .18 * np.sin(z[3])
    r[5] += .12 * z[0] * z[4]
    return r


def run(seed, guided, eval_budget=500, p=8, pop=20):
    rng = np.random.default_rng(seed)
    Z = rng.uniform(-3., 3., (pop, p)); R = np.array([residual(z) for z in Z])
    trace_z, trace_r, curve = list(Z), list(R), []
    while len(trace_z) < eval_budget:
        ib = np.argmin(np.sum(R * R, axis=1)); z0, r0 = Z[ib].copy(), R[ib].copy()
        nnew = min(pop, eval_budget - len(trace_z)); candidates = []
        for k in range(nnew):
            use = guided and k < max(1, nnew // 4) and len(trace_z) >= 2 * p
            if use:
                idx = np.argsort([np.linalg.norm(x-z0) for x in trace_z])[:min(len(trace_z), 80)]
                J, _ = fit_jacobian(z0, r0, np.array([trace_z[i] for i in idx]), np.array([trace_r[i] for i in idx]), eta=2e-3, h=2.)
                dgn, de, cond = geometry_steps(J, r0, .08, .8, rng)
                d = dgn if k == 0 else de
                if not np.isfinite(cond) or cond > 1e10:
                    d = rng.normal(size=p); d *= .8 / max(np.linalg.norm(d), 1e-12)
            else:
                d = rng.normal(size=p); d *= .8 / max(np.linalg.norm(d), 1e-12)
            candidates.append(np.clip(z0 + d, -3., 3.))
        CR = np.array([residual(z) for z in candidates])
        allz, allr = np.vstack([Z, candidates]), np.vstack([R, CR])
        keep = np.argsort(np.sum(allr * allr, axis=1))[:pop]
        Z, R = allz[keep], allr[keep]; trace_z.extend(candidates); trace_r.extend(CR)
        curve.append(float(np.min(np.sum(R * R, axis=1))))
    return curve


def main():
    base, guided = [], []
    for seed in [11, 22, 33, 44, 55, 66, 77, 88]:
        base.append(run(seed, False)); guided.append(run(seed, True))
    checkpoints = [100, 200, 300, 400, 500]
    def at(curves, ev):
        i = max(0, min(len(curves[0])-1, (ev-20)//20-1))
        return float(np.mean([c[i] for c in curves]))
    out = {'verification': verification(), 'benchmark': {
        'seeds': 8, 'checkpoints': checkpoints,
        'baseline_best_loss': [at(base, e) for e in checkpoints],
        'guided_best_loss': [at(guided, e) for e in checkpoints],
        'final_baseline_mean': float(np.mean([c[-1] for c in base])),
        'final_guided_mean': float(np.mean([c[-1] for c in guided])),
        'final_baseline_std': float(np.std([c[-1] for c in base])),
        'final_guided_std': float(np.std([c[-1] for c in guided]))}}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()