"""Phase Difference of Arrival (PDOA) observable computation.

PDOA measures phase differences between antenna elements in an interferometer
array. For a signal arriving from direction unit vector u (in array frame),
the phase difference at element k relative to the reference element is:

    phi_k = (2*pi*f / c) * dot(element_k_position, u)

where element_k_position is in the array coordinate frame (reference at origin),
f is signal frequency, and c is the speed of light.

This is mathematically related to AOA - both encode direction information, but
PDOA provides N scalar measurements from N phase differences.

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.

Note:
    Phase ambiguity (phi + 2*pi*n) is inherent to PDOA. For element spacing
    > wavelength/2, multiple signal directions produce the same phase differences
    modulo 2*pi. Ambiguity handling is left to the caller.
"""

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
from .aoa_quat_to_xyz import aoa_quat_to_xyz
from .pdoa_to_aoa import pdoa_to_aoa
from .rotation_utils import rotation_to_quat_array

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray


def get_pdoa(
    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,
) -> NDArray[np.floating]:
    """Compute Phase Difference of Arrival for interferometer array elements.

    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:
        Phase differences in radians. Shape (N,) for single source,
        (M, N) for M sources. Raw values, not wrapped to [-pi, pi].

    Note:
        Phase ambiguity (phi + 2*pi*n) is not handled. For element spacing
        > wavelength/2, multiple signal directions may produce the same
        phase differences modulo 2*pi.
    """
    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))
    )

    # Unit vector from collector to source in ECEF
    _, direction_ecef = _range_and_unit(collector_xyz, source_xyz)

    # Rotate to array frame
    direction_array = rotation.apply(direction_ecef)

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

    if direction_array.ndim == 1:
        # Single source: (N,) = (N, 3) @ (3,)
        pdoa = scale * (element_positions @ direction_array)
    else:
        # Multiple sources: (M, N) = (M, 3) @ (N, 3).T
        pdoa = scale * (direction_array @ element_positions.T)

    return np.asarray(pdoa)


def get_pdoa_residual(
    collector_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
    element_positions: NDArray[np.floating] | Sequence,
    frequency_hz: float,
    pdoa_radians: NDArray[np.floating] | Sequence,
) -> Callable[[NDArray[np.floating]], NDArray[np.floating]]:
    """Create a vectorized PDOA angular residual function.

    PDOA is fundamentally an AOA measurement. This converts the measured phase
    differences to a bearing direction, then returns a residual function that
    computes angular deviation from that bearing for candidate emitter positions.
    The zero-crossing defines the PDOA bearing line.

    For ambiguous arrays (element spacing > wavelength/2), the principal solution
    (closest to boresight) is used.

    Args:
        collector_xyz: Reference element position in ECEF meters. Shape (3,).
        collector_rotation: Rotation from ECEF to array frame. Quaternion
            [x, y, z, w] (scipy convention) or scipy Rotation object.
        element_positions: Non-reference element positions in array frame,
            in meters. Shape (N, 3). Reference element is at origin.
        frequency_hz: Signal frequency in Hz.
        pdoa_radians: Measured phase differences in radians. Shape (N,).

    Returns:
        Vectorized residual function: takes points (M, 3) or (3,), returns
        angular deviation in radians (M,) or (). Zero along the bearing line.
    """
    col = np.asarray(collector_xyz, dtype=np.float64)

    # Convert PDOA to AOA direction cosines (principal solution)
    aoa_candidates = pdoa_to_aoa(
        np.asarray(pdoa_radians),
        np.asarray(element_positions),
        frequency_hz,
    )
    aoa = aoa_candidates[0]

    # Convert array-frame AOA to ECEF bearing direction
    bearing_ecef = aoa_quat_to_xyz(aoa, collector_rotation)
    bearing_ecef = bearing_ecef / np.linalg.norm(bearing_ecef)

    def residual(points: NDArray[np.floating]) -> NDArray[np.floating]:
        pts = np.asarray(points)
        vec = pts - col
        dist = np.linalg.norm(vec, axis=-1, keepdims=True)
        unit_vec = vec / np.maximum(dist, 1e-12)

        cos_angle = np.sum(unit_vec * bearing_ecef, axis=-1)
        cos_angle = np.clip(cos_angle, -1.0, 1.0)
        return np.asarray(np.arccos(cos_angle))

    return residual


def get_pdoa_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,
) -> NDArray[np.floating]:
    """Broadcast-capable wrapper for get_pdoa.

    Follows numpy broadcasting semantics for batch dimensions. Inputs with shape
    (..., M) are broadcast together on the (...) dimensions, then flattened for
    computation and reshaped back.

    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:
        Phase differences in radians. Shape (..., N) matching the broadcast
        batch shape of inputs. Single inputs (shape (M,)) return shape (N,).
    """
    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 (everything except last dim)
    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
    # Note: get_pdoa handles single source but not batch collector/rotation
    # So we loop over batch dimension
    results = np.empty((n_batch, n_elements))
    for i in range(n_batch):
        results[i] = get_pdoa(
            collector_flat[i],
            source_flat[i],
            quat_flat[i],
            element_positions,
            frequency_hz,
        )

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