#!/usr/bin/env python3
"""Independent orbit/transport invariants and actual renderer CSV consumption."""
import importlib.util
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location('track', ROOT / 'scripts/schwarzschild_camera_track.py')
track = importlib.util.module_from_spec(spec)
spec.loader.exec_module(track)


class TrackTests(unittest.TestCase):
    def test_sampling_endpoints(self):
        state = track.initial_state([8, 0, 0], [0, 0, 0])
        for duration, count in ((0., 1), (.29, 30), (.295, 30)):
            tau, states, stopped = track.integrate(state, duration, 100)
            self.assertEqual(len(tau), count)
            self.assertLessEqual(tau[-1], duration)
            self.assertIsNone(stopped)

    def test_radial_infall_through_horizon(self):
        # E=1 infall from rest at infinity: dr/dtau=-sqrt(2/r).
        r = 8.
        w = np.sqrt(2/r)
        ut = (1 + w + w*w) / (1+w)
        state = track.initial_state([r, 0, 0], [-w/ut, 0, 0])
        tau, states, stopped = track.integrate(state, 20, 20, stop_radius=.01)
        expected_stop = 2/(3*np.sqrt(2)) * (r**1.5 - .01**1.5)
        self.assertAlmostEqual(stopped, expected_stop, delta=2e-8)
        radii = np.linalg.norm(states[:, 1:4], axis=1)
        np.testing.assert_allclose(radii, (r**1.5 - 1.5*np.sqrt(2)*tau)**(2/3), atol=2e-8, rtol=2e-8)
        self.assertLess(radii[-1], 2)
        for s in states:
            g, _ = track.metric_connection(s[1:4])
            e = s[4:].reshape(4, 4)
            np.testing.assert_allclose(e @ g @ e.T, track.ETA, atol=2e-8)
            self.assertAlmostEqual(-(g @ e[0])[0], 1, delta=2e-8)

    def test_circular_orbit_and_transport_convergence(self):
        r = 8.
        omega = r**-1.5
        state = track.initial_state([r, 0, 0], [0, r*omega, 0])
        # Choose a pure Schwarzschild radial leg, transformed to KS time.
        # Projecting a zero-KS-time radial seed would produce a different leg.
        e = state[4:].reshape(4, 4)
        root = np.sqrt(1-2/r)
        e[1] = [-2/r/root, -root, 0, 0]
        e[2] = [0, 0, 0, 1]
        e[3] = [e[0, 2]/root, 0, e[0, 0]*root, 0]
        duration = 2*np.pi/omega*np.sqrt(1-3/r)
        results = []
        for tol in (1e-6, 1e-10):
            tau, states, stopped = track.integrate(state, duration, 2, rtol=tol, atol=tol*.01)
            self.assertIsNone(stopped)
            angle = omega*tau/np.sqrt(1-3/r)
            expected = r*np.column_stack([np.cos(angle), np.sin(angle), np.zeros_like(angle)])
            results.append(np.max(np.abs(states[:, 1:4]-expected)))
        self.assertLess(results[1], 2e-7)
        self.assertLess(results[1], results[0]/100)
        s = states[-1]
        e = s[4:].reshape(4, 4)
        g, _ = track.metric_connection(s[1:4])
        np.testing.assert_allclose(e @ g @ e.T, track.ETA, atol=2e-8)
        # Analytic parallel transport of initially inward radial e1 on circular orbit.
        # In Schwarzschild components e1^r=-sqrt(1-2/r) cos(omega*tau).
        n = s[1:4]/np.linalg.norm(s[1:4])
        self.assertAlmostEqual(n @ e[1, 1:], -np.sqrt(1-2/r)*np.cos(omega*tau[-1]), delta=2e-8)

    def test_initial_tetrad_and_invalid_velocity(self):
        state = track.initial_state([2, 0, 0], [-.5, .1, 0], 40, 30, 17)
        explicit = track.initial_state([2, 0, 0], [-.5, .1, 0], tetrad=state[4:].reshape(4, 4))
        np.testing.assert_array_equal(state, explicit)
        with self.assertRaises(ValueError):
            track.initial_state([2, 0, 0], [0, 0, 0])
        with self.assertRaises(ValueError):
            track.initial_state([8, 0, 0], [0, 0, 0], tetrad=np.eye(4))

    def test_csv_movie(self):
        binary = ROOT / 'build/Release/schwarzschild_sky'
        if not binary.exists():
            self.fail('build Schwarzschild Release renderer before running this test')
        with tempfile.TemporaryDirectory() as directory:
            directory = Path(directory)
            csv = directory / 'camera.csv'
            command = [sys.executable, str(ROOT/'scripts/schwarzschild_camera_track.py'),
                       '--output', str(csv), '--position', '8', '0', '0',
                       '--look-ra-deg', '0', '--look-dec-deg', '0',
                       '--fps', '10', '--duration', '.21', '--t0', '7']
            subprocess.run(command, check=True, capture_output=True, text=True)
            rows = np.loadtxt(csv, delimiter=',', skiprows=1)
            self.assertEqual(rows.shape, (3, 21))
            np.testing.assert_allclose(rows[:, 1], [0, .1, .2])
            result = subprocess.run([str(binary), '--observer-track', str(csv),
                '--movie-track-samples', '--frames-dir', str(directory),
                '--catalog', str(ROOT/'assets/sky_grid_5deg.csv'), '--width', '16',
                '--height', '16', '--coarse-cell-pixels', '8', '--refine-max-level', '0'],
                env=dict(os.environ, OMP_NUM_THREADS='2'), capture_output=True, text=True)
            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertEqual(len(list(directory.glob('frame_*.png'))), 3)
            for png in directory.glob('frame_*.png'):
                self.assertEqual(png.read_bytes()[:8], b'\x89PNG\r\n\x1a\n')
            # Error must happen before writing an output file.
            csv.unlink()
            bad = subprocess.run(command + ['--velocity', '2', '0', '0'], capture_output=True)
            self.assertNotEqual(bad.returncode, 0)
            self.assertFalse(csv.exists())


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