from typing import (
    Optional,
    TYPE_CHECKING,
    Iterable,
    cast,
    Any,
    Sequence,
    Tuple,
    Union,
    List,
)

import numpy as np
import enum
from numbers import Real

from uhi.typing.plottable import PlottableHistogram, PlottableAxis
from uhi.numpy_plottable import ensure_plottable_histogram

if TYPE_CHECKING:
    from numpy.typing import ArrayLike
else:
    ArrayLike = Any


def get_plottable_protocol_bins(axis: PlottableAxis) -> np.ndarray:
    out = np.empty(len(axis) + 1)
    assert isinstance(
        axis[0], tuple
    ), f"Currently only support non-discrete axes {axis}"
    # TODO: Support discrete axes
    out[0] = axis[0][0]
    out[1:] = [axis[i][1] for i in range(len(axis))]  # type: ignore
    return out


def hist_object_handler(
    hist: Union[
        ArrayLike,
        PlottableHistogram,
        Tuple[Union[ArrayLike, None], ...],
        List[ArrayLike],
    ],
    *bins: Sequence[Optional[float]],
) -> PlottableHistogram:

    if not bins or all(b is None for b in bins):
        if isinstance(hist, list):
            if not bins and len(hist) > 0 and not isinstance(hist[0], (list, Real)):
                hist = tuple(hist)
            else:
                hist = (np.asarray(hist), None)
        hist_obj = ensure_plottable_histogram(hist)
    elif isinstance(hist, PlottableHistogram):
        raise TypeError("Cannot give bins with existing histogram")
    else:
        hist_obj = ensure_plottable_histogram((hist, *bins))

    if len(hist_obj.axes) not in {1, 2}:
        raise ValueError("Must have only 1 or 2 axes")

    return hist_obj


def process_histogram_parts(
    H: Union[
        ArrayLike, PlottableHistogram, Iterable[ArrayLike], List[PlottableHistogram]
    ],
    *bins: Sequence[Optional[float]],
):
    """
    Parameters
    ----------

        H : object
            Histogram object with containing values and optionally bins. Can be:

            - `np.histogram` tuple
            - PlottableHistogram object
            - `boost_histogram` classic (<0.13) histogram object
            - raw histogram values, provided `bins` is specified.

        A list of any of the above.

        *bins : Sequence[float], optional
            Histogram bins, if not part of ``h``. One iterable per histogram dimension.

    Returns
    -------
        values, bins: Iterator[Tuple[np.ndarray, np.ndarray]]
    """

    # Try to understand input
    if isinstance(H, list) and not isinstance(H[0], (Real)):
        return _process_histogram_parts_iter(H, *bins)
    else:
        return _process_histogram_parts_iter((H,), *bins)  # type: ignore


def _process_histogram_parts_iter(
    hists: Union[Iterable[ArrayLike], Iterable[PlottableHistogram]],
    *bins: Sequence[Optional[float]],
) -> Iterable[PlottableHistogram]:
    original_bins: Tuple[Sequence[float], ...] = bins  # type: ignore

    for hist in hists:
        h = hist_object_handler(hist, *bins)
        current_bins: Tuple[Sequence[float], ...] = tuple(get_plottable_protocol_bins(a) for a in h.axes)  # type: ignore
        if any(b is None for b in original_bins):
            original_bins = current_bins

        if len(current_bins) != len(original_bins):
            msg = "Plotting multiple histograms must have the same dimensionality"
            raise ValueError(msg)
        for i in range(len(current_bins)):
            diff_lengths = len(original_bins[i]) != len(current_bins[i])
            if diff_lengths or not np.allclose(original_bins[i], current_bins[i]):
                msg = "Plotting multiple histograms with different binning is not supported"
                raise ValueError(msg)

        yield h


def get_histogram_axes_title(axis: Any) -> str:

    if hasattr(axis, "label"):
        return axis.label
    # Classic support for older hist, deprecated
    elif hasattr(axis, "title"):
        return axis.title
    elif hasattr(axis, "name"):
        return axis.name

    # No axis title found
    return ""
