"""Confidence interval for Poisson-distributed count data using the Garwood method.

The Garwood method (1936) provides exact confidence intervals for Poisson rates
using the chi-squared distribution. Given an observed count k from a Poisson
process, the confidence interval bounds the true rate parameter lambda.

Reference:
    Garwood, F. (1936). Fiducial Limits for the Poisson Distribution.
    Biometrika, 28(3/4), 437-442.
    https://www.ine.pt/revstat/pdf/rs120203.pdf
"""

from __future__ import annotations

from typing import TYPE_CHECKING, overload

import numpy as np
from scipy.stats import chi2

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


@overload
def poisson_ci(
    counts: float,
    confidence: float = 0.95,
) -> tuple[float, float]: ...


@overload
def poisson_ci(
    counts: NDArray[np.floating] | Sequence,
    confidence: float = 0.95,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]: ...


def poisson_ci(
    counts: float | NDArray[np.floating] | Sequence,
    confidence: float = 0.95,
) -> tuple[float, float] | tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Calculate confidence interval for Poisson count data using the Garwood method.

    Given observed counts from a Poisson distribution, computes exact confidence
    interval bounds for the true Poisson rate parameter. The method uses the
    relationship between Poisson and chi-squared distributions.

    Note:
        This is NOT the same as scipy.stats.poisson.interval(), which returns
        probability intervals (quantiles), not confidence intervals for the rate.

    Args:
        counts: Observed count(s). Can be a scalar, array, or sequence.
            All values should be non-negative integers (though floats are accepted).
        confidence: Confidence level in (0, 1). Default is 0.95 for 95% CI.

    Returns:
        Tuple of (lower_bound, upper_bound) arrays. Each has the same shape as
        the input counts. For count=0, lower bound is 0.

    Example:
        >>> lower, upper = poisson_ci(10)  # Single count
        >>> lower, upper = poisson_ci([5, 10, 15])  # Multiple counts
        >>> lower, upper = poisson_ci(10, confidence=0.99)  # 99% CI
    """
    counts = np.asarray(counts, dtype=np.float64)
    scalar_input = counts.ndim == 0
    counts = np.atleast_1d(counts)

    alpha = 1 - confidence

    # Garwood method using chi-squared distribution
    # Lower bound: chi2(alpha/2, 2k) / 2
    # Upper bound: chi2(1-alpha/2, 2(k+1)) / 2
    lower: NDArray[np.floating] = np.asarray(chi2.ppf(alpha / 2, 2 * counts) / 2)
    upper: NDArray[np.floating] = np.asarray(
        chi2.ppf(1 - alpha / 2, 2 * (counts + 1)) / 2,
    )

    # Handle count=0 case: lower bound should be 0
    lower = np.where(counts == 0, 0.0, lower)

    if scalar_input:
        return float(lower.item(0)), float(upper.item(0))
    return lower, upper
