"""Tests for PDOA gradient computation.

Tests verify that pdoa_with_gradient returns correct Jacobian matrices
for use in optimization and least-squares fitting algorithms.
"""

import numpy as np
from scipy.spatial.transform import Rotation

from gri_utils.constants import C
from gri_utils.conversion import xyz_enu_rotation
from gri_utils.observables.pdoa import get_pdoa
from gri_utils.observables.pdoa_gradient import (
    pdoa_with_gradient,
    pdoa_with_gradient_broadcast,
)

# Gradient correctness via finite differences


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

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array(
        [
            [0.025, 0.0, 0.0],
            [0.0, 0.025, 0.0],
            [0.025, 0.025, 0.0],
        ],
    )

    _, gradient = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    # Compute numerical gradient using central differences
    eps = 1e-6
    n_elements = element_positions.shape[0]
    numerical_gradient = np.zeros((n_elements, 3))

    for j in range(3):
        source_plus = source_xyz.copy()
        source_plus[j] += eps
        source_minus = source_xyz.copy()
        source_minus[j] -= eps

        pdoa_plus = get_pdoa(
            collector_xyz,
            source_plus,
            quat,
            element_positions,
            frequency_hz,
        )
        pdoa_minus = get_pdoa(
            collector_xyz,
            source_minus,
            quat,
            element_positions,
            frequency_hz,
        )

        numerical_gradient[:, j] = (pdoa_plus - pdoa_minus) / (2 * eps)

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


def test_gradient_at_different_source_positions() -> None:
    """Verify gradient at multiple source positions."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array(
        [
            [0.025, 0.0, 0.0],
            [0.0, 0.025, 0.0],
        ],
    )

    # Test at various positions
    test_sources = [
        np.array([6378137.0, 5000.0, 0.0]),  # East
        np.array([6378137.0, 0.0, 5000.0]),  # North
        np.array([6388137.0, 0.0, 0.0]),  # Up
        np.array([6378137.0, 3000.0, 4000.0]),  # Diagonal
    ]

    eps = 1e-6
    for source_xyz in test_sources:
        _, gradient = pdoa_with_gradient(
            collector_xyz,
            source_xyz,
            quat,
            element_positions,
            frequency_hz,
        )

        n_elements = element_positions.shape[0]
        numerical_gradient = np.zeros((n_elements, 3))
        for j in range(3):
            source_plus = source_xyz.copy()
            source_plus[j] += eps
            source_minus = source_xyz.copy()
            source_minus[j] -= eps

            pdoa_plus = get_pdoa(
                collector_xyz,
                source_plus,
                quat,
                element_positions,
                frequency_hz,
            )
            pdoa_minus = get_pdoa(
                collector_xyz,
                source_minus,
                quat,
                element_positions,
                frequency_hz,
            )

            numerical_gradient[:, j] = (pdoa_plus - pdoa_minus) / (2 * eps)

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


def test_gradient_with_identity_rotation() -> None:
    """Verify gradient with identity rotation (simpler geometry)."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 1000.0, 1000.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    frequency_hz = C / 1.0
    element_positions = np.array(
        [
            [0.5, 0.0, 0.0],
            [0.0, 0.5, 0.0],
            [0.0, 0.0, 0.5],
        ],
    )

    _, gradient = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    eps = 1e-6
    n_elements = element_positions.shape[0]
    numerical_gradient = np.zeros((n_elements, 3))

    for j in range(3):
        source_plus = source_xyz.copy()
        source_plus[j] += eps
        source_minus = source_xyz.copy()
        source_minus[j] -= eps

        pdoa_plus = get_pdoa(
            collector_xyz,
            source_plus,
            quat,
            element_positions,
            frequency_hz,
        )
        pdoa_minus = get_pdoa(
            collector_xyz,
            source_minus,
            quat,
            element_positions,
            frequency_hz,
        )

        numerical_gradient[:, j] = (pdoa_plus - pdoa_minus) / (2 * eps)

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


# PDOA value consistency


def test_pdoa_matches_non_gradient_version() -> None:
    """Verify gradient function returns same PDOA as non-gradient version."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6378137.0, 1000.0, 500.0])

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array([[0.025, 0.0, 0.0], [0.0, 0.025, 0.0]])

    pdoa_simple = get_pdoa(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )
    pdoa_grad, _ = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert np.allclose(pdoa_simple, pdoa_grad, atol=1e-10)


# Batch processing for optimization


def test_batch_gradient_for_multiple_sources() -> None:
    """Batch gradient computation for least-squares fitting scenario."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # Multiple sources
    sources_xyz = np.array(
        [
            [6378137.0, 1000.0, 0.0],
            [6378137.0, 0.0, 1000.0],
            [6388137.0, 500.0, 500.0],
        ],
    )

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array([[0.025, 0.0, 0.0], [0.0, 0.025, 0.0]])

    pdoas, gradients = pdoa_with_gradient(
        collector_xyz,
        sources_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    n_sources = sources_xyz.shape[0]
    n_elements = element_positions.shape[0]

    assert pdoas.shape == (n_sources, n_elements)
    assert gradients.shape == (n_sources, n_elements, 3)

    # Verify each gradient individually
    eps = 1e-6
    for i in range(n_sources):
        numerical_gradient = np.zeros((n_elements, 3))
        for j in range(3):
            source_plus = sources_xyz[i].copy()
            source_plus[j] += eps
            source_minus = sources_xyz[i].copy()
            source_minus[j] -= eps

            pdoa_plus = get_pdoa(
                collector_xyz,
                source_plus,
                quat,
                element_positions,
                frequency_hz,
            )
            pdoa_minus = get_pdoa(
                collector_xyz,
                source_minus,
                quat,
                element_positions,
                frequency_hz,
            )

            numerical_gradient[:, j] = (pdoa_plus - pdoa_minus) / (2 * eps)

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


# Broadcast functionality


def test_broadcast_gradient_single_inputs() -> None:
    """Broadcast wrapper returns same result as non-broadcast for single inputs."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6378137.0, 1000.0, 500.0])

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array([[0.025, 0.0, 0.0], [0.0, 0.025, 0.0]])

    pdoa_ref, grad_ref = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )
    pdoa_bc, grad_bc = pdoa_with_gradient_broadcast(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert pdoa_bc.shape == (2,)
    assert grad_bc.shape == (2, 3)
    assert np.allclose(pdoa_bc, pdoa_ref, atol=1e-10)
    assert np.allclose(grad_bc, grad_ref, atol=1e-10)


def test_broadcast_gradient_multidimensional() -> None:
    """Broadcast gradient over multidimensional source array."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # 2x3 grid of sources
    sources_xyz = np.zeros((2, 3, 3))
    sources_xyz[..., 0] = collector_xyz[0] + 10000.0
    sources_xyz[..., 1] = np.arange(6).reshape(2, 3) * 200
    sources_xyz[..., 2] = 500.0

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array([[0.025, 0.0, 0.0], [0.0, 0.025, 0.0]])

    pdoas, gradients = pdoa_with_gradient_broadcast(
        collector_xyz,
        sources_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert pdoas.shape == (2, 3, 2)
    assert gradients.shape == (2, 3, 2, 3)


# Input format flexibility


def test_accepts_scipy_rotation_object() -> None:
    """Verify scipy Rotation objects work directly."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6378137.0, 1000.0, 500.0])

    rot_matrix = xyz_enu_rotation(collector_xyz)
    rotation = Rotation.from_matrix(rot_matrix)
    quat = rotation.as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array([[0.025, 0.0, 0.0]])

    pdoa_rot, grad_rot = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        rotation,
        element_positions,
        frequency_hz,
    )
    pdoa_quat, grad_quat = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert np.allclose(pdoa_rot, pdoa_quat, atol=1e-10)
    assert np.allclose(grad_rot, grad_quat, atol=1e-10)


def test_accepts_list_inputs() -> None:
    """Verify list inputs are accepted."""
    collector_xyz = [6378137.0, 0.0, 0.0]
    source_xyz = [6388137.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    element_positions = [[0.5, 0.0, 0.0]]
    frequency_hz = C / 1.0

    pdoa, gradient = pdoa_with_gradient(
        collector_xyz,
        source_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

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