import os
from unittest import TestCase

import numpy as np
from   numpy.testing import assert_allclose, assert_equal

from scipy.integrate import cumulative_trapezoid

from math_utils.array_utils import array_reflect as ar
from nr_analysis_axi import null, invars

class TestNull(TestCase):

    def setUp(self):
        datafile = os.path.splitext(__file__)[0] + '.npz'

        self.data = np.load(datafile)
        self.t   = ar(self.data['t'], -1)
        self.x   = ar(self.data['x'], -1)
        self.gxx = ar(ar(self.data['gxx'], axis = 1), axis = 0)
        self.gtt = ar(ar(self.data['gtt'], axis = 1), axis = 0)

    def test_null_geodesics(self):
        g = null.null_geodesics(self.t, self.x, self.gxx, self.gtt, integrate_times = [0.1])[0]

        l = self.data['l']

        assert_allclose(g(l), self.data['g'], 1e-12)

    def test_coord_similarity(self):
        lapse_origin = np.sqrt(-self.data['gtt'][:, 0])
        tau          = ar(cumulative_trapezoid(lapse_origin, self.data['t'], initial = 0.0), -1)
        T_null       = -np.log(5.4 - tau)

        g = null.null_geodesics(self.t, self.x, self.gxx, self.gtt)

        lambda_max = 0.01
        ret = null.coord_similarity(self.t, self.x, T_null, None, None, g = g,
                                    lambda_max = lambda_max, lambda_points = 32)

        # returned T_null should be the same as source in this case
        assert_equal(ret[0], T_null)

        # lambdas are from 0 - lambda_max
        assert_allclose(ret[1][0],         0.0, 1e-12)
        assert_allclose(ret[1][-1], lambda_max, 1e-12)
        # lambdas are equidistant
        assert_allclose(ret[1][1:] - ret[1][:-1], ret[1][1], 1e-12)

        assert_allclose(ret[2][:-10], self.data['x_of_Tnull_lambda'], 1e-12)
        assert_allclose(ret[3][:-10], self.data['t_of_Tnull_lambda'], 1e-12)