"""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, grid_limit=160, note_off=-1.):
    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(grid_limit,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))
    # A finite patch of the string is damped when the finger returns.
    stop_patch=np.zeros(n)
    for j in range(n):stop_patch[j]=np.exp(-.5*(((j+1)/(n+1)-pluck)/.055)**2)
    # A separate neck coordinate receives fingerboard reactions. It is
    # coupled passively to body modes, not mistaken for bridge force.
    neck_weight=weight*.045
    neck_mass=.35;neck_stiffness=neck_mass*(2*np.pi*480.)**2
    neck_damping=neck_mass*100.
    neck_position=np.sum(q*neck_weight);neck_velocity=0.
    # Distributed contact coefficient includes grid-cell length.
    contact_density=3.3e6/np.sqrt(.00005)
    contact_loss=5.
    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.
            neck_relative=neck_position-np.sum(q*neck_weight)
            neck_relative_velocity=neck_velocity-np.sum(vbody*neck_weight)
            neck_link=neck_stiffness*neck_relative+neck_damping*neck_relative_velocity
            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=neck_position-clearance-y[j]
                if contact and x<.72 and penetration>0:
                    # Lossy, distributed compliant contact; no independent noise.
                    cf=contact_density*h*penetration**1.5*max(0.,1+contact_loss*(neck_velocity-v[j]))
                    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)
            neck_force=-total_contact-neck_link
            for j in range(8):bf[j]=weight[j]*bridgeforce+neck_weight[j]*neck_link-bm[j]*(2*np.pi*freq[j])**2*q[j]
            neck_velocity+=dt*.5*neck_force/neck_mass
            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:
                neck_position+=dt*neck_velocity
                # Exact velocity damping, split from conservative forces.
                stop_amount=0.
                if note_off>=0 and frame*dt>note_off:
                    progress=min(1.,(frame*dt-note_off)/.022)
                    stop_amount=.5-.5*np.cos(np.pi*progress)
                for j in range(n):
                    v[j]*=np.exp(-dt*(.9+700.*stop_amount*stop_patch[j]))
                    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)
            energy+=.5*neck_mass*neck_velocity**2+.5*neck_stiffness*(neck_position-np.sum(q*neck_weight))**2
            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.,neck_position-(.0007+.0024*x)-y[j])
                if contact and x<.72:energy+=contact_density*h*pen**2.5/2.5
            energy+=.5*k*(bridge-y[-1])**2
            energies[frame//882]=energy
    return out,energies,contacts,peak_penetration,n