"""Numerical tests for the dimensionally scaled H2 solvers."""

from __future__ import annotations

import sys
import unittest
from pathlib import Path

import numpy as np


CODE_ROOT = Path(__file__).resolve().parents[1]
if str(CODE_ROOT) not in sys.path:
    sys.path.insert(0, str(CODE_ROOT))

from core.wrench_solver import (  # noqa: E402
    ScaledDLSSolver,
    UndampedSVDSolver,
    WrenchSolveStatus,
)


class WrenchSolverTest(unittest.TestCase):
    def setUp(self) -> None:
        rng = np.random.default_rng(240319)
        self.jacobian = rng.normal(size=(6, 7))
        self.wrench = np.array([8.0, -3.0, 5.0, 0.7, -0.4, 0.2])
        self.length = 0.37
        self.residual = self.jacobian.T @ self.wrench

    def test_scaling_preserves_virtual_work(self) -> None:
        solver = UndampedSVDSolver(self.length)
        scaled_jacobian = solver.scaled_jacobian(self.jacobian)
        scaled_wrench = np.concatenate(
            (self.length * self.wrench[:3], self.wrench[3:])
        )
        np.testing.assert_allclose(
            scaled_jacobian.T @ scaled_wrench,
            self.jacobian.T @ self.wrench,
            rtol=1e-14,
            atol=1e-14,
        )

    def test_undamped_full_rank_recovers_known_wrench(self) -> None:
        result = UndampedSVDSolver(
            self.length, relative_rank_tolerance=1e-10
        ).solve(self.jacobian, self.residual)
        self.assertIs(result.status, WrenchSolveStatus.FULL_RANK)
        self.assertEqual(result.rank, 6)
        np.testing.assert_allclose(
            result.wrench, self.wrench, rtol=2e-13, atol=2e-13
        )
        np.testing.assert_allclose(
            result.reconstructed_residual,
            self.residual,
            rtol=2e-13,
            atol=2e-13,
        )

    def test_dls_matches_scaled_svd_filter_formula(self) -> None:
        damping = 0.08
        solver = ScaledDLSSolver(self.length, damping)
        result = solver.solve(self.jacobian, self.residual)

        scaled_jacobian = solver.scaled_jacobian(self.jacobian)
        u, singular_values, vt = np.linalg.svd(
            scaled_jacobian, full_matrices=False
        )
        expected_scaled = u @ (
            singular_values
            / (singular_values**2 + damping**2)
            * (vt @ self.residual)
        )
        expected_wrench = np.diag(
            [1.0 / self.length] * 3 + [1.0] * 3
        ) @ expected_scaled
        np.testing.assert_allclose(result.wrench, expected_wrench, atol=1e-14)
        np.testing.assert_allclose(
            result.singular_values, singular_values, atol=0.0
        )
        self.assertAlmostEqual(result.damping, damping)

    def test_rank_deficiency_is_retained_and_reported(self) -> None:
        jacobian = self.jacobian.copy()
        jacobian[-1, :] = jacobian[-2, :]
        residual = jacobian.T @ self.wrench
        for solver in (
            UndampedSVDSolver(self.length),
            ScaledDLSSolver(self.length, 0.05),
        ):
            with self.subTest(method=solver.name):
                result = solver.solve(jacobian, residual)
                self.assertIs(
                    result.status, WrenchSolveStatus.RANK_DEFICIENT
                )
                self.assertEqual(result.rank, 5)
                self.assertTrue(np.isinf(result.condition_number))
                self.assertTrue(np.all(np.isfinite(result.wrench)))

    def test_invalid_units_and_shapes_are_rejected(self) -> None:
        with self.assertRaises(ValueError):
            ScaledDLSSolver(0.0, 0.1)
        with self.assertRaises(ValueError):
            ScaledDLSSolver(self.length, 0.0)
        with self.assertRaises(ValueError):
            UndampedSVDSolver(self.length, relative_rank_tolerance=1.0)
        solver = UndampedSVDSolver(self.length)
        with self.assertRaisesRegex(ValueError, "jacobian"):
            solver.solve(np.eye(7), np.ones(7))
        with self.assertRaisesRegex(ValueError, "joint_residual"):
            solver.solve(self.jacobian, np.ones(6))


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