"""Convert PDOA measurements to AOA direction cosines.

Solves the inverse problem: given phase differences and array geometry,
estimate the signal direction of arrival.

For the forward problem (direction to phases), see pdoa.py.

Frame Convention:
    The returned direction cosines are in the array frame, which is the same
    frame used by get_aoa_quat() and get_pdoa(). Element positions must be defined
    in this same array frame.

Ambiguity Handling:
    For element spacing > lambda/2, phase measurements are ambiguous (wrap modulo
    2*pi). The function enumerates all possible phase unwrappings, solves the
    least-squares problem for each, and keeps only solutions that:
    1. Are physically valid (direction cosines in [-1, 1], sum of squares <= 1)
    2. Are consistent with the measurements (low residual error)

    For spacing d where lambda/2 < d < 3*lambda/2, there is exactly one wrap
    ambiguity per baseline, but the overdetermined system constrains the solution
    to exactly 2 valid AOAs (true direction and one grating lobe).
"""

from __future__ import annotations

from itertools import product
from typing import TYPE_CHECKING

import numpy as np

from gri_utils.constants import C

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def pdoa_to_aoa(
    pdoa: NDArray[np.floating] | Sequence,
    element_positions: NDArray[np.floating] | Sequence,
    frequency_hz: float,
) -> NDArray[np.floating]:
    """Convert phase differences to AOA direction cosines.

    Solves for the direction unit vector u that satisfies:
        pdoa = (2*pi*f/c) * element_positions @ u

    This is a least-squares problem for the direction cosines. The function
    detects the array geometry and returns the direction cosines corresponding
    to axes that have element variation.

    For ambiguous arrays (spacing > lambda/2), the function enumerates all
    valid phase unwrappings and returns only those that produce consistent
    solutions across all baselines. The overdetermined system naturally
    filters out invalid wrap combinations.

    Args:
        pdoa: Phase differences in radians. Shape (N,) for N elements.
        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.

    Returns:
        AOA direction cosines. Shape (K, 2) where K is the number of valid
        candidates. For unambiguous arrays (spacing < lambda/2), K=1.
        For single-wrap ambiguous arrays, typically K=2 (true and grating lobe).

        For planar arrays, returns the two direction cosines for axes with
        non-zero element positions (e.g., X-Y for a horizontal array).
        For full 3D arrays, returns X and Y components. The third component
        can be recovered as sqrt(1 - cx^2 - cy^2) for sources in the
        forward hemisphere.

        Candidates are sorted by sum of squared direction cosines (smallest
        first, i.e., closest to boresight).

    Raises:
        ValueError: If element_positions has fewer than 2 rows (need at least
            2 elements for a unique 2D direction estimate).
        ValueError: If array geometry is degenerate (fewer than 2 axes have
            non-zero element positions).

    Example:
        # Unambiguous array (d = lambda/4) - returns single solution
        aoas = pdoa_to_aoa(pdoa, elements, freq)  # Shape (1, 2)
        aoa = aoas[0]

        # Ambiguous array (d = lambda) - returns 2 candidates
        aoas = pdoa_to_aoa(pdoa, elements, freq)  # Shape (2, 2)
        # Use external information to select correct candidate
    """
    pdoa = np.asarray(pdoa)
    element_positions = np.asarray(element_positions)

    # Check minimum element count
    min_elements = 2
    n_elements = element_positions.shape[0]
    if n_elements < min_elements:
        raise ValueError(
            f"Need at least 2 elements for direction estimation, got {n_elements}",
        )

    wavelength = C / frequency_hz

    # Detect array geometry: which axes have element positions?
    col_has_variation = np.any(element_positions != 0, axis=0)
    active_cols = np.where(col_has_variation)[0]
    n_dims = len(active_cols)

    if n_dims < 2:  # noqa: PLR2004
        msg = (
            f"Array geometry is degenerate: only {n_dims} axis has "
            "variation. Need at least 2 axes with non-zero element positions."
        )
        raise ValueError(msg)

    # Compute max wraps per element from baseline length
    baseline_lengths = np.linalg.norm(element_positions, axis=1)
    per_element_max_wraps = np.floor(baseline_lengths / wavelength).astype(int)

    # If no element has wraps, return single solution
    if np.all(per_element_max_wraps == 0):
        aoa = _solve_aoa(pdoa, element_positions, frequency_hz, active_cols)
        return aoa.reshape(1, 2)

    # Enumerate all phase unwrappings
    wrap_ranges = [range(-mw, mw + 1) for mw in per_element_max_wraps]
    wrap_combinations = list(product(*wrap_ranges))

    # For each unwrapping, solve and keep valid solutions with low residual
    candidates = []
    residuals = []

    for wraps in wrap_combinations:
        # Unwrap phases
        unwrapped_pdoa = pdoa + 2 * np.pi * np.array(wraps)

        # Solve for direction cosines
        aoa, residual = _solve_aoa_with_residual(
            unwrapped_pdoa,
            element_positions,
            frequency_hz,
            active_cols,
        )

        # Check physical validity
        if np.all(np.abs(aoa) <= 1.0) and np.sum(aoa**2) <= 1.0:
            candidates.append(aoa)
            residuals.append(residual)

    if not candidates:
        # Fallback: return principal solution
        aoa = _solve_aoa(pdoa, element_positions, frequency_hz, active_cols)
        return aoa.reshape(1, 2)

    # Filter by residual: keep only solutions with residual near the minimum
    # This filters out wrap combinations that don't satisfy all baselines
    candidates_array = np.array(candidates)
    residuals_array = np.array(residuals)

    min_residual = np.min(residuals_array)
    # Allow residuals within 10x of minimum (accounts for numerical precision)
    residual_threshold = max(min_residual * 10, 1e-10)
    valid_mask = residuals_array <= residual_threshold

    valid_candidates = candidates_array[valid_mask]

    # Remove near-duplicates
    unique_candidates = _remove_duplicate_aoas(valid_candidates)

    # Sort by distance from boresight (sum of squared cosines)
    distances = np.sum(unique_candidates**2, axis=1)
    sort_idx = np.argsort(distances)

    return unique_candidates[sort_idx]


def _solve_aoa(
    pdoa: NDArray[np.floating],
    element_positions: NDArray[np.floating],
    frequency_hz: float,
    active_cols: NDArray[np.integer] | None = None,
) -> NDArray[np.floating]:
    """Solve for AOA direction cosines (internal helper).

    Args:
        pdoa: Phase differences in radians. Shape (N,).
        element_positions: Element positions in array frame. Shape (N, 3).
        frequency_hz: Signal frequency in Hz.
        active_cols: Pre-computed active column indices. If None, will be
            computed from element_positions.

    Returns:
        AOA direction cosines. Shape (2,).

    Raises:
        ValueError: If array geometry is degenerate.
    """
    aoa, _residual = _solve_aoa_with_residual(
        pdoa,
        element_positions,
        frequency_hz,
        active_cols,
    )
    return aoa


def _solve_aoa_with_residual(
    pdoa: NDArray[np.floating],
    element_positions: NDArray[np.floating],
    frequency_hz: float,
    active_cols: NDArray[np.integer] | None = None,
) -> tuple[NDArray[np.floating], float]:
    """Solve for AOA direction cosines and return residual.

    Args:
        pdoa: Phase differences in radians. Shape (N,).
        element_positions: Element positions in array frame. Shape (N, 3).
        frequency_hz: Signal frequency in Hz.
        active_cols: Pre-computed active column indices. If None, will be
            computed from element_positions.

    Returns:
        Tuple of (aoa, residual) where aoa is shape (2,) and residual is the
        sum of squared errors from the least-squares fit.

    Raises:
        ValueError: If array geometry is degenerate.
    """
    min_elements = 2
    n_elements = element_positions.shape[0]
    if n_elements < min_elements:
        raise ValueError(
            f"Need at least 2 elements for direction estimation, got {n_elements}",
        )

    # Convert phase differences to path length differences
    # pdoa = (2*pi*f/c) * P @ u  =>  P @ u = pdoa * c / (2*pi*f)
    scale = C / (2.0 * np.pi * frequency_hz)
    path_diffs = pdoa * scale

    # Detect which columns have variation (non-zero elements)
    if active_cols is None:
        col_has_variation = np.any(element_positions != 0, axis=0)
        active_cols = np.where(col_has_variation)[0]

    if len(active_cols) < 2:  # noqa: PLR2004
        msg = (
            f"Array geometry is degenerate: only {len(active_cols)} axis has "
            "variation. Need at least 2 axes with non-zero element positions."
        )
        raise ValueError(msg)

    if len(active_cols) == 2:  # noqa: PLR2004
        # Planar array: solve reduced 2D system for better conditioning
        positions_2d = element_positions[:, active_cols]
        u_2d, residuals, _rank, _s = np.linalg.lstsq(
            positions_2d,
            path_diffs,
            rcond=None,
        )
        # Compute residual if not returned (happens when system is not overdetermined)
        if len(residuals) == 0:
            predicted = positions_2d @ u_2d
            residual = float(np.sum((path_diffs - predicted) ** 2))
        else:
            residual = float(residuals[0])
        return u_2d, residual

    # Full 3D array: solve full system, return first two components
    # (assumes standard X-Y-Z ordering where Z is the "up" direction)
    u_raw, residuals, _rank, _s = np.linalg.lstsq(
        element_positions,
        path_diffs,
        rcond=None,
    )
    if len(residuals) == 0:
        predicted = element_positions @ u_raw
        residual = float(np.sum((path_diffs - predicted) ** 2))
    else:
        residual = float(residuals[0])
    return u_raw[:2], residual


def _remove_duplicate_aoas(
    candidates: NDArray[np.floating],
    tol: float = 1e-6,
) -> NDArray[np.floating]:
    """Remove near-duplicate AOA candidates.

    Args:
        candidates: Array of shape (K, 2) containing candidate AOAs.
        tol: Tolerance for considering two AOAs as duplicates.

    Returns:
        Array of unique candidates, shape (M, 2) where M <= K.
    """
    if len(candidates) == 0:
        return candidates

    unique = [candidates[0]]
    for candidate in candidates[1:]:
        is_duplicate = False
        for existing in unique:
            if np.allclose(candidate, existing, atol=tol):
                is_duplicate = True
                break
        if not is_duplicate:
            unique.append(candidate)

    return np.array(unique)
