"""Tests for AOA gradient computation.

Tests verify that aoa_quat_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.conversion import xyz_enu_rotation
from gri_utils.observables import (
    aoa_quat_with_gradient,
    aoa_quat_with_gradient_broadcast,
    get_aoa_quat,
)

# Gradient correctness via finite differences


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

    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    _, gradient = aoa_quat_with_gradient(sensor_xyz, target_xyz, sensor_quat)

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

    for j in range(3):
        target_plus = target_xyz.copy()
        target_plus[j] += eps
        target_minus = target_xyz.copy()
        target_minus[j] -= eps

        aoa_plus = get_aoa_quat(sensor_xyz, target_plus, sensor_quat)
        aoa_minus = get_aoa_quat(sensor_xyz, target_minus, sensor_quat)

        numerical_gradient[:, j] = (aoa_plus - aoa_minus) / (2 * eps)

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


def test_gradient_at_different_target_positions() -> None:
    """Verify gradient at multiple target positions."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    # Test at various positions
    test_targets = [
        np.array([6378137.0, 500.0, 0.0]),  # East
        np.array([6378137.0, 0.0, 500.0]),  # North
        np.array([6388137.0, 0.0, 0.0]),  # Up
        np.array([6378137.0, 300.0, 400.0]),  # Diagonal
    ]

    eps = 1e-6
    for target_xyz in test_targets:
        _, gradient = aoa_quat_with_gradient(sensor_xyz, target_xyz, sensor_quat)

        numerical_gradient = np.zeros((2, 3))
        for j in range(3):
            target_plus = target_xyz.copy()
            target_plus[j] += eps
            target_minus = target_xyz.copy()
            target_minus[j] -= eps

            aoa_plus = get_aoa_quat(sensor_xyz, target_plus, sensor_quat)
            aoa_minus = get_aoa_quat(sensor_xyz, target_minus, sensor_quat)

            numerical_gradient[:, j] = (aoa_plus - aoa_minus) / (2 * eps)

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


# AOA value consistency


def test_aoa_matches_non_gradient_version() -> None:
    """Verify gradient function returns same AOA as non-gradient version."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    target_xyz = np.array([6378137.0, 1000.0, 500.0])

    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    aoa_simple = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)
    aoa_grad, _ = aoa_quat_with_gradient(sensor_xyz, target_xyz, sensor_quat)

    assert np.allclose(aoa_simple, aoa_grad, atol=1e-10)


# Batch processing for optimization


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

    # Multiple observations of different targets
    targets_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(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    sensors_xyz = np.tile(sensor_xyz, (3, 1))
    quats = np.tile(sensor_quat, (3, 1))

    aoas, gradients = aoa_quat_with_gradient(sensors_xyz, targets_xyz, quats)

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

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

            aoa_plus = get_aoa_quat(sensors_xyz, targets_plus, quats)
            aoa_minus = get_aoa_quat(sensors_xyz, targets_minus, quats)

            numerical_gradient[:, j] = (aoa_plus[i] - aoa_minus[i]) / (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."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    target_xyz = np.array([6378137.0, 1000.0, 500.0])

    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    aoa_ref, grad_ref = aoa_quat_with_gradient(sensor_xyz, target_xyz, sensor_quat)
    aoa_bc, grad_bc = aoa_quat_with_gradient_broadcast(
        sensor_xyz,
        target_xyz,
        sensor_quat,
    )

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


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

    # 2x3 grid of targets
    targets_xyz = np.zeros((2, 3, 3))
    targets_xyz[..., 0] = sensor_xyz[0]
    targets_xyz[..., 1] = np.arange(6).reshape(2, 3) * 200
    targets_xyz[..., 2] = 500.0

    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    aoas, gradients = aoa_quat_with_gradient_broadcast(
        sensor_xyz,
        targets_xyz,
        sensor_quat,
    )

    assert aoas.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."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    target_xyz = np.array([6378137.0, 1000.0, 500.0])

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

    aoa_rot, grad_rot = aoa_quat_with_gradient(sensor_xyz, target_xyz, rotation)
    aoa_quat_arr, grad_quat = aoa_quat_with_gradient(sensor_xyz, target_xyz, quat)

    assert np.allclose(aoa_rot, aoa_quat_arr, atol=1e-10)
    assert np.allclose(grad_rot, grad_quat, atol=1e-10)
