"""Exploratory mass-string/body/contact model. No recorded audio.
SI units; uncalibrated dimensions/masses. Velocity Verlet at 88.2 kHz.
String and body exchange equal/opposite bridge and fingerboard forces.
This is a reduced model, not a validated acoustic double bass.
"""
import numpy as np
from numba import njit
from scipy.signal import resample_poly
from scipy.io.wavfile import write
import json
from pathlib import Path

@njit
def simulate(pitch=38, amplitude=.002, contact=True, seconds=3., release_seconds=.003):
    sr=88200; dt=1/sr; f=440*2**((pitch-69)/12)
    open_pitch=43 if pitch>=43 else 38 if pitch>=38 else 33 if pitch>=33 else 28
    tension=180.; speed=2*1.05*440*2**((open_pitch-69)/12); density=tension/speed**2
    length=speed/(2*f); n=min(160,int(sr*.36/(2*f)))
    h=length/(n+1); mass=density*h; k=tension/h
    # Modal masses and frequencies are provisional, not measured responses.
    freq=np.array([68.,108.,175.,285.,460.,735.,1200.,2100.])
    bm=np.array([.9,.65,.5,.35,.26,.18,.12,.08])
    weight=np.array([.6,.8,.7,.5,.35,.24,.16,.1])
    damping=np.array([18.,24.,32.,42.,58.,85.,130.,210.])
    q=np.zeros(8);vbody=np.zeros(8);y=np.zeros(n);v=np.zeros(n)
    pluck_index=int(.76*(n+1))-1
    pluck=(pluck_index+1)/(n+1)
    for j in range(n):
        x=(j+1)/(n+1);y[j]=amplitude*(x/pluck if x<pluck else (1-x)/(1-pluck))
    # Equilibrate body and string under the held pluck before releasing it.
    compliance=np.sum(weight*weight/(bm*(2*np.pi*freq)**2))
    held_bridge_force=k*y[-1]/(1+k*compliance/(n+1))
    for j in range(8):q[j]=weight[j]*held_bridge_force/(bm[j]*(2*np.pi*freq[j])**2)
    bridge=np.sum(q*weight)
    for j in range(n):y[j]+=(j+1)/(n+1)*bridge
    held_finger_force=tension*amplitude/(length*pluck*(1-pluck))
    out=np.zeros(int(seconds*sr)); force=np.zeros(n);bf=np.zeros(8)
    # String pulled away from the board; collisions arise on its return.
    contacts=0; peak_penetration=0.; energies=np.zeros(1+len(out)//882)
    for frame in range(len(out)):
        for half in range(2):
            bridge=np.sum(q*weight)
            total_contact=0.
            for j in range(n):
                left=0. if j==0 else y[j-1]
                right=bridge if j==n-1 else y[j+1]
                force[j]=k*(left-2*y[j]+right)
                x=(j+1)/(n+1)
                clearance=.0007+.0024*x
                penetration=bridge-clearance-y[j]
                if contact and x<.72 and penetration>0:
                    # Elastic distributed penalty spring; no random noise.
                    cf=22000.*penetration
                    force[j]+=cf;total_contact+=cf
                    if half==0:contacts+=1
                    peak_penetration=max(peak_penetration,penetration)
            # Finite release of the force holding the string, not an audio envelope.
            time=(frame+half)*dt
            if time<release_seconds:
                force[pluck_index]+=held_finger_force*.5*(1+np.cos(np.pi*time/release_seconds))
            bridgeforce=k*(y[-1]-bridge)-total_contact
            for j in range(8):bf[j]=weight[j]*bridgeforce-bm[j]*(2*np.pi*freq[j])**2*q[j]
            for j in range(n):v[j]+=dt*.5*force[j]/mass
            for j in range(8):vbody[j]+=dt*.5*bf[j]/bm[j]
            if half==0:
                # Exact velocity damping, split from conservative forces.
                for j in range(n):v[j]*=np.exp(-dt*.9);y[j]+=dt*v[j]
                for j in range(8):vbody[j]*=np.exp(-dt*damping[j]);q[j]+=dt*vbody[j]
        out[frame]=np.sum(vbody*weight)
        if frame%882==0:
            bridge=np.sum(q*weight)
            energy=.5*mass*np.sum(v*v)+.5*np.sum(bm*vbody*vbody)+.5*np.sum(bm*(2*np.pi*freq)**2*q*q)
            prev=0.
            for j in range(n):
                energy+=.5*k*(y[j]-prev)**2;prev=y[j]
                x=(j+1)/(n+1);pen=max(0.,bridge-(.0007+.0024*x)-y[j])
                if contact and x<.72:energy+=.5*22000.*pen*pen
            energy+=.5*k*(bridge-y[-1])**2
            energies[frame//882]=energy
    return out,energies,contacts,peak_penetration,n


if __name__=='__main__':
    root=Path(__file__).parent;cases={};renders={}
    for name,amp,on in [('soft',.0007,True),('firm',.003,True),('firm-no-board',.003,False)]:
        x,e,c,p,n=simulate(amplitude=amp,contact=on)
        assert np.isfinite(x).all()
        e=e[e>0];assert max(e)<e[0]*1.03, (name,max(e)/e[0])
        renders[name]=resample_poly(x,1,2)
        cases[name]={'contacts':c,'maxPenetrationMeters':p,'gridPoints':n,'peakEnergyRatio':float(max(e)/e[0]),'finalEnergyRatio':float(e[-1]/e[0])}
    gain=.8/max(np.max(abs(x)) for x in renders.values())
    for name,x in renders.items():write(root/(name+'.wav'),44100,(x*gain*32767).astype(np.int16))
    (root/'report.json').write_text(json.dumps({'cases':cases,'commonGain':gain,'status':'Uncalibrated isolated-note experiment; not a realism result'},indent=2));print(json.dumps(cases,indent=2))