import json
import time
import numpy as np
from scipy.optimize import minimize

SEED = 7

def dirs(theta):
    return np.stack((np.cos(theta), np.sin(theta)), axis=1)

def Bmat(x, theta, k=6.0):
    return np.exp(1j * k * (x @ dirs(theta).T))

def ridge(B, y, lam=0.0, w=None):
    if w is None:
        w = np.ones(len(y))
    A = B.conj().T @ (w[:, None] * B) + lam * np.eye(B.shape[1])
    b = B.conj().T @ (w * y)
    c = np.linalg.solve(A, b)
    r = B @ c - y
    phi = float(np.sum(w * np.abs(r)**2) + lam * np.vdot(c, c).real)
    return c, phi, r

def vp_obj(theta, x, y, k, lam):
    return ridge(Bmat(x, theta, k), y, lam)[1]

def vp_fit(x, y, k, lam, init, maxiter=220):
    result = minimize(lambda t: vp_obj(t, x, y, k, lam), init,
                      method='Nelder-Mead',
                      options={'maxiter': maxiter, 'xatol': 2e-7, 'fatol': 1e-11})
    B = Bmat(x, result.x, k)
    c, phi, r = ridge(B, y, lam)
    return result, c, phi, r

def relative_error(B, c, y):
    return float(np.linalg.norm(B @ c-y) / np.linalg.norm(y))

def checks():
    rng=np.random.default_rng(SEED)
    n=180; k=6.; x=rng.uniform(-1,1,(n,2)); theta=np.array([.1,1.,2.2])
    ctrue=np.array([1+.2j,-.7+.4j,.4-.3j])
    y=Bmat(x,theta,k)@ctrue + .03*(rng.normal(size=n)+1j*rng.normal(size=n))
    ls=np.array([0.,1e-4,1e-3,1e-2,.1,1.])
    ph=[]; cn=[]
    for lam in ls:
        c,p,_=ridge(Bmat(x,theta,k),y,lam); ph.append(p); cn.append(np.vdot(c,c).real)
    ph=np.array(ph); cn=np.array(cn)
    # Envelope prediction: ridge objective nondecreasing; coefficient norm nonincreasing.
    ridge_ok=bool(np.all(np.diff(ph)>=-1e-8) and np.all(np.diff(cn)<=1e-8))

    # Gramian prediction: for two nearly duplicate directions, smallest eigenvalue
    # scales quadratically in angular separation (Taylor expansion of columns).
    base=.65; deltas=np.array([.01,.02,.04,.08,.16]); eig=[]
    for d in deltas:
        G=Bmat(x,np.array([base,base+d]),k).conj().T @ Bmat(x,np.array([base,base+d]),k)
        eig.append(np.linalg.eigvalsh(G)[0]/n)
    eig=np.array(eig)
    slope=float(np.polyfit(np.log(deltas),np.log(np.maximum(eig,1e-16)),1)[0])
    gram_ok=bool(1.5 < slope < 2.5)

    # Direction identifiability prediction: increasing sample count improves recovery.
    recovery=[]
    planted=np.array([.35,1.45]); cc=np.array([1.1-.3j,.65+.5j])
    for nn in [30,60,120,240]:
        xx=rng.uniform(-1,1,(nn,2)); yy=Bmat(xx,planted,k)@cc
        init=np.array([2.7,-1.1])
        out,_,_,_=vp_fit(xx,yy,k,1e-8,init,180)
        # account for 2pi periodicity and column permutation
        est=np.sort(np.mod(out.x,2*np.pi)); true=np.sort(np.mod(planted,2*np.pi))
        recovery.append(float(np.linalg.norm(np.angle(np.exp(1j*(est-true))))))
    # prediction is lower error with more data, allowing optimizer noise.
    recover_ok=bool(recovery[-1] < recovery[0] and recovery[-1] < .25)
    return {
        'ridge_lambdas':ls.tolist(),'ridge_objectives':ph.tolist(),'coefficient_norms':cn.tolist(),
        'ridge_monotonic_predicted':True,'ridge_monotonic_observed':ridge_ok,
        'duplicate_deltas':deltas.tolist(),'duplicate_small_eigenvalues':eig.tolist(),
        'gramian_scaling_predicted_exponent':2.0,'gramian_scaling_observed_exponent':slope,
        'gramian_quadratic_predicted':True,'gramian_quadratic_observed':gram_ok,
        'sample_counts':[30,60,120,240],'direction_errors':recovery,
        'recovery_improves_predicted':True,'recovery_improves_observed':recover_ok,
        'all_checks_pass':ridge_ok and gram_ok and recover_ok}

def experiment():
    rng=np.random.default_rng(SEED+1); k=6.; n=180
    x=rng.uniform(-1,1,(n,2)); planted=np.array([.35,1.45,.35+0.65])
    # Duplicate-looking third component intentionally tests adaptive fitting under redundancy.
    ctrue=np.array([1.0-.25j,.7+.4j,.15-.1j]); y=Bmat(x,planted,k)@ctrue
    init=rng.uniform(0,2*np.pi,3)
    t=time.perf_counter(); out,c,phi,r=vp_fit(x,y,k,1e-7,init,300); tv=time.perf_counter()-t
    idea_err=relative_error(Bmat(x,out.x,k),c,y)
    # Standard fixed random Fourier features: same number of columns, coefficients fitted only.
    fixed=rng.uniform(0,2*np.pi,3); t=time.perf_counter(); cf,pf,rf=ridge(Bmat(x,fixed,k),y,1e-7); tf=time.perf_counter()-t
    base_err=relative_error(Bmat(x,fixed,k),cf,y)
    # A fairer fixed baseline averages several random dictionaries.
    errs=[]
    for _ in range(12):
        th=rng.uniform(0,2*np.pi,3); cb,_,_=ridge(Bmat(x,th,k),y,1e-7)
        errs.append(relative_error(Bmat(x,th,k),cb,y))
    G=Bmat(x,out.x,k).conj().T@Bmat(x,out.x,k)
    ev=np.linalg.eigvalsh(G); retained=int(np.sum(ev >= 1e-3*ev.max()))
    return {'n':n,'features':3,'planted_directions':planted.tolist(),
            'optimized_directions':out.x.tolist(),'idea_relative_error':idea_err,
            'fixed_random_relative_error_one_draw':base_err,'fixed_random_relative_error_mean':float(np.mean(errs)),
            'fixed_random_relative_error_std':float(np.std(errs)), 'retained_by_gramian_cutoff':retained,
            'idea_seconds':tv,'fixed_seconds':tf,'optimizer_success':bool(out.success),'optimizer_iterations':int(out.nit)}

def main():
    result={'checks':checks(),'experiment':experiment()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()