import json, math
from pathlib import Path
import numpy as np
from scipy.optimize import least_squares
SEED=2295

def mixture(p,w,mu,sig):
    z=p[:,None,:]-np.asarray(mu)[None,:,:]
    q=np.sum(z*z,axis=2); s=np.asarray(sig)
    return np.sum(np.asarray(w)*(2*np.pi*s*s)**-1*np.exp(-q/(2*s*s)),axis=1)

def make_grid(n=7,extent=4.):
    a=np.linspace(-extent,extent,n); x,y=np.meshgrid(a,a,indexing='ij')
    return np.c_[x.ravel(),y.ravel()],(a[1]-a[0])**2

def projected_collision(p,gamma):
    n=len(p); A=np.c_[np.ones(n),np.sum(p*p,axis=1)]
    P=A@np.linalg.inv(A.T@A)@A.T
    return gamma*(np.eye(n)-P)

def response(L,p,k,omega):
    n=len(p); M=L+1j*(k*p[:,0]-omega)*np.eye(n)
    return np.mean(np.linalg.solve(M,np.ones(n)))

def fit_test():
    p,_=make_grid(11,4.5)
    tw=np.array([.58,.27,.15]); tm=np.array([[.4,-.25],[-1,.8],[1.2,1.] ]); ts=np.array([.55,.8,.38])
    y=mixture(p,tw,tm,ts)
    def unpack(z):
        q=z[:3]; w=np.exp(q-q.max()); w/=w.sum(); mu=z[3:9].reshape(3,2); sig=np.log1p(np.exp(z[9:12]))+1e-4
        return w,mu,sig
    def fun(z):
        w,m,s=unpack(z); return mixture(p,w,m,s)-y
    z=np.zeros(12); z[3:9]=np.array([[.2,0],[-.5,.4],[.7,.5]]).ravel(); z[9:]=np.log(np.expm1(.7))
    sol=least_squares(fun,z,max_nfev=60)
    pred=mixture(p,*unpack(sol.x));
    c,_=make_grid(7,4.5); yc=mixture(c,tw,tm,ts)
    ii=np.clip(np.round((p[:,0]+4.5)/(9/6)).astype(int),0,6); jj=np.clip(np.round((p[:,1]+4.5)/(9/6)).astype(int),0,6)
    gp=yc[ii*7+jj]
    return {'mixture_rmse':float(np.sqrt(np.mean((pred-y)**2))), 'grid_rmse':float(np.sqrt(np.mean((gp-y)**2))), 'mixture_min':float(pred.min()),'grid_min':float(gp.min())}

def main():
    p,dp=make_grid(); w=np.array([.2,.5,.3]); mu=np.array([[.4,.2],[-.8,.5],[1.1,-.7]]); sig=np.array([.45,.7,.35])
    f=mixture(p,w,mu,sig)
    eig=[]
    for g in [.25,.5,1.,2.]:
        ev=np.linalg.eigvalsh(projected_collision(p,g)); nz=ev[ev>1e-8]
        eig.append({'gamma':g,'zero_modes':int((ev<1e-8).sum()),'median_nonzero':float(np.median(nz)),'relative_error':float(abs(np.median(nz)-g)/g)})
    L=projected_collision(p,1.); ratios=np.array([.25,.5,1.,2.,4.]); widths=[]; peak_freq=[]
    ws=np.linspace(-3,3,25)
    for r in ratios:
        amp=np.array([abs(response(L,p,r,om)) for om in ws]); peak=amp.max(); ii=np.where(amp>=peak/2)[0]
        widths.append(float(ws[ii[-1]]-ws[ii[0]])); peak_freq.append(float(ws[amp.argmax()]))
    out={'seed':SEED,'quadrature_points':len(p),'normalization_integral_truncated':float(f.sum()*dp),'normalization_error_due_to_box':float(abs(f.sum()*dp-1)),'min_density':float(f.min()),'eigen_sweep':eig,'crossover':{'predicted_ratio':1.,'ratios':ratios.tolist(),'fwhm':widths,'peak_frequency':peak_freq},'fit':fit_test()}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()