from __future__ import annotations

from os import PathLike

import fitsio  # type: ignore [import-not-found]
from numpy.typing import NDArray

TYPE_CHECKING = False
if TYPE_CHECKING:
    from typing import TypeAlias, Any

    _DictKey: TypeAlias = str | int | tuple["_DictKey", ...]


def _key_from_string(s: str) -> tuple[str, str, int, int] | None:
    """return key from string representation"""
    tuple_key_dict = None
    if "POSPOS" in s:
        name_split = s.split("_")
        tuple_key_dict = ("POS", "POS", int(name_split[-2]), int(name_split[-1]))
    elif "POSSHEAR" in s:
        name_split = s.split("_")
        tuple_key_dict = ("POS", "SHE", int(name_split[-2]), int(name_split[-1]))
    elif "SHEAR" in s:
        name_split = s.split("_")
        tuple_key_dict = ("SHE", "SHE", int(name_split[-2]), int(name_split[-1]))
    return tuple_key_dict


def correlation_functions(path: str | PathLike[str]) -> dict[_DictKey, NDArray[Any]]:
    """
    Reads 2D correlation functions from a Euclid data product.

    Parameters
    ----------
    path : str | PathLike[str]
        The path to the FITS file containing the Euclid data product.

    Returns
    -------
    dict[_DictKey, NDArray[Any]]
        A dictionary where each key is a tuple representing the HDU name (generated by
        `_key_from_string`), and each value is a NumPy array containing the data of
        the corresponding HDU.

    Notes
    -----
    - Only HDUs whose extension name contains '2D' are considered.
    - The `fitsio.FITS` context manager is used to read the FITS file efficiently.
    - The function depends on `_key_from_string()` to convert the HDU names into keys
      for the output dictionary.
    """

    xi: dict[_DictKey, NDArray[Any]] = {}
    with fitsio.FITS(path) as fits:
        for hdu in fits:
            if "2D" not in hdu.get_extname():
                continue
            key = _key_from_string(hdu.get_extname())
            if key is None:
                continue
            xi[key] = hdu.read()
    return xi


def bandpowers(path: str | PathLike[str]) -> dict[_DictKey, NDArray[Any]]:
    """
    Reads 2D bandpowers from a Euclid data product.

    Parameters
    ----------
    path : str | PathLike[str]
        The path to the FITS file containing the Euclid bandpower data product.

    Returns
    -------
    dict[_DictKey, NDArray[Any]]
        A dictionary where each key is a tuple representing the HDU name (generated by
        `_key_from_string`), and each value is a NumPy array containing the data of
        the corresponding HDU.

    Notes
    -----
    - Only HDUs whose extension name contains '2D' are considered.
    - The `fitsio.FITS` context manager is used to read the FITS file efficiently.
    - The function depends on `_key_from_string()` to convert the HDU names into keys
      for the output dictionary.
    """
    bandp: dict[_DictKey, NDArray[Any]] = {}
    with fitsio.FITS(path) as fits:
        for hdu in fits:
            if "2D" not in hdu.get_extname():
                continue
            key = _key_from_string(hdu.get_extname())
            if key is None:
                continue
            data = hdu.read()
            if key[:2] == ("POS", "POS"):
                data.dtype.names = ["L", "CL", "LMIN", "LMAX"]
            elif key[:2] == ("POS", "SHE"):
                data.dtype.names = ["L", "CL_E", "CL_B", "LMIN", "LMAX"]
            elif key[:2] == ("SHE", "SHE"):
                data.dtype.names = ["L", "CL_E", "CL_B", "LMIN", "LMAX"]
            bandp[key] = data
    return bandp


def cosebis(path: str | PathLike[str]) -> dict[_DictKey, NDArray[Any]]:
    """
    Reads 2D bandpowers from a Euclid data product.

    Parameters
    ----------
    path : str | PathLike[str]
        The path to the FITS file containing the Euclid bandpower data product.

    Returns
    -------
    dict[_DictKey, NDArray[Any]]
        A dictionary where each key is a tuple representing the HDU name (generated by
        `_key_from_string`), and each value is a NumPy array containing the data of
        the corresponding HDU.

    Notes
    -----
    - Only HDUs whose extension name contains '2D' are considered.
    - The `fitsio.FITS` context manager is used to read the FITS file efficiently.
    - The function depends on `_key_from_string()` to convert the HDU names into keys
      for the output dictionary.
    """

    cb: dict[_DictKey, NDArray[Any]] = {}
    with fitsio.FITS(path) as fits:
        for hdu in fits:
            if "2D" not in hdu.get_extname():
                continue
            key = _key_from_string(hdu.get_extname())
            if key is None:
                continue
            cb[key] = hdu.read()
    return cb
