"""Simple conversions of 1d, 2d, or 3d covariance between 1 sigma and 95%."""

from __future__ import annotations

from typing import TYPE_CHECKING, overload

import numpy as np
from scipy.stats import chi2

from gri_utils import constants

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray

PPF_95_VALS = {
    1: constants.SIG_TO_95_1D,
    2: constants.SIG_TO_95_2D,
    3: constants.SIG_TO_95_3D,
}


def _ppf(conf: float, dof: int) -> float:
    """Return the ppf / chi2inv function value."""
    # Keep some values so we don't have to calculate them
    if abs(conf - 0.95) < 1e-8 and dof in PPF_95_VALS:  # noqa: PLR2004
        return PPF_95_VALS[dof]
    return float(chi2.ppf(conf, dof))


@overload
def scale_chi2(
    n: float,
    *,
    conf: float = 0.95,
    dof: int | None = None,
    power: float | None = None,
) -> float: ...
@overload
def scale_chi2(
    n: Sequence | NDArray[np.floating],
    *,
    conf: float = 0.95,
    dof: int | None = None,
    power: float | None = None,
) -> NDArray[np.floating]: ...
def scale_chi2(
    n: float | Sequence | NDArray[np.floating],
    *,
    conf: float = 0.95,
    dof: int | None = None,
    power: float | None = None,
) -> float | NDArray[np.floating]:
    """Convert a 1 sigma scalar, vector, or matrix to a confidence value.

    Helper function for our most common conversions.

    By default returns 95% confidence. For the different inputs, assumes that a float is
    1 degree of freedom and scales by (SIG_TO_95_1D)**0.5, a vector scales by
    (SIG_TO_95_2D)**0.5, a 2D matrix scales by SIG_TO_95_2D, and a 3D matrix scales by
    SIG_TO_95_3D. These can be overriden with the dimensionality (1 to 3) for dof and
    conf and power.

    conf is confidence, usually between 0 and 1

    dof is degrees of freedom

    power is power, where usually it is 0.5 or 1
    """
    n = np.asarray(n)
    if dof is None:
        if n.ndim == 0:
            dof = 1
        elif n.ndim == 1 or n.shape == (2, 2):
            dof = 2
        elif n.shape == (3, 3):
            dof = 3
        else:
            raise ValueError("dof must be defined")
    if power is None:
        if n.ndim in {0, 1}:
            power = 0.5
        elif n.shape in ((2, 2), (3, 3)):
            power = 1
        else:
            raise ValueError("power must be defined")
    ppf = _ppf(conf, dof)
    return n * ppf**power
