"""Production LSTP event regression against analytic dynamics and Amesim."""
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
import json
import math
from pathlib import Path
import re
import subprocess
import sys
import time
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.contacts import contact_table
from app.simulation.native_codegen.extended import compile_extended_program
from app.simulation.native_codegen.runner import execute_native
from tests.test_mechanical_event_experiment import contact_pair
from tests.test_native_catalog import Circuit

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT/'test/lstp-mainline-20260917/unit'


class NativeContactEventsTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        OUT.mkdir(parents=True, exist_ok=True)
        cls.tag = str(time.time_ns())

    def run_case(self, b, method='BDF', stop=.7, h=.02, compiler=compile_native_program,
                 disabled=False, label='', record=True, step=.002):
        program = compiler(b.net)
        if disabled:
            program = replace(program, header=re.sub(r'#define NCONTACTS \d+', '#define NCONTACTS 0', program.header))
        path = OUT/self._testMethodName/self.tag/(method+compiler.__name__+str(disabled)+label)
        with patch.object(builder, 'CACHE', OUT/'cache'), \
             patch.object(result_storage, 'RESULT_ROOT', OUT/'results'), \
             patch.object(builder, 'ThreadPoolExecutor', lambda **kw: ThreadPoolExecutor(max_workers=1)):
            build = builder.build_native(program)
            try:
                r = execute_native(build, SolveIVPConfig(t_stop=stop, method=method, max_step=h, rtol=1e-9),
                    step, run_dir=path, record_samples=record, timeout=60)
            finally:
                build.close()
        self.assertTrue(r['success'], r['message'])
        events = []
        for line in (path/'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 r, events

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

    def test_nonnegative_force_clips_before_detachment(self):
        for method in ('BDF','RK45'):
            r,events=self.run_case(contact_pair(10,2),method)
            clips=[e for e in events if e['kind']=='force-clip']
            self.assertEqual(len(clips),1)
            self.assertAlmostEqual(clips[0]['time'],.2+math.pi/20,delta=3e-6)
            self.assertGreater(clips[0]['penetration'],0)
            self.assertGreaterEqual(min(r['series']['contact.force']),0)

    def test_geometry_stiffness_uses_same_event_parameters(self):
        b=contact_pair();contact=b.net.components['contact']
        # Independent geometry: G*d^4/(8*D^3*n) = 100 N/m exactly.
        for name,value in dict(stiffmode=2,G=51200,sdiam=.5,wdiam=.25,na=2,kcont=1e9).items():
            setattr(contact,name,value)
        contact.set_parameter_values({p.name:getattr(contact,p.name) for p in contact.PARAMETERS})
        for method in ('BDF','RK45'):
            r,events=self.run_case(b,method)
            np.testing.assert_allclose([e['time'] for e in events],[.2,.2+math.pi/math.sqrt(200)],atol=2e-6,rtol=0)
            self.assertAlmostEqual(r['final']['left.v'],0,delta=2e-6)
            self.assertAlmostEqual(r['final']['right.v'],1,delta=2e-6)

    def test_exponential_damping_modes(self):
        for option in (1,2):
            for method in ('BDF','RK45'):
                r,events=self.run_case(contact_pair(2,option,.01),method,stop=1.1,label=str(option))
                self.assertGreaterEqual(len(events),2)
                self.assertLess(len(events),15)
                self.assertNotIn('experimentalEvents',r)

    def test_no_events_preserve_trajectory(self):
        for label,b,stop in [('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)]:
            r,events=self.run_case(b,stop=stop,label=label)
            base,_=self.run_case(b,stop=stop,disabled=True,label=label)
            self.assertEqual(events,[])
            for key in ('series','finalState','solverStarts','nfev'):
                self.assertEqual(r[key],base[key])

    def test_output_replay_does_not_change_events_or_state(self):
        for method in ('BDF','RK45'):
            r,events=self.run_case(contact_pair(),method)
            bare,bare_events=self.run_case(contact_pair(),method,label='bare',record=False)
            for key in ('finalState','contactEvents','nfev','acceptedSteps','solverStarts'):
                self.assertEqual(r[key],bare[key])
            self.assertEqual(events,bare_events)
            times=r['series']['time']
            self.assertTrue(all(a<b for a,b in zip(times,times[1:])))
            for event in events:
                self.assertIn(event['time'],times)
                self.assertIn(math.nextafter(event['time'],-math.inf),times)

    def test_existing_amesim_joint_mechanical(self):
        cases=json.loads((ROOT/'tests/fixtures/amesim-component-semantics-2404.json').read_bytes())['cases']
        case=next(c for c in cases if c['name']=='joint_mechanical')
        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)
        r,_=self.run_case(b,stop=3,h=.002,step=.001)
        for key,expected in case['series'].items():
            actual=np.interp(case['time'],r['series']['time'],r['series'][key])
            atol=.002 if key.split('.')[-1] in ('a','Ffric','Fmin','Fmax','force') else 3e-6
            np.testing.assert_allclose(actual,expected,atol=atol,rtol=0,err_msg=key)

    def test_dense_excursion_tangent_and_expression_association(self):
        path=OUT/self._testMethodName/self.tag;path.mkdir(parents=True)
        _,header=contact_table([(0,2,0,100,1,0,1,0)])
        (path/'model.h').write_text('#define NSTATES 4\n#define NOUTPUTS 1\n#define NSTOPS 0\n'+header,encoding='ascii')
        code=r'''
#include "runtime.h"
#include <math.h>
#include "contact_events.h"
const NativeContact model_contacts[1]={{0,2,0,100,1,0,1,0}};
static int polynomial(void *context,double t,double *y) {
    y[0]=1-2*t;y[1]=(*(int*)context?0:.01)-(t-.5)*(t-.5);y[2]=y[3]=0;return 1;
}
int main(void) {
    for(int tangent=0;tangent<2;tangent++) {
        NativeRun r={0};double a[4],b[4],when[8],t=0;int indices[8],kinds[8],count=0;
        polynomial(&tangent,t,a);polynomial(&tangent,1,b);
        if(!contact_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);contact_commit(&r,t,a,when,indices,kinds,count);
        count=0;
        if(!contact_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;
}
'''
        (path/'check.c').write_text(code,encoding='ascii')
        cc,_,_=builder.toolchain();exe=path/('check.exe' if sys.platform=='win32' else 'check')
        builder._command([cc,*builder.COMPILER_FLAGS,'-I',str(path),'-I',str(ROOT/'native/include'),
            '-I',str(ROOT/'native/runtime'),str(path/'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)
        # Both schedules must register their own arithmetic association.
        compact=compile_native_program(contact_pair().net)
        extended=compile_extended_program(contact_pair().net)
        self.assertIn('0.0,100.0,0.0,0.0,1,1}',compact.source)
        self.assertIn('0.0,100.0,0.0,0.0,1,0}',extended.source)


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