import json, math, time
import numpy as np


def make_features(z, omega, kappa):
    z = np.asarray(z)
    return np.exp(1j * kappa * (z @ omega.T)) / np.sqrt(len(omega))


def projector_from_sources(sources, omega, kappa, noise=0.0, rng=None):
    A = make_features(sources, omega, kappa).T
    if noise:
        if rng is None: rng = np.random.default_rng(0)
        A = A + noise * (rng.normal(size=A.shape) + 1j*rng.normal(size=A.shape)) / np.sqrt(2*len(omega))
    U, _, _ = np.linalg.svd(A, full_matrices=False)
    return U[:, :len(sources)]


def score_grad(z, U, omega, kappa):
    """q=1-||U^*phi||^2 and its exact real spatial gradient."""
    phi = make_features(np.asarray(z)[None, :], omega, kappa)[0]
    Pphi = U @ (U.conj().T @ phi)
    q = float(1.0 - np.vdot(phi, Pphi).real)
    dphi = (1j * kappa * omega) * phi[:, None]
    grad = -2.0 * np.real(np.conj(dphi).T @ Pphi)
    return q, grad


def project(z, lo, hi):
    return np.minimum(np.maximum(z, lo), hi)


def refine(z0, U, omega, kappa, lo, hi, steps=40, eta=0.1, tol=1e-7):
    z = np.array(z0, dtype=float)
    used = 0
    for t in range(steps):
        q, g = score_grad(z, U, omega, kappa)
        h = eta / (kappa*kappa)
        # Backtracking makes the stated fixed scaling safe under finite samples.
        while True:
            zn = project(z - h*g, lo, hi)
            qn, _ = score_grad(zn, U, omega, kappa)
            if qn <= q + 1e-12 or h < 1e-8/(kappa*kappa): break
            h *= 0.5
        z = zn; used += 1
        if np.linalg.norm(g) < tol: break
    return z, score_grad(z, U, omega, kappa)[0], used


def nms(points, values, radius):
    order = np.argsort(values)
    kept = []
    for i in order:
        if all(np.linalg.norm(points[i]-points[j]) > radius for j in kept): kept.append(i)
    return kept


def nearest_errors(found, sources):
    return [float(np.min(np.linalg.norm(np.asarray(found)-s, axis=1))) if len(found) else float('inf') for s in sources]


def main():
    rng = np.random.default_rng(3138)
    d, M, kappa = 2, 96, 18.0
    lo, hi = np.zeros(d), np.ones(d)
    sources = np.array([[.23,.31],[.68,.72],[.79,.25]])
    omega = rng.normal(size=(M,d)); omega /= np.linalg.norm(omega, axis=1, keepdims=True)
    U = projector_from_sources(sources, omega, kappa, noise=0.015, rng=rng)

    # Stage 1: exact gradient versus central differences, and h~k^-2 stability.
    ztest = np.array([.41,.57]); q,g = score_grad(ztest,U,omega,kappa)
    eps=1e-5; fd=[]
    for j in range(d):
        a=ztest.copy(); b=ztest.copy(); a[j]+=eps; b[j]-=eps
        fd.append((score_grad(a,U,omega,kappa)[0]-score_grad(b,U,omega,kappa)[0])/(2*eps))
    grad_rel_err=float(np.linalg.norm(g-np.array(fd))/(np.linalg.norm(g)+1e-12))
    stability=[]
    for kk in [kappa/2,kappa,2*kappa]:
        # same dimensionless eta, measure average one-step decrease from random points
        drops=[]
        for z in rng.uniform(.05,.95,size=(100,d)):
            qq,gg=score_grad(z,U,omega,kk) # intentionally U fixed: tests step scaling only
            zn=project(z-.1*gg/(kk*kk),lo,hi)
            drops.append(score_grad(zn,U,omega,kk)[0]-qq)
        stability.append(float(np.mean(np.array(drops)<=1e-10)))

    # A grid at c_g/kappa; coarse stage keeps only low-score cells and NMSes them.
    spacing=0.8/kappa
    axes=[np.arange(0,1+1e-9,spacing) for _ in range(d)]
    grid=np.array(np.meshgrid(*axes,indexing='ij')).reshape(d,-1).T
    vals=np.array([score_grad(z,U,omega,kappa)[0] for z in grid])
    labels=np.min(np.linalg.norm(grid[:,None,:]-sources[None,:,:],axis=2),axis=1)<0.65/kappa
    pos, neg = vals[labels], vals[~labels]
    tau=float(.5*(np.max(pos)+np.min(neg))) if np.max(pos)<np.min(neg) else float(np.quantile(vals,.12))
    accepted=np.flatnonzero(vals<=tau)
    survivors=nms(grid[accepted], vals[accepted], spacing/2)
    starts_coarse=grid[accepted][survivors]

    def run(starts):
        out=[]; steps=0
        t=time.perf_counter()
        for z in starts:
            x,v,n=refine(z,U,omega,kappa,lo,hi)
            out.append(x); steps+=n
        elapsed=time.perf_counter()-t
        keep=nms(np.array(out),np.array([score_grad(x,U,omega,kappa)[0] for x in out]),.55/kappa) if out else []
        return np.array(out), len(keep), steps, elapsed
    all_out, all_unique, all_steps, all_time=run(grid)
    coarse_out, coarse_unique, coarse_steps, coarse_time=run(starts_coarse)
    result={
      'config':{'kappa':kappa,'M':M,'grid_points':len(grid),'sources':len(sources),'grid_spacing':spacing,'threshold':tau},
      'math_check':{'gradient_relative_error':grad_rel_err,'decrease_fraction_eta_over_k2':[{'kappa':x,'fraction':y} for x,y in zip([kappa/2,kappa,2*kappa],stability)]},
      'baseline_exhaustive':{'starts':len(grid),'refinement_steps':all_steps,'unique_outputs':all_unique,'source_errors':nearest_errors(all_out, sources),'seconds':all_time},
      'idea_coarse_to_fine':{'accepted':len(accepted),'survivor_starts':len(starts_coarse),'refinement_steps':coarse_steps,'unique_outputs':coarse_unique,'source_errors':nearest_errors(coarse_out,sources),'seconds':coarse_time},
    }
    print(json.dumps(result, indent=2))

if __name__=='__main__': main()