"""Tests for FDOA gradient computation.

Tests verify that fdoa_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 fdoa_with_gradient, get_fdoa

# Gradient correctness via finite differences


def test_position_gradient_matches_finite_differences() -> None:
    """Verify analytical position gradient matches numerical finite differences."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    emit_vel = np.array([50.0, -100.0, 25.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    _, grad_pos, _ = fdoa_with_gradient(emit_xyz, emit_vel, c1_xyz, c2_xyz, freq_hz)

    # 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

        fdoa_plus = get_fdoa(emit_plus, emit_vel, c1_xyz, c2_xyz, freq_hz)
        fdoa_minus = get_fdoa(emit_minus, emit_vel, c1_xyz, c2_xyz, freq_hz)

        numerical_gradient[j] = (fdoa_plus - fdoa_minus) / (2 * eps)

    assert np.allclose(grad_pos, numerical_gradient, rtol=1e-3, atol=1e-6)


def test_velocity_gradient_matches_finite_differences() -> None:
    """Verify analytical velocity gradient matches numerical finite differences."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    emit_vel = np.array([50.0, -100.0, 25.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    _, _, grad_vel = fdoa_with_gradient(emit_xyz, emit_vel, c1_xyz, c2_xyz, freq_hz)

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

    for j in range(3):
        vel_plus = emit_vel.copy()
        vel_plus[j] += eps
        vel_minus = emit_vel.copy()
        vel_minus[j] -= eps

        fdoa_plus = get_fdoa(emit_xyz, vel_plus, c1_xyz, c2_xyz, freq_hz)
        fdoa_minus = get_fdoa(emit_xyz, vel_minus, c1_xyz, c2_xyz, freq_hz)

        numerical_gradient[j] = (fdoa_plus - fdoa_minus) / (2 * eps)

    assert np.allclose(grad_vel, numerical_gradient, rtol=1e-3, atol=1e-6)


def test_gradients_at_different_positions() -> None:
    """Verify gradients at multiple emitter positions and velocities."""
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    test_cases = [
        # (emit_xyz, emit_vel)
        (np.array([6378137.0, 0.0, 0.0]), np.array([0.0, -100.0, 0.0])),
        (np.array([6378137.0, 500.0, 0.0]), np.array([100.0, 0.0, 0.0])),
        (np.array([6378137.0, -500.0, 0.0]), np.array([0.0, 0.0, 100.0])),
        (np.array([6378137.0, 0.0, 1000.0]), np.array([50.0, 50.0, 50.0])),
        (np.array([6388137.0, 300.0, 400.0]), np.array([-30.0, 80.0, -20.0])),
    ]

    eps = 1e-6
    for emit_xyz, emit_vel in test_cases:
        _, grad_pos, grad_vel = fdoa_with_gradient(
            emit_xyz,
            emit_vel,
            c1_xyz,
            c2_xyz,
            freq_hz,
        )

        # Check position gradient
        numerical_grad_pos = 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

            fdoa_plus = get_fdoa(emit_plus, emit_vel, c1_xyz, c2_xyz, freq_hz)
            fdoa_minus = get_fdoa(emit_minus, emit_vel, c1_xyz, c2_xyz, freq_hz)

            numerical_grad_pos[j] = (fdoa_plus - fdoa_minus) / (2 * eps)

        assert np.allclose(grad_pos, numerical_grad_pos, rtol=1e-3, atol=1e-6)

        # Check velocity gradient
        numerical_grad_vel = np.zeros(3)
        for j in range(3):
            vel_plus = emit_vel.copy()
            vel_plus[j] += eps
            vel_minus = emit_vel.copy()
            vel_minus[j] -= eps

            fdoa_plus = get_fdoa(emit_xyz, vel_plus, c1_xyz, c2_xyz, freq_hz)
            fdoa_minus = get_fdoa(emit_xyz, vel_minus, c1_xyz, c2_xyz, freq_hz)

            numerical_grad_vel[j] = (fdoa_plus - fdoa_minus) / (2 * eps)

        assert np.allclose(grad_vel, numerical_grad_vel, rtol=1e-3, atol=1e-6)


# FDOA value consistency


def test_fdoa_matches_non_gradient_version() -> None:
    """Verify gradient function returns same FDOA as non-gradient version."""
    emit_xyz = np.array([6378137.0, 500.0, 300.0])
    emit_vel = np.array([50.0, -100.0, 25.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    fdoa_simple = get_fdoa(emit_xyz, emit_vel, c1_xyz, c2_xyz, freq_hz)
    fdoa_grad, _, _ = fdoa_with_gradient(emit_xyz, emit_vel, c1_xyz, c2_xyz, freq_hz)

    assert np.isclose(fdoa_simple, fdoa_grad, atol=1e-15)


def test_velocity_gradient_is_unit_vector_difference() -> None:
    """Verify velocity gradient equals f0/c * (unit1 - unit2)."""
    emit_xyz = np.array([6378137.0, 0.0, 0.0])
    emit_vel = np.array([50.0, -100.0, 25.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    _, _, grad_vel = fdoa_with_gradient(emit_xyz, emit_vel, c1_xyz, c2_xyz, freq_hz)

    # Compute expected: f0/c * (unit1 - unit2)
    diff1 = c1_xyz - emit_xyz
    diff2 = c2_xyz - emit_xyz
    unit1 = diff1 / np.linalg.norm(diff1)
    unit2 = diff2 / np.linalg.norm(diff2)
    expected_grad_vel = (freq_hz / C) * (unit1 - unit2)

    assert np.allclose(grad_vel, expected_grad_vel, atol=1e-15)


# Gradient units


def test_gradient_units() -> None:
    """Verify gradients have correct units."""
    emit_xyz = np.array([6378137.0, 500.0, 0.0])
    emit_vel = np.array([0.0, -100.0, 0.0])
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    _, grad_pos, grad_vel = fdoa_with_gradient(
        emit_xyz,
        emit_vel,
        c1_xyz,
        c2_xyz,
        freq_hz,
    )

    # grad_vel has units Hz/(m/s) = s^-1 / (m/s) = 1/m
    # Max magnitude is 2 * f0 / c (when unit vectors are opposite)
    max_grad_vel_magnitude = 2.0 * freq_hz / C
    assert np.linalg.norm(grad_vel) <= max_grad_vel_magnitude + 1e-10

    # grad_pos units are more complex (Hz/m), depends on geometry and velocity
    # Just verify it's finite and reasonable
    assert np.all(np.isfinite(grad_pos))


# 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],
        ],
    )
    emits_vel = np.array(
        [
            [0.0, -100.0, 0.0],
            [50.0, 50.0, 0.0],
            [0.0, 0.0, -50.0],
        ],
    )
    c1_xyz = np.array([6378137.0, -1000.0, 0.0])
    c2_xyz = np.array([6378137.0, 1000.0, 0.0])
    freq_hz = 1e9

    fdoas, grad_pos, grad_vel = fdoa_with_gradient(
        emits_xyz,
        emits_vel,
        c1_xyz,
        c2_xyz,
        freq_hz,
    )

    assert fdoas.shape == (3,)
    assert grad_pos.shape == (3, 3)
    assert grad_vel.shape == (3, 3)

    # Verify each gradient individually via finite differences
    eps = 1e-6
    for i in range(3):
        # Position gradient
        numerical_grad_pos = 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

            fdoa_plus = get_fdoa(emits_plus, emits_vel, c1_xyz, c2_xyz, freq_hz)
            fdoa_minus = get_fdoa(emits_minus, emits_vel, c1_xyz, c2_xyz, freq_hz)

            numerical_grad_pos[j] = (fdoa_plus[i] - fdoa_minus[i]) / (2 * eps)

        assert np.allclose(grad_pos[i], numerical_grad_pos, rtol=1e-3, atol=1e-6)

        # Velocity gradient
        numerical_grad_vel = np.zeros(3)
        for j in range(3):
            vels_plus = emits_vel.copy()
            vels_plus[i, j] += eps
            vels_minus = emits_vel.copy()
            vels_minus[i, j] -= eps

            fdoa_plus = get_fdoa(emits_xyz, vels_plus, c1_xyz, c2_xyz, freq_hz)
            fdoa_minus = get_fdoa(emits_xyz, vels_minus, c1_xyz, c2_xyz, freq_hz)

            numerical_grad_vel[j] = (fdoa_plus[i] - fdoa_minus[i]) / (2 * eps)

        assert np.allclose(grad_vel[i], numerical_grad_vel, rtol=1e-3, atol=1e-6)


# Input format flexibility


def test_accepts_lists() -> None:
    """Verify function accepts Python lists."""
    emit_xyz = [6378137.0, 500.0, 300.0]
    emit_vel = [50.0, -100.0, 25.0]
    c1_xyz = [6378137.0, -1000.0, 0.0]
    c2_xyz = [6378137.0, 1000.0, 0.0]
    freq_hz = 1e9

    fdoa, grad_pos, grad_vel = fdoa_with_gradient(
        emit_xyz,
        emit_vel,
        c1_xyz,
        c2_xyz,
        freq_hz,
    )

    assert isinstance(fdoa, np.ndarray)
    assert isinstance(grad_pos, np.ndarray)
    assert isinstance(grad_vel, np.ndarray)
    assert grad_pos.shape == (3,)
    assert grad_vel.shape == (3,)
