"""Tests for TDOA gradient computation.

Tests verify that tdoa_with_gradient returns correct gradients
for use in optimization and least-squares fitting algorithms.
"""

import numpy as np

from gri_utils.constants import C
from gri_utils.observables import get_tdoa, tdoa_with_gradient

# Gradient correctness via finite differences


def test_gradient_matches_finite_differences() -> None:
    """Verify analytical gradient matches numerical finite differences."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    _, gradient = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

    # Compute numerical gradient using central differences
    eps = 1e-6
    numerical_gradient = np.zeros(3)

    for j in range(3):
        emit_plus = emit_xyz.copy()
        emit_plus[j] += eps
        emit_minus = emit_xyz.copy()
        emit_minus[j] -= eps

        tdoa_plus = get_tdoa(emit_plus, c1_xyz, c2_xyz)
        tdoa_minus = get_tdoa(emit_minus, c1_xyz, c2_xyz)

        numerical_gradient[j] = (tdoa_plus - tdoa_minus) / (2 * eps)

    assert np.allclose(gradient, numerical_gradient, atol=1e-12)


def test_gradient_at_different_positions() -> None:
    """Verify gradient at multiple emitter positions."""
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    test_emitters = [
        np.array([6378137.0, 0.0, 0.0]),  # Equidistant
        np.array([6378137.0, 500.0, 0.0]),  # Closer to c2
        np.array([6378137.0, -500.0, 0.0]),  # Closer to c1
        np.array([6378137.0, 0.0, 1000.0]),  # Above midpoint
        np.array([6388137.0, 300.0, 400.0]),  # 3D offset
    ]

    eps = 1e-6
    for emit_xyz in test_emitters:
        _, gradient = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

        numerical_gradient = np.zeros(3)
        for j in range(3):
            emit_plus = emit_xyz.copy()
            emit_plus[j] += eps
            emit_minus = emit_xyz.copy()
            emit_minus[j] -= eps

            tdoa_plus = get_tdoa(emit_plus, c1_xyz, c2_xyz)
            tdoa_minus = get_tdoa(emit_minus, c1_xyz, c2_xyz)

            numerical_gradient[j] = (tdoa_plus - tdoa_minus) / (2 * eps)

        assert np.allclose(gradient, numerical_gradient, atol=1e-12)


# TDOA value consistency


def test_tdoa_matches_non_gradient_version() -> None:
    """Verify gradient function returns same TDOA as non-gradient version."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    tdoa_simple = get_tdoa(emit_xyz, c1_xyz, c2_xyz)
    tdoa_grad, _ = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

    assert np.isclose(tdoa_simple, tdoa_grad, atol=1e-15)


def test_gradient_equals_unit_vector_diff_over_c() -> None:
    """Verify TDOA gradient equals (unit1 - unit2) / c."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    _, tdoa_gradient = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

    diff1 = emit_xyz - c1_xyz
    diff2 = emit_xyz - c2_xyz
    unit1 = diff1 / np.linalg.norm(diff1)
    unit2 = diff2 / np.linalg.norm(diff2)
    expected = (unit1 - unit2) / C

    assert np.allclose(tdoa_gradient, expected, atol=1e-15)


# Gradient units


def test_gradient_units_seconds_per_meter() -> None:
    """Verify gradient has correct units (seconds/meter)."""
    emit_xyz = np.array([6378137.0, 500.0, 0.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    _, gradient = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

    # Gradient magnitude is around 1/c for unit vectors in same direction.
    # Max gradient magnitude is 2/c when unit vectors are opposite.
    max_gradient_magnitude = 2.0 / C
    assert np.linalg.norm(gradient) <= max_gradient_magnitude + 1e-15


# Batch processing


def test_batch_gradient_for_multiple_emitters() -> None:
    """Batch gradient computation for least-squares fitting scenario."""
    emits_xyz = np.array(
        [
            [6378137.0, 0.0, 0.0],
            [6378137.0, 500.0, 0.0],
            [6378137.0, 0.0, 1000.0],
        ],
    )
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])

    tdoas, gradients = tdoa_with_gradient(emits_xyz, c1_xyz, c2_xyz)

    assert tdoas.shape == (3,)
    assert gradients.shape == (3, 3)

    # Verify each gradient individually
    eps = 1e-6
    for i in range(3):
        numerical_gradient = np.zeros(3)
        for j in range(3):
            emits_plus = emits_xyz.copy()
            emits_plus[i, j] += eps
            emits_minus = emits_xyz.copy()
            emits_minus[i, j] -= eps

            tdoa_plus = get_tdoa(emits_plus, c1_xyz, c2_xyz)
            tdoa_minus = get_tdoa(emits_minus, c1_xyz, c2_xyz)

            numerical_gradient[j] = (tdoa_plus[i] - tdoa_minus[i]) / (2 * eps)

        assert np.allclose(gradients[i], numerical_gradient, atol=1e-12)


# Input format flexibility


def test_accepts_lists() -> None:
    """Verify function accepts Python lists."""
    emit_xyz = [6378137.0, 500.0, 300.0]
    c1_xyz = [6378137.0, -1000.0, 0.0]
    c2_xyz = [6378137.0, 1000.0, 0.0]

    tdoa, gradient = tdoa_with_gradient(emit_xyz, c1_xyz, c2_xyz)

    assert isinstance(tdoa, np.ndarray)
    assert isinstance(gradient, np.ndarray)
    assert gradient.shape == (3,)
