"""Confidence interval for sample quantiles using normal approximation to binomial.

Provides confidence intervals for any quantile (e.g., median) of a distribution
by using the normal approximation to the binomial distribution to find the
appropriate order statistics.

Note:
    For 1D data, scipy.stats.quantile_test().confidence_interval() provides
    equivalent functionality with additional statistical testing capabilities.

Example using scipy (1D equivalent):
    >>> from scipy import stats
    >>> result = stats.quantile_test(data, q=0.5, p=0.5)
    >>> ci = result.confidence_interval(confidence_level=0.95)
"""

from __future__ import annotations

from typing import TYPE_CHECKING, overload

import numpy as np
from scipy.stats import norm

if TYPE_CHECKING:
    from numpy.typing import NDArray

# Minimum samples required for reliable quantile CI
_MIN_SAMPLES = 10


@overload
def quantile_ci(
    data: NDArray[np.floating],
    quantile: float = ...,
    confidence: float = ...,
    *,
    axis: None = ...,
    min_samples: int = ...,
) -> tuple[float, float]: ...


@overload
def quantile_ci(
    data: NDArray[np.floating],
    quantile: float = ...,
    confidence: float = ...,
    *,
    axis: int,
    min_samples: int = ...,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]: ...


def quantile_ci(
    data: NDArray[np.floating],
    quantile: float = 0.5,
    confidence: float = 0.95,
    *,
    axis: int | None = None,
    min_samples: int = _MIN_SAMPLES,
) -> tuple[float, float] | tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Calculate confidence interval for a sample quantile.

    Uses normal approximation to the binomial distribution to determine which
    order statistics bound the confidence interval for the specified quantile.

    Args:
        data: Array of sample data. Can be multi-dimensional.
        quantile: The quantile of interest, in (0, 1). Default 0.5 for median.
        confidence: Confidence level in (0, 1). Default 0.95 for 95% CI.
        axis: Axis along which to compute the CI. Default None uses axis 0,
            except for 2D row vectors (shape (1, n)) which use axis 1.
        min_samples: Minimum samples required along the operating axis.
            Returns NaN if fewer samples provided. Default 10.

    Returns:
        Tuple of (lower_bound, upper_bound). For 1D input or axis=None, returns
        scalar floats. For multi-dimensional input with explicit axis, returns
        arrays with the operating axis removed.

    Example:
        >>> data = np.random.randn(100)
        >>> lower, upper = quantile_ci(data)  # Median CI, scalar output

        >>> data_2d = np.random.randn(100, 3)  # 100 samples, 3 variables
        >>> lower, upper = quantile_ci(data_2d, axis=0)  # CI for each variable

        >>> data_3d = np.random.randn(50, 4, 3)
        >>> lower, upper = quantile_ci(data_3d, quantile=0.75, axis=0)
    """
    data = np.asarray(data)

    # Handle 1D case: always return scalars
    if data.ndim <= 1:
        data = data.ravel()
        n = len(data)

        if n < min_samples:
            return np.nan, np.nan

        z = abs(norm.ppf((1 - confidence) / 2))
        spread = z * np.sqrt(quantile * (1 - quantile) * n)
        idx_lower = n * quantile - spread
        idx_upper = n * quantile + spread

        sorted_data = np.sort(data)
        indices = np.arange(1, n + 1)

        ci_lower = float(np.interp(idx_lower, indices, sorted_data))
        ci_upper = float(np.interp(idx_upper, indices, sorted_data))
        return ci_lower, ci_upper

    # Multi-dimensional case
    # Default axis: 0, unless 2D row vector (shape (1, n))
    if axis is None:
        is_row_vector = data.ndim == 2 and data.shape[0] == 1  # noqa: PLR2004
        axis = 1 if is_row_vector else 0

    # Normalize negative axis
    axis = axis % data.ndim

    n = data.shape[axis]

    # Output shape: input shape with operating axis removed
    out_shape = list(data.shape)
    out_shape.pop(axis)
    out_shape = tuple(out_shape)

    # Check min_samples
    if n < min_samples:
        nan_result = np.full(out_shape, np.nan)
        return nan_result, nan_result.copy()

    # Normal approximation to binomial
    z = abs(norm.ppf((1 - confidence) / 2))
    spread = z * np.sqrt(quantile * (1 - quantile) * n)
    idx_lower = n * quantile - spread
    idx_upper = n * quantile + spread

    # Sort along the specified axis
    sorted_data = np.sort(data, axis=axis)

    # Move operating axis to position 0 for vectorized interpolation
    sorted_data = np.moveaxis(sorted_data, axis, 0)

    # 1-based indices (matching statistical convention)
    indices = np.arange(1, n + 1)

    # Flatten non-operating dimensions for vectorized interp
    remaining_shape = sorted_data.shape[1:]
    sorted_flat = sorted_data.reshape(n, -1)

    # Interpolate each column
    ci_lower_flat = np.array(
        [np.interp(idx_lower, indices, col) for col in sorted_flat.T],
    )
    ci_upper_flat = np.array(
        [np.interp(idx_upper, indices, col) for col in sorted_flat.T],
    )

    # Reshape back to output shape
    ci_lower = ci_lower_flat.reshape(remaining_shape)
    ci_upper = ci_upper_flat.reshape(remaining_shape)

    return ci_lower, ci_upper
