"""Replay verified MN9 spike trains into MuJoCo; condition names never enter solve()."""
from pathlib import Path
import hashlib,json,time,resource
import numpy as np
import mujoco as mj

HERE=Path(__file__).resolve().parent
ROOT=HERE.parent.parent
SOURCE=ROOT/'artifacts/digital-fly/results/results.json'
PARAMS={'dt':0.0001,'duration':1.3,'settle':0.4,'decaySeconds':0.05,'spikeScale':10.,'maxTorque':0.06,'sampleEvery':50}

def sha(path):return hashlib.sha256(path.read_bytes()).hexdigest()

def solve(motor_spikes, *, coupling=1., stiffness_scale=1., timestep=None):
    """Input: two numeric spike-time arrays (ms), no condition or stimulus metadata.

    Causal exponential activation, a fixed saturating torque, and an actual
    dynamic joint. No qpos writes occur during integration.
    """
    m=mj.MjModel.from_xml_path(str(HERE/'model/embodied.xml'))
    m.jnt_stiffness[0]*=stiffness_scale
    if timestep is not None:m.opt.timestep=timestep
    dt=m.opt.timestep;d=mj.MjData(m)
    for _ in range(round(PARAMS['settle']/dt)):mj.mj_step(m,d)
    mj.mj_forward(m,d);baseline=float(d.qpos[0]);d.time=0
    spikes=np.sort(np.concatenate([np.array(s,dtype=float) for s in motor_spikes]))/1000
    assert np.isfinite(spikes).all() and (spikes>=0).all()
    n=round(PARAMS['duration']/dt);arr=np.zeros((n+1,8));j=0;a=0.
    moving=[m.geom('c_rostrum').id,m.geom('c_haustellum').id]
    transforms=[]
    every=round(.005/dt)
    for step in range(n+1):
        t=step*dt
        while j<len(spikes) and spikes[j]<=t+1e-12:
            a+=np.exp(-(t-spikes[j])/PARAMS['decaySeconds']);j+=1
        activation=1-np.exp(-a/PARAMS['spikeScale'])
        d.ctrl[0]=-coupling*PARAMS['maxTorque']*activation
        mj.mj_forward(m,d)
        arr[step]=[t,d.qpos[0],d.qvel[0],d.ctrl[0],d.actuator_force[0],activation,*d.energy]
        if step%every==0:
            transforms.append([np.concatenate((d.geom_xpos[g],d.geom_xmat[g])).tolist() for g in moving])
        if step<n:mj.mj_step(m,d)
        a*=np.exp(-dt/PARAMS['decaySeconds'])
    assert np.isfinite(arr).all() and not any(w.number for w in d.warning)
    return arr,baseline,transforms,m,d

def export_scene(m,d):
    meshes=[]
    for mid in range(m.nmesh):
        v,n=m.mesh_vertadr[mid],m.mesh_vertnum[mid];f,k=m.mesh_faceadr[mid],m.mesh_facenum[mid]
        meshes.append({'vertices':np.round(m.mesh_vert[v:v+n],7).flatten().tolist(),'faces':m.mesh_face[f:f+k].flatten().tolist()})
    geoms=[];zmin=1e9
    for g in range(m.ngeom):
        if m.geom_type[g]!=mj.mjtGeom.mjGEOM_MESH:continue
        mid=int(m.geom_dataid[g]);v,n=m.mesh_vertadr[mid],m.mesh_vertnum[mid]
        zmin=min(zmin,float((m.mesh_vert[v:v+n]@d.geom_xmat[g].reshape(3,3).T+d.geom_xpos[g])[:,2].min()))
        geoms.append({'id':g,'name':mj.mj_id2name(m,mj.mjtObj.mjOBJ_GEOM,g),'mesh':mid,'frame':np.concatenate((d.geom_xpos[g],d.geom_xmat[g])).tolist()})
    return {'meshes':meshes,'geoms':geoms,'floorZ':zmin-.02,'anchor':d.xanchor[0].tolist(),'moving':[m.geom('c_rostrum').id,m.geom('c_haustellum').id]}

def main():
    start=time.perf_counter();source=json.loads(SOURCE.read_text())
    assert source['referenceCheck']['allSpikeTrainsMatch']
    neural_raw_hashes={}
    for trial in source['trials']:
        path=SOURCE.parent/trial['rawFile'];raw=np.load(path)
        for rootid,exported in zip(source['inputIds']['ids_mn9'],trial['motorSpikes']):
            index=int(np.flatnonzero(raw['root_ids']==int(rootid))[0])
            actual=raw['times_ms'][raw['indices']==index]
            assert len(actual)==len(exported) and np.allclose(actual,exported,atol=.00005,rtol=0)
        neural_raw_hashes[trial['id']]=sha(path)
    report={'schema':1,'parameters':PARAMS,'neuralSourceSHA256':sha(SOURCE),'modelSHA256':sha(HERE/'model/embodied.xml'),'flygymCommit':'38c8ec61034cd59bc5ba0de20688d4a3c0000d60','mujocoVersion':mj.__version__,'neuralRawHashes':neural_raw_hashes,'trials':[]}
    traces={};base={};scene=None
    def run(trial,**kwargs):
        nonlocal scene
        a,b,frames,m,d=solve(trial['motorSpikes'],**kwargs)
        name=trial['id'];traces[name]=a;base[name]=b
        path=HERE/'results'/f'{name}.npz';np.savez_compressed(path,trace=a,baseline=b,motor1=trial['motorSpikes'][0],motor2=trial['motorSpikes'][1])
        samples=a[::PARAMS['sampleEvery']]
        r={'id':name,'condition':trial['condition'],'seed':trial['seed'],'motorSpikes':trial['motorSpikes'],'baselineRad':b,'peakDeflectionDeg':float(np.max(np.abs(a[:,1]-b))*180/np.pi),'peakTorque':float(np.max(np.abs(a[:,3]))),'rawSHA256':sha(path),'samples':np.round(samples,9).tolist(),'frames':np.round(frames,8).tolist()}
        report['trials'].append(r)
        if scene is None:scene=export_scene(m,d)
        print(name,'spikes',sum(map(len,trial['motorSpikes'])),'peak degrees',round(r['peakDeflectionDeg'],4),flush=True)
    for trial in source['trials']:run(trial)
    sugar=next(t for t in source['trials'] if t['id']=='sugar-1729')
    run({**sugar,'id':'bridge-off-1729','condition':'bridge-off'},coupling=0.)
    run({**sugar,'id':'stiffer-1729','condition':'stiffer'},stiffness_scale=2.)
    # Causal controls and numerical refinement, not aesthetic fitting.
    zero=traces['rest-1729']
    for t in source['trials']:
        if t['condition'] in ('rest','disconnected','bitter'):
            assert np.array_equal(traces[t['id']],zero)
    assert np.array_equal(traces['bridge-off-1729'][:,:5],zero[:,:5])
    for seed in (1729,1730,1731):
        s=next(t for t in report['trials'] if t['id']==f'sugar-{seed}')
        mix=next(t for t in report['trials'] if t['id']==f'mixed-{seed}')
        assert s['peakDeflectionDeg']>mix['peakDeflectionDeg']>0
    assert np.array_equal(traces['stiffer-1729'][:,3:6],traces['sugar-1729'][:,3:6])
    assert next(t for t in report['trials'] if t['id']=='stiffer-1729')['peakDeflectionDeg']<next(t for t in report['trials'] if t['id']=='sugar-1729')['peakDeflectionDeg']
    first=min(sum(sugar['motorSpikes'],[]))/1000
    assert np.max(np.abs(traces['sugar-1729'][traces['sugar-1729'][:,0]<first,1]-zero[zero[:,0]<first,1]))<1e-12
    finer,_,_,_,_=solve(sugar['motorSpikes'],timestep=.00005)
    error=float(np.max(np.abs(finer[::2,1]-traces['sugar-1729'][:,1]))*180/np.pi)
    assert error<.2
    checks={'neuralInputsMatchOriginalRawNPZ':True,'allFiniteNoSolverWarnings':True,'zeroSpikeConditionsIdentical':True,'bridgeOffMatchesRest':True,'sugarExceedsMixedAllSeeds':True,'sameSpikesStifferJointMovesLess':True,'noResponseBeforeFirstMotorSpike':True,'halfTimestepMaxAngleDifferenceDeg':error,'zeroInputResidualDeflectionDeg':float(np.max(np.abs(zero[:,1]-base['rest-1729']))*180/np.pi)}
    report['checks']=checks;report['wallSeconds']=time.perf_counter()-start;report['peakRssMiB']=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
    (HERE/'results/summary.json').write_text(json.dumps({**report,'trials':[{k:v for k,v in t.items() if k not in ('samples','frames')} for t in report['trials']]},indent=2)+'\n')
    (HERE/'results/playback.json').write_text(json.dumps(report,separators=(',',':')))
    (HERE/'results/scene.json').write_text(json.dumps(scene,separators=(',',':')))
    print(json.dumps(checks),flush=True)
    print('Wall seconds',report['wallSeconds'],'peak MiB',report['peakRssMiB'],flush=True)
if __name__=='__main__':main()