"""Numerical tests for the TRACE method-specific components (no training data needed)."""
import sys
import unittest
from pathlib import Path

import numpy as np
from scipy.interpolate import PchipInterpolator

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from biochime_regulon_core import (  # noqa: E402
    pchip_slopes, hermite_eval, project_shape_safe_slopes,
    _match_bio_vector_norm, fit_ridge_map, permute_tf_channels,
)


class TestTRACECore(unittest.TestCase):
    def test_pchip_matches_scipy(self):
        rng = np.random.default_rng(17)
        for k in (2, 3, 4):
            x = np.cumsum(np.r_[0.0, rng.uniform(0.1, 0.7, k-1)])
            y = rng.normal(size=(k, 15, 6))
            m = pchip_slopes(x, y)
            ref = PchipInterpolator(x, y, axis=0)
            np.testing.assert_allclose(m, ref.derivative()(x), atol=1e-10)
            t = rng.uniform(x[0], x[-1], len(y[0]))
            pos, vel = hermite_eval(x, y, m, t)
            np.testing.assert_allclose(pos, np.stack([ref(tt)[i] for i, tt in enumerate(t)]), atol=1e-10)
            np.testing.assert_allclose(vel, np.stack([ref.derivative()(tt)[i] for i, tt in enumerate(t)]), atol=1e-10)

    def test_projected_slope_preserves_monotone_intervals(self):
        x = np.array([0.0, 0.4, 1.0])
        z = np.array([[[0., 0.]], [[1., 2.]], [[2., 3.]]])
        raw = np.array([[[20., 20.]], [[-50., 10.]], [[20., -15.]]])
        m = project_shape_safe_slopes(x, z, raw)
        for lo, hi in [(0., 0.4), (0.4, 1.0)]:
            t = np.linspace(lo, hi, 29)
            # evaluate one query time for one trajectory at a time
            for tt in t:
                pos, _ = hermite_eval(x, z, m, np.array([tt]))
                self.assertTrue(np.all(pos[0] >= z[0, 0] - 1e-9))
                self.assertTrue(np.all(pos[0] <= z[-1, 0] + 1e-9))

    def test_ridge_activity_to_state_direction(self):
        rng = np.random.default_rng(23)
        a = rng.normal(size=(200, 7))
        w = rng.normal(size=(7, 4))
        b = rng.normal(size=4)
        z = a @ w + b
        fitted, intercept = fit_ridge_map(a, z, lam=0.0)
        np.testing.assert_allclose(fitted, w, atol=1e-5)
        np.testing.assert_allclose(intercept, b, atol=1e-5)

    def test_norm_clipping_and_permutation(self):
        m_geom = np.ones((3, 12, 5), dtype=float)
        m_bio = np.ones_like(m_geom) * 100.0
        matched = _match_bio_vector_norm(m_geom, m_bio)
        np.testing.assert_allclose(matched, m_bio * 0.25)
        a = {0.0: np.arange(18).reshape(6, 3), 1.0: np.arange(18).reshape(6, 3) + 1}
        perm = permute_tf_channels(a, seed=4)
        self.assertEqual(perm[0.0].shape, a[0.0].shape)
        self.assertEqual(sorted(perm[0.0].T.tolist()), sorted(a[0.0].T.tolist()))


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