"""Time-domain tap test: leapfrog wave equation on each drum, point tap, point mic, FFT."""
from geom import *; import fd, numpy as np, json
N=48; h=1/N; dt=0.5*h; T=400.0; steps=int(T/dt)
def run(P,tap,mic):
    A,p=fd.lap(P,N)
    it=np.argmin(np.linalg.norm(p-tap,axis=1)); im=np.argmin(np.linalg.norm(p-mic,axis=1))
    u0=np.zeros(len(p)); v=np.exp(-np.sum((p-p[it])**2,1)/(2*(1.5*h)**2))   # mallet: velocity kick
    u_prev=u0-dt*v; u=u0.copy(); rec=np.empty(steps); vol=np.empty(steps)
    for n in range(steps):
        u_next=2*u-u_prev-dt*dt*(A@u); u_prev,u=u,u_next; rec[n]=u[im]; vol[n]=u.sum()*h*h
    return rec,vol,p[it],p[im]
def spectrum(x):
    w=np.hanning(len(x)); X=np.abs(np.fft.rfft(x*w,8*len(x))); om=2*np.pi*np.fft.rfftfreq(8*len(x),dt); return om,X
out={}
cases={'tapA':{'D1':((0.5,0.3),(1.5,1.2)),'D2':((0.75,0.6),(1.5,1.2))},'tapB':{'D1':((2.2,1.8),(0.5,0.3)),'D2':((2.2,1.8),(0.75,0.6))}}
for cname,cc in cases.items():
    for dn,P in (('D1',D1),('D2',D2)):
        tap,mic=cc[dn]
        # same tap/mic coordinates must lie inside both; nudge to nearest node anyway
        rec,vol,pt,pm=run(P,np.array(tap),np.array(mic))
        om,X=spectrum(vol)
        m=om<9.0; out[f'{cname}_{dn}']={'om':om[m][::4].tolist(),'X':(X[m][::4]/X[m].max()).tolist(),'tap':pt.tolist(),'mic':pm.tolist()}
        # peaks
        Xm=X[m]; dw=int(0.09/(om[1]-om[0])); pk=[i for i in range(dw,len(Xm)-dw) if Xm[i]==Xm[i-dw:i+dw+1].max() and Xm[i]>0.003*Xm.max()]
        out[f'{cname}_{dn}']['peaks']=[float(om[m][i]) for i in pk]; out[f'{cname}_{dn}']['heights']=[float(Xm[i]/Xm.max()) for i in pk]
        print(cname,dn,'inside?',inside(P,np.array([tap,mic])),'peaks',np.round(out[f'{cname}_{dn}']['peaks'],3),np.round(out[f'{cname}_{dn}']['heights'],2),flush=True)
lam=fd.solve(D1,N,12); wd=np.arccos(1-dt*dt*lam/2)/dt
print('predicted (discrete leapfrog) omegas',np.round(wd,3))
out['pred']=wd.tolist(); json.dump(out,open('tap.json','w'))