"""Event output sides, replay, and phase pairing independent of output errors."""
import math
from pathlib import Path
import tempfile
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.runner import execute_native
from tests.manual.event_phase_comparison import pair_saved_phases
from tests.test_native_catalog import Circuit


class PhasePairingTests(unittest.TestCase):
    def test_roundoff_different_sides_match_reference_phase_using_one_row(self):
        boundary = 32.400000000000006
        nt = [32.4, boundary, 32.41]
        rt = [32.40000000000221, 32.41]
        ni, ai, records = pair_saved_phases(nt, rt,
            [[49000., 2.], [1e17, 2.], [1e17, 2.]], [[1e17, 2.], [1e17, 2.]], [32.4, 32.41], .01)
        self.assertEqual(ni.tolist(), [1, 2])
        self.assertEqual(ai.tolist(), [0, 1])
        self.assertEqual(records[0]['status'], 'matched-other-event-side')

    def test_reference_before_event_uses_saved_left_side(self):
        left = math.nextafter(.8, -math.inf)
        ni, _, _ = pair_saved_phases([left, .8], [.8], [[1e17], [49000.]], [[1e17]], [.8], .01)
        self.assertEqual(ni.tolist(), [0])

    def test_missing_side_or_one_different_signal_is_not_interpolated_or_hidden(self):
        ni, _, records = pair_saved_phases([.79, .8, .81], [.8],
            [[1., 1.], [2., 1.], [2., 2.]], [[2., 2.]], [.8], .01)
        self.assertEqual(ni.tolist(), [-1])
        self.assertEqual(records[0]['status'], 'unmatched-forcing-phase')
        ni, _, records = pair_saved_phases([.79, .81], [.8], [[1.], [1.]], [[1.]], [.8], .01)
        self.assertEqual(ni.tolist(), [-1])
        self.assertEqual(records[0]['status'], 'missing-saved-sample')


class NativeEventOutputTests(unittest.TestCase):
    def test_off_grid_step_keeps_regular_samples_and_both_sides(self):
        b = Circuit()
        b.add('amesim_step0', 'step', time=.055, initial=2., final=4.)
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            with patch.object(builder, 'CACHE', root/'build'), patch.object(result_storage, 'RESULT_ROOT', root/'results'):
                build = builder.build_native(compile_native_program(b.net))
                try:
                    for method in ('BDF', 'RK45'):
                        result = execute_native(build, SolveIVPConfig(t_stop=.2, method=method, max_step=1.),
                                                .01, run_dir=root/method)
                        self.assertTrue(result['success'], result['message'])
                        times = result['series']['time']
                        expected = sorted({i*.01 for i in range(21)} | {.055, math.nextafter(.055, -math.inf)})
                        self.assertEqual(times, expected)
                        self.assertEqual(result['series']['step.y'], [2. if t<.055 else 4. for t in times])
                        self.assertEqual(result['solverStarts'], 2)
                finally:
                    build.close()

    def test_collision_sides_and_recording_do_not_change_integration(self):
        b = Circuit()
        mass = b.add('amesim_mecmas21', 'mass', mass=1, stoptype=3,
                     xmin=0, xmax=1, x0=.9, v0=1, restcoeff=.5, restdvel=.001, useFriction=1)
        for i in (1, 2):
            free = b.add('amesim_f000', 'free'+str(i))
            b.connect(mass, 'port_'+str(i), free, 'port_1')
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            with patch.object(builder, 'CACHE', root/'build'), patch.object(result_storage, 'RESULT_ROOT', root/'results'):
                build = builder.build_native(compile_native_program(b.net))
                try:
                    for method in ('BDF', 'RK45'):
                        config = SolveIVPConfig(t_stop=.25, method=method, max_step=.03, rtol=1e-8)
                        result = execute_native(build, config, .007, run_dir=root/method)
                        control = execute_native(build, config, .007, run_dir=root/(method+'-solve-only'), record_samples=False)
                        self.assertTrue(result['success'], result['message'])
                        self.assertEqual(result['finalState'], control['finalState'])
                        for key in ('acceptedSteps', 'rejectedSteps', 'solverStarts', 'stateTransitions', 'nfev'):
                            self.assertEqual(result[key], control[key], key)
                        series = result['series']
                        times, speeds = series['time'], series['mass.v']
                        self.assertTrue(all(a < z for a, z in zip(times, times[1:])))
                        idx = next(i for i in range(1, len(times)) if speeds[i] < 0 <= speeds[i-1])
                        self.assertEqual(times[idx-1], math.nextafter(times[idx], -math.inf))
                        self.assertAlmostEqual(times[idx], .1, delta=1e-6)
                        self.assertAlmostEqual(speeds[idx-1], 1., places=10)
                        self.assertAlmostEqual(speeds[idx], -.5, places=10)
                        self.assertAlmostEqual(series['mass.x'][idx], 1., places=12)
                finally:
                    build.close()


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