import os
from unittest import TestCase

import numpy as np
from   numpy.testing import assert_allclose

from   scipy.integrate   import cumulative_trapezoid
from   scipy.interpolate import interp1d

from math_utils      import array_utils
from nr_analysis_axi import doublenull, invars

class TestDoubleNull(TestCase):

    def _data_path(self):
        return os.path.splitext(__file__)[0] + '.npz'

    def update_refs(self, **kwargs):
        data = dict(self.data) | kwargs
        np.savez(self._data_path(), **data)

    def setUp(self):
        self.data = np.load(self._data_path())

    def test_null_curves(self):
        rays = doublenull.null_curves(self.data['t'], self.data['x'],
                                      self.data['gxx'], np.zeros_like(self.data['gxx']),
                                      -(self.data['alpha'] ** 2))

        assert_allclose(rays[0], self.data['ray_times'], 1e-12)
        assert_allclose(rays[1], self.data['rays_pos'],  1e-12)
        assert_allclose(rays[2], self.data['rays_neg'],  1e-12)

        #self.update_refs(ray_times = rays[0], rays_pos = rays[1], rays_neg = rays[2])

    def test_null_coordinates(self):
        x = self.data['x']
        idx = x.shape[0] // 2
        gxx = self.data['gxx']
        uv = 0.5 * array_utils.array_reflect(invars.dist_proper(x[idx:], gxx[0, idx:]), -1)

        coords = doublenull.null_coordinates(self.data['t'], x, uv, uv,
                                             gxx, np.zeros_like(gxx),
                                             -(self.data['alpha'] ** 2))
        assert_allclose(coords[0], self.data['uxt'], 1e-12)
        assert_allclose(coords[1], self.data['vxt'], 1e-12)

        #self.update_refs(uxt = coords[0], vxt = coords[1])

    def test_null_coordinates_inv(self):
        t = self.data['t']
        x = self.data['x']
        idx = x.shape[0] // 2

        gxx   = self.data['gxx']
        alpha = self.data['alpha']

        # proper time τ(t)
        tau_t = cumulative_trapezoid(alpha[:, idx], t, initial = 0.0)

        # uniform grid in τ
        tau_uniform = np.linspace(tau_t[0], tau_t[-1], t.shape[0])

        # t on the uniform τ grid
        t_tau = interp1d(tau_t, t)(tau_uniform)

        Tuv, Xuv = doublenull.null_coordinates_inv(t, x, tau_uniform,
                                                   gxx, None, -(alpha ** 2),
                                                   uv_times = t_tau)
        assert_allclose(Tuv, self.data['Tuv'], 1e-12)
        assert_allclose(Xuv, self.data['Xuv'], 1e-12)

        #self.update_refs(Tuv = Tuv, Xuv = Xuv)