"""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=58, amplitude=.0008, contact=False, seconds=3., release_seconds=.0015):
    sr=88200; dt=1/sr; f=440*2**((pitch-69)/12)
    tension=55.; speed=2*.65*(195.997718 if pitch>=55 else 146.832384); 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([102.,195.,340.,520.,700.,1050.,1700.,2600.])
    bm=np.array([.4,.15,.1,.07,.06,.05,.04,.03])
    weight=np.array([.8,1.,.42,.22,.4,.1,.06,.03])
    damping=np.array([32.,42.,55.,75.,95.,135.,190.,270.])
    q=np.zeros(8);vbody=np.zeros(8);y=np.zeros(n);v=np.zeros(n)
    pluck_index=int(.74*(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=.0004+.0014*x
                penetration=bridge-clearance-y[j]
                if contact and x<.72 and penetration>0:
                    # Elastic distributed penalty spring; no random noise.
                    cf=1800.*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-(.0004+.0014*x)-y[j])
                if contact and x<.72:energy+=.5*1800.*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
    x,e,c,p,n=simulate()
    assert np.isfinite(x).all();e=e[e>0];assert max(e)<e[0]*1.03
    x=resample_poly(x,1,2);gain=.75/max(abs(x));write(root/'guitar-bb3.wav',44100,(x*gain*32767).astype(np.int16))
    (root/'report.json').write_text(json.dumps({'gridPoints':n,'peakEnergyRatio':float(max(e)/e[0]),'gain':gain,'status':'Uncalibrated coupled guitar/body experiment'},indent=2))