"""Independent analytic and Amesim-reference checks for opt-in event variants."""
import json
import math
import time
import subprocess
import sys
from pathlib import Path
import unittest
from unittest.mock import patch

import numpy as np

from app.simulation.config import SolveIVPConfig
from app.simulation.native_codegen import build as builder, result_storage
from app.simulation.native_codegen.compiler import compile_native_program
from app.simulation.native_codegen.extended import compile_extended_program
from app.simulation.native_codegen.runner import execute_native
from tests.manual.mechanical_event_variant import event_program, prepare_runtime
from tests.test_native_catalog import Circuit

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT/'test/mechanical-events-20260917/isolated'
CONTROL = ROOT/'test/mechanical-events-20260917/source-control/native'
if not CONTROL.exists():
    CONTROL = ROOT/'native'


def free_ports(b, mass, connected=()):
    for port in ('port_1','port_2'):
        if port not in connected:
            free=b.add('amesim_f000',mass.name+'_'+port+'_free')
            b.connect(mass,port,free,'port_1')


def elastic_mass(damping=1, signed=1, pdis=0):
    b=Circuit()
    mass=b.add('amesim_mecmas21','mass',mass=1,useFriction=1,stoptype=2,x0=.5,v0=1,
        xmin=0,xmax=1,Kbmin=100,Kbmax=100,Dbmin=damping,Dbmax=damping,
        Pdmin=pdis,Pdmax=pdis,discContactOption=signed)
    free_ports(b,mass)
    return b


def contact_pair(damping=0,signed=1,pdis=0,x2=.2,v1=1,v2=0):
    b=Circuit()
    a=b.add('amesim_mecmas21','left',mass=1,useFriction=1,stoptype=4,x0=0,v0=v1,
            fcoul=0,fstick=0,rvisc=0,wind=0,theta=0)
    z=b.add('amesim_mecmas21','right',mass=1,useFriction=1,stoptype=4,x0=x2,v0=v2,
            fcoul=0,fstick=0,rvisc=0,wind=0,theta=0)
    c=b.add('amesim_lstp00a','contact',gap0=0,stiffmode=1,kcont=100,rcont=damping,
            Pdis=pdis,discContactOption=signed)
    b.connect(a,'port_2',c,'port_1');b.connect(z,'port_1',c,'port_2')
    free_ports(b,a,('port_2',));free_ports(b,z,('port_1',))
    return b


class MechanicalEventExperimentTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        OUT.mkdir(parents=True,exist_ok=True)
        cls.runtime=OUT/'experimental-native'
        prepare_runtime(CONTROL,cls.runtime)
        cls.run_tag=str(time.time_ns())

    def run_case(self,b,mode,method,*,stop=1,max_step=.05,label='',compiler=compile_native_program,sample_step=.002):
        original=compiler(b.net)
        program,metadata=event_program(original,b.net,mode)
        self.assertEqual(program.state_keys,original.state_keys)
        self.assertEqual(program.jacobian_structure,original.jacobian_structure)
        self.assertEqual(program.evaluation_schedule,original.evaluation_schedule)
        self.assertTrue(program.source.startswith(original.source))
        directory=OUT/self._testMethodName/(mode+'-'+method+label)/self.run_tag
        directory.mkdir(parents=True,exist_ok=True)
        (directory/'event-descriptors.json').write_text(json.dumps(metadata,indent=2)+'\n',encoding='utf-8')
        with patch.object(builder,'NATIVE',CONTROL if mode=='off' else self.runtime), \
             patch.object(builder,'CACHE',OUT/'build-cache'),patch.object(result_storage,'RESULT_ROOT',OUT/'results'):
            build=builder.build_native(program)
            try:
                result=execute_native(build,SolveIVPConfig(t_stop=stop,method=method,max_step=max_step,rtol=1e-9),
                    sample_step,run_dir=directory,timeout=60)
            finally:
                build.close()
        self.assertTrue(result['success'],result['message'])
        self.assertEqual(result['simulatedUntil'],stop)
        times=result['series']['time']
        self.assertTrue(all(a<c for a,c in zip(times,times[1:])))
        events=[]
        for line in (directory/'worker.log').read_text(encoding='utf-8').splitlines():
            try: e=json.loads(line)
            except ValueError: continue
            if e.get('phase')=='mechanical-event':events.append(e)
        return result,events

    def test_mass_elastic_upper_and_lower_contacts_analytic(self):
        duration=math.pi/math.sqrt(100-.25)
        restitution=math.exp(-.5*duration)
        expected=[.5,.5+duration,.5+duration+1/restitution,.5+2*duration+1/restitution]
        for method in ('BDF','RK45'):
            for h in (.08,.02):
                r,events=self.run_case(elastic_mass(),'mass',method,stop=2.5,max_step=h,label=str(h))
                hits=[e for e in events if e['kind']=='mass-contact']
                self.assertEqual([e['direction'] for e in hits],[1,-1,1,-1])
                np.testing.assert_allclose([e['time'] for e in hits],expected,rtol=0,atol=2e-6)
                # Elastic contact must not apply a hard-stop velocity reset.
                np.testing.assert_allclose([abs(e['relativeVelocity']) for e in hits],
                    [1,restitution,restitution,restitution**2],rtol=0,atol=2e-6)

    def test_mass_ideal_release_both_sides_analytic(self):
        for upper in (False,True):
            b=Circuit();sign=-1 if upper else 1
            drive=b.add('amesim_ud00','drive',nstages=1,iscyclic=0,start1=-sign,end1=sign,t1=2)
            force=b.add('amesim_forc','force',direction=1)
            mass=b.add('amesim_mecmas21','mass',mass=1,useFriction=1,stoptype=1,
                       xmin=0,xmax=1,x0=float(upper),v0=0)
            b.connect(drive,'out',force,'res');b.connect(force,'port_2',mass,'port_1');free_ports(b,mass,('port_1',))
            for method in ('BDF','RK45'):
                r,events=self.run_case(b,'mass',method,stop=1.5,max_step=.4,label=str(upper))
                releases=[e for e in events if e['kind']=='mass-release']
                self.assertEqual(len(releases),1)
                self.assertAlmostEqual(releases[0]['time'],1,delta=1e-8)
                self.assertAlmostEqual(r['final']['mass.v'],sign*.125,delta=2e-6)
                self.assertAlmostEqual(r['final']['mass.x'],float(upper)+sign/48,delta=2e-6)

    def test_mass_existing_restitution_impact_is_preserved(self):
        b=Circuit()
        mass=b.add('amesim_mecmas21','mass',mass=1,useFriction=1,stoptype=3,
            xmin=0,xmax=1,x0=.9,v0=1,restcoeff=.5,restdvel=.001)
        free_ports(b,mass)
        for method in ('BDF','RK45'):
            baseline,_=self.run_case(b,'off',method,stop=.4,max_step=.02)
            r,events=self.run_case(b,'mass',method,stop=.4,max_step=.02)
            self.assertEqual(r['series'],baseline['series'])
            self.assertEqual(r['finalState'],baseline['finalState'])
            self.assertEqual(r['solverStarts'],baseline['solverStarts'])
            self.assertEqual(events,[])
            np.testing.assert_allclose(r['finalState'],[-.5,.85],rtol=0,atol=2e-6)

    def test_lstp_elastic_collision_conserves_momentum_and_energy(self):
        expected=[.2,.2+math.pi/math.sqrt(200)]
        for method in ('BDF','RK45'):
            for h in (.06,.015):
                r,events=self.run_case(contact_pair(),'lstp',method,stop=.7,max_step=h,label=str(h))
                hits=[e for e in events if e['kind']=='lstp-contact']
                self.assertEqual([e['direction'] for e in hits],[1,-1])
                np.testing.assert_allclose([e['time'] for e in hits],expected,rtol=0,atol=2e-6)
                v1=np.array(r['series']['left.v']);v2=np.array(r['series']['right.v'])
                penetration=np.maximum(-np.array(r['series']['contact.gap']),0)
                np.testing.assert_allclose(v1+v2,1,rtol=0,atol=1e-9)
                np.testing.assert_allclose(.5*(v1*v1+v2*v2)+50*penetration**2,.5,rtol=0,atol=2e-6)
                self.assertAlmostEqual(r['final']['left.v'],0,delta=2e-6)
                self.assertAlmostEqual(r['final']['right.v'],1,delta=2e-6)

    def test_nonnegative_force_clip_is_local_event_mass_and_lstp(self):
        for label,b,mode in (('mass',elastic_mass(10,2),'mass'),('lstp',contact_pair(10,2),'lstp')):
            for method in ('BDF','RK45'):
                r,events=self.run_case(b,mode,method,stop=1.05,label=label)
                clips=[e for e in events if e['kind']=='force-clip']
                self.assertTrue(clips)
                self.assertTrue(all(e['penetration']>0 for e in clips))
                key='mass.Fmax' if label=='mass' else 'contact.force'
                self.assertGreaterEqual(min(r['series'][key]),0)
                if label=='mass':
                    expected=.5+(math.pi-math.atan(math.sqrt(3)))/math.sqrt(75)
                    self.assertAlmostEqual(clips[0]['time'],expected,delta=3e-6)

    def test_no_contact_initial_touch_and_continuous_contact_do_not_restart(self):
        cases=[('apart',contact_pair(v1=.1,v2=.1),.3),('touch',contact_pair(x2=0,v1=0,v2=0),.3),
               ('inside',contact_pair(x2=-.1,v1=0,v2=0),.01)]
        for label,b,stop in cases:
            a,_=self.run_case(b,'off','BDF',stop=stop,label=label)
            z,events=self.run_case(b,'lstp','BDF',stop=stop,label=label)
            self.assertEqual(events,[])
            self.assertEqual(a['finalState'],z['finalState'])
            self.assertEqual(a['series'],z['series'])
            self.assertEqual(a['solverStarts'],z['solverStarts'])

    def test_exponential_damping_modes_and_extended_lowering(self):
        for signed in (1,2):
            for mode,b in (('mass',elastic_mass(2,signed,.01)),('lstp',contact_pair(2,signed,.01))):
                for method in ('BDF','RK45'):
                    r,events=self.run_case(b,mode,method,stop=1.1,label=str(signed),compiler=compile_extended_program)
                    self.assertGreaterEqual(len(events),2)
                    self.assertLess(len(events),15)
                    self.assertEqual(r['experimentalEvents']['releaseRhsCalls'],0)

    def test_existing_amesim_mechanical_reference_circuits(self):
        cases=json.loads((ROOT/'tests/fixtures/amesim-component-semantics-2404.json').read_text(encoding='utf-8'))['cases']
        for case in cases:
            if case['name'] not in ('friction_ideal','friction_elastic','joint_mechanical'):continue
            b=Circuit()
            for c in case['components']:b.add(c['type'],c['name'],**c['parameters'])
            for left,right in case['connections']:b.net.connect(*left,*right)
            modes=('mass',) if case['name']!='joint_mechanical' else ('lstp','all')
            for mode in modes:
                r,_=self.run_case(b,mode,'BDF',stop=3,max_step=.002,label=case['name'],sample_step=.001)
                for key,expected in case['series'].items():
                    actual=np.interp(case['time'],r['series']['time'],r['series'][key])
                    field=key.split('.')[-1]
                    atol=.002 if field in ('a','Ffric','Fmin','Fmax','force') else 3e-6
                    np.testing.assert_allclose(actual,expected,rtol=0,atol=atol,err_msg=case['name']+'/'+mode+'/'+key)

    def test_same_sign_gap_excursion_and_tangent_with_analytic_dense_output(self):
        directory=OUT/self._testMethodName/self.run_tag;directory.mkdir(parents=True)
        header='''#include "kernels.h"
#define NSTATES 2
#define NOUTPUTS 1
#define NSTOPS 0
#define NEXPERIMENT_CONTACTS 1
#define NEXPERIMENT_RELEASES 0
typedef struct { int kind,v1,v2;double boundary,stiffness,damping,pdis;int signed_force; } ExperimentContact;
extern const ExperimentContact experiment_contacts[1];
'''
        (directory/'model.h').write_text(header,encoding='ascii')
        harness=r'''
#include "runtime.h"
#include <math.h>
#include <string.h>
#include "experimental_events.c"
const ExperimentContact experiment_contacts[1]={{2,0,-1,0,100,1,0,1}};
static int polynomial(void *context,double t,double *y) {
    y[0]=1-2*t;y[1]=(*(int*)context?0:.01)-(t-.5)*(t-.5);return 1;
}
int main(void) {
    for(int tangent=0;tangent<2;tangent++) {
        NativeRun r={0};double a[2],b[2],when[8],t=0;int indices[8],kinds[8],count=0;
        polynomial(&tangent,t,a);polynomial(&tangent,1,b);
        if(!experiment_candidates(&r,t,1,a,b,polynomial,&tangent,when,indices,kinds,&count))return 1;
        if(tangent) {if(count)return 2;continue;}
        if(count!=1 || fabs(when[0]-.4)>1e-12)return 3;
        t=when[0];polynomial(&tangent,t,a);experiment_commit(&r,t,a,when,indices,kinds,count);
        count=0;
        if(!experiment_candidates(&r,t,1,a,b,polynomial,&tangent,when,indices,kinds,&count))return 4;
        if(count!=1 || fabs(when[0]-.6)>1e-12)return 5;
    }
    return 0;
}
'''
        (directory/'check.c').write_text(harness,encoding='ascii')
        cc,_,_=builder.toolchain();exe=directory/('check.exe' if sys.platform=='win32' else 'check')
        builder._command([cc,*builder.COMPILER_FLAGS,'-I',str(directory),'-I',str(self.runtime/'include'),
            '-I',str(self.runtime/'runtime'),str(directory/'check.c'),'-lm','-o',str(exe)],log=[])
        r=subprocess.run([str(exe)],capture_output=True,text=True,timeout=15)
        self.assertEqual(r.returncode,0,r.stderr)


if __name__=='__main__':unittest.main()