"""Phase Difference of Arrival (PDOA) with gradient computation.

Provides PDOA computation along with Jacobian matrices for use in optimization
and fitting algorithms.

For basic PDOA computation without gradients, see pdoa.py.

Frame Convention:
    The array frame is defined by collector_rotation, which rotates vectors from
    ECEF to the local array frame. This is the same frame convention used by
    get_aoa_quat(). Element positions must be defined in this same array frame.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

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

from gri_utils.constants import C

from ._range_utils import _range_and_unit, _unit_vector_jacobian
from .rotation_utils import rotation_to_quat_array

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def pdoa_with_gradient(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
    element_positions: NDArray[np.floating] | Sequence,
    frequency_hz: float,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Compute PDOA and gradient w.r.t. source position.

    Computes phase differences for each element along with the Jacobian
    matrix of PDOA with respect to the source position.

    Args:
        collector_xyz: Reference element position in ECEF meters. Shape (3,).
        source_xyz: Source position in ECEF meters. Shape (3,) or (M, 3).
        collector_rotation: Rotation from ECEF to array frame. Quaternion
            [x, y, z, w] (scipy convention) or scipy Rotation object. This
            defines the array frame, same convention as get_aoa_quat().
        element_positions: Non-reference element positions in array frame,
            in meters. Shape (N, 3). Reference element is at origin. Must be
            defined in the same frame that collector_rotation rotates into.
        frequency_hz: Signal frequency in Hz.

    Returns:
        Tuple of (pdoa, gradient):
            - pdoa: Phase differences in radians. Shape (N,) for single source,
              (M, N) for M sources.
            - gradient: Jacobian of PDOA w.r.t. source position.
              Shape (N, 3) for single source, (M, N, 3) for M sources.
              gradient[k, j] = d(phi_k)/d(source_xyz[j]).
    """
    collector_xyz = np.asarray(collector_xyz)
    source_xyz = np.asarray(source_xyz)
    element_positions = np.asarray(element_positions)

    rotation: Rotation = (
        collector_rotation
        if isinstance(collector_rotation, Rotation)
        else Rotation.from_quat(np.asarray(collector_rotation))
    )

    # Determine if single point or batch
    single_point = source_xyz.ndim == 1

    # Ensure 2D for consistent processing
    if single_point:
        source_xyz = source_xyz[np.newaxis, :]

    # Get rotation matrix
    rot_matrix = rotation.as_matrix()  # (3, 3)

    # Compute range and unit vector for each source
    # Need to broadcast collector to match sources
    collector_broadcast = np.broadcast_to(collector_xyz, source_xyz.shape)
    range_val, direction_ecef = _range_and_unit(collector_broadcast, source_xyz)

    # Rotate direction to array frame
    direction_array = rotation.apply(direction_ecef)  # (M, 3)

    # Phase difference for each element: (2*pi*f/c) * p_k . u
    scale = 2.0 * np.pi * frequency_hz / C

    # PDOA: (M, N) = (M, 3) @ (N, 3).T
    pdoa = scale * (direction_array @ element_positions.T)

    # Gradient computation:
    # phi_k = (2*pi*f/c) * p_k^T @ R @ u_ecef
    # d(phi_k)/d(source) = (2*pi*f/c) * p_k^T @ R @ d(u_ecef)/d(source)
    #
    # where d(u)/d(source) = (I - u @ u.T) / r

    # Get Jacobian of unit vector in ECEF frame
    j_ecef = _unit_vector_jacobian(direction_ecef, range_val)  # (M, 3, 3)

    # Rotate Jacobian to get d(u_array)/d(source):
    # d(u_array)/d(source) = R @ d(u_ecef)/d(source)
    # j_rotated[m, i, j] = sum_k R[i,k] * j_ecef[m, k, j]
    j_rotated = np.einsum("ik,mkj->mij", rot_matrix, j_ecef)  # (M, 3, 3)

    # Gradient for each element:
    # grad[m, n, j] = scale * sum_i element_positions[n, i] * j_rotated[m, i, j]
    gradient = scale * np.einsum("ni,mij->mnj", element_positions, j_rotated)

    # Return to original shape if single point
    if single_point:
        pdoa = pdoa[0]  # (N,)
        gradient = gradient[0]  # (N, 3)

    return pdoa, gradient


def pdoa_with_gradient_broadcast(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
    element_positions: NDArray[np.floating] | Sequence,
    frequency_hz: float,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Broadcast-capable wrapper for pdoa_with_gradient.

    Follows numpy broadcasting semantics for batch dimensions.

    Args:
        collector_xyz: Reference element position(s) in ECEF meters. Shape (..., 3).
        source_xyz: Source position(s) in ECEF meters. Shape (..., 3).
        collector_rotation: Rotation from ECEF to array frame. Quaternion
            array with shape (..., 4) or scipy Rotation object. This defines
            the array frame, same convention as get_aoa_quat().
        element_positions: Non-reference element positions in array frame,
            in meters. Shape (N, 3). Reference element is at origin. Must be
            defined in the same frame that collector_rotation rotates into.
        frequency_hz: Signal frequency in Hz.

    Returns:
        Tuple of (pdoa, gradient):
            - pdoa: Phase differences. Shape (..., N).
            - gradient: Jacobian. Shape (..., N, 3).
        Single inputs (shape (M,)) return shapes (N,) and (N, 3).
    """
    collector_xyz = np.asarray(collector_xyz)
    source_xyz = np.asarray(source_xyz)
    element_positions = np.asarray(element_positions)
    collector_quat = rotation_to_quat_array(collector_rotation)

    n_elements = element_positions.shape[0]

    # Track if all inputs are 1D (scalar case)
    all_1d = (
        collector_xyz.ndim == 1 and source_xyz.ndim == 1 and collector_quat.ndim == 1
    )

    # Ensure at least 2D for batch dimension handling
    collector_xyz = np.atleast_2d(collector_xyz)
    source_xyz = np.atleast_2d(source_xyz)
    collector_quat = np.atleast_2d(collector_quat)

    # Extract batch shapes
    batch_c = collector_xyz.shape[:-1]
    batch_s = source_xyz.shape[:-1]
    batch_q = collector_quat.shape[:-1]

    # Compute broadcast batch shape
    batch_shape = np.broadcast_shapes(batch_c, batch_s, batch_q)

    # Broadcast to common shape
    collector_xyz = np.broadcast_to(collector_xyz, (*batch_shape, 3))
    source_xyz = np.broadcast_to(source_xyz, (*batch_shape, 3))
    collector_quat = np.broadcast_to(collector_quat, (*batch_shape, 4))

    # Flatten batch dimensions
    n_batch = int(np.prod(batch_shape)) if batch_shape else 1
    collector_flat = collector_xyz.reshape(n_batch, 3)
    source_flat = source_xyz.reshape(n_batch, 3)
    quat_flat = collector_quat.reshape(n_batch, 4)

    # Compute for each batch element
    pdoa_results = np.empty((n_batch, n_elements))
    gradient_results = np.empty((n_batch, n_elements, 3))

    for i in range(n_batch):
        pdoa_i, grad_i = pdoa_with_gradient(
            collector_flat[i],
            source_flat[i],
            quat_flat[i],
            element_positions,
            frequency_hz,
        )
        pdoa_results[i] = pdoa_i
        gradient_results[i] = grad_i

    # Return scalar shape if all inputs were 1D
    if all_1d:
        return pdoa_results.reshape(n_elements), gradient_results.reshape(n_elements, 3)
    return (
        pdoa_results.reshape(*batch_shape, n_elements),
        gradient_results.reshape(*batch_shape, n_elements, 3),
    )
