import numpy
import numpy.ma as ma
import pandas
from numpy.lib.histograms import _hist_bin_selectors
from numpy.ma.core import nomask as _nomask

from .colors import hexcolor


def robust_min_max(
    s, threshold_for_dropping_zeros=0.35, cut_after_stdevs=5, percentiles_to_cut=1
):
    # Get s as an array (it should be already though)
    s = numpy.asarray(s)
    # mask out infs and nans
    sx = ma.masked_array(s, mask=~numpy.isfinite(s))
    n_values = sx.count()
    # Check if data is mostly zeros
    n_zeros = (s == 0).sum()
    if n_zeros / n_values > threshold_for_dropping_zeros:
        sx.mask |= sx == 0
        zeros_are_dropped = True
        n_values = sx.count()
    else:
        zeros_are_dropped = False
    if n_values == 0:
        # no data is left, default values
        bottom, top, range_is_truncated = 0, 1, 100
    else:
        mean = sx.mean()
        stdev = sx.std()
        bottom = sx.min()
        top = sx.max()
        range_is_truncated = 0
        cut_bottom = bottom < mean - cut_after_stdevs * stdev
        cut_top = top > mean + cut_after_stdevs * stdev
        if percentiles_to_cut > 0:
            if cut_bottom:
                bottom_ = numpy.percentile(sx, percentiles_to_cut)
                if not numpy.isnan(bottom_):
                    bottom = bottom_
                    range_is_truncated += percentiles_to_cut
            if cut_top:
                top_ = numpy.percentile(sx, 100 - percentiles_to_cut)
                if not numpy.isnan(top_):
                    top = top_
                    range_is_truncated += percentiles_to_cut
    return bottom, top, zeros_are_dropped, range_is_truncated, sx


def fixed_min_max(s, fixed_min, fixed_max, threshold_for_dropping_zeros=0.35):
    # Get s as an array (it should be already though)
    s = numpy.asarray(s)
    # mask out infs and nans
    sx = ma.masked_array(s, mask=numpy.isnan(s))
    n_values_finite = n_values = sx.count()

    n_zeros = (s == 0).sum()
    if n_zeros / n_values > threshold_for_dropping_zeros:
        sx.mask |= sx == 0
        zeros_are_dropped = True
        n_values = sx.count()
    else:
        zeros_are_dropped = False

    # Over max
    if fixed_max is not None:
        sx.mask |= sx >= fixed_max
    n_values_overmax = n_values - sx.count()
    n_values = sx.count()

    # Under Min
    if fixed_min is not None:
        sx.mask |= sx <= fixed_min
    n_values_undermin = n_values - sx.count()
    n_values = sx.count()

    pct_over = n_values_overmax / n_values_finite
    pct_under = n_values_undermin / n_values_finite

    return sx, pct_under, pct_over, zeros_are_dropped


def get_histogram_bins(
    data_to_bin,
    bins="sturges",
    percentiles_to_cut=1,
    maxrange=None,
    fixedrange=None,
    threshold_for_dropping_zeros=0.35,
):
    if fixedrange is None:
        bottom, top, zeros_are_dropped, range_is_truncated, sx = robust_min_max(
            data_to_bin, percentiles_to_cut=percentiles_to_cut
        )
        pct_under, pct_over = 0, 0
    else:
        bottom, top = fixedrange[0], fixedrange[1]
        sx, pct_under, pct_over, zeros_are_dropped = fixed_min_max(
            data_to_bin,
            bottom,
            top,
            threshold_for_dropping_zeros=threshold_for_dropping_zeros,
        )
        range_is_truncated = (pct_under + pct_over) * 100
    bottom = float(bottom) if bottom is not None else sx.min() - 1e-8
    top = float(top) if top is not None else sx.max() + 1e-8

    if numpy.isinf(bottom):
        bottom = sx.min()
    if numpy.isinf(top):
        top = sx.max()

    if maxrange is not None:
        if maxrange[0] > bottom:
            bottom = float(maxrange[0])
        if maxrange[1] < top:
            top = float(maxrange[1])

    if bottom > top:  # unclear how this would happen, but just in case
        bottom, top = top, bottom

    # generate number of bins based on named algorithm
    if isinstance(bins, str):
        try:
            width = _hist_bin_selectors[bins](sx, (bottom, top))
        except IndexError:
            width = 1
        if numpy.isnan(width):
            width = 1
        try:
            if width:
                bins = int(numpy.ceil((top - bottom) / width))
            else:
                bins = 1
        except OverflowError:
            bins = 10
        # The spark graphs get hard to read if the bin slices are too thin, so we will max out at 50 bins
        if bins > 50:
            bins = 50
        # The spark graphs look weird if the bin slices are too fat, so we will min at 10 bins
        if bins < 10:
            bins = 10

    # generate actual bins based on integer number of bins
    if isinstance(bins, int):
        try:
            bins = numpy.linspace(
                bottom, top, num=bins, endpoint=True, retstep=False, dtype=None
            )
        except UnboundLocalError:
            raise

    return (
        bins,
        bottom,
        top,
        zeros_are_dropped,
        range_is_truncated,
        sx,
        pct_under,
        pct_over,
    )


def _compress_by_mask(arr, mask=None):
    if arr is None:
        return None
    if isinstance(arr, numpy.ma.core.MaskedArray):
        if mask is None:
            mask = arr
        arr = arr._data
    if isinstance(arr, pandas.Series):
        arr = arr.data
    if isinstance(arr, memoryview):
        arr = numpy.asarray(arr)
    data = numpy.ndarray.ravel(arr)
    if mask._mask is not _nomask:
        data = data.compress(numpy.logical_not(numpy.ndarray.ravel(mask._mask)))
    return data


def histogram_data(
    data_to_bin,
    bins="sturges",
    weights=None,
    ch_weights=None,
    density=None,
    percentiles_to_cut=1,
    maxrange=None,
    fixedrange=None,
    threshold_for_dropping_zeros=0.35,
):
    (
        bins,
        bottom,
        top,
        zeros_are_dropped,
        range_is_truncated,
        sx,
        pct_under,
        pct_over,
    ) = get_histogram_bins(
        data_to_bin,
        bins=bins,
        percentiles_to_cut=percentiles_to_cut,
        maxrange=maxrange,
        fixedrange=fixedrange,
        threshold_for_dropping_zeros=threshold_for_dropping_zeros,
    )

    h, bin_edges = numpy.histogram(
        _compress_by_mask(sx),
        bins=bins,
        range=(bottom, top),
        weights=_compress_by_mask(weights, sx),
        density=False,
    )
    if ch_weights is not None:
        ch_h, _ = numpy.histogram(
            _compress_by_mask(sx),
            bins=bin_edges,
            weights=_compress_by_mask(ch_weights, sx),
            density=False,
            range=(bottom, top),
        )
    else:
        ch_h = None

    if density:
        if fixedrange is not None and ch_weights is not None:
            # Get adjusted density for choice heights
            db = numpy.array(numpy.diff(bin_edges), float)
            if weights is None:
                h = h / db / sx._data.size
            else:
                h = h / db / weights.sum()
            ch_h = ch_h / db / ch_weights.sum()
        else:
            db = numpy.array(numpy.diff(bin_edges), float)
            h = h / db / h.sum()
            ch_h = ch_h / db / ch_h.sum()

    return (
        h,
        bin_edges,
        zeros_are_dropped,
        range_is_truncated,
        ch_h,
        pct_under,
        pct_over,
    )


_histogram_notes = {
    hexcolor(
        "orange"
    ): "Histograms are orange if the zeros are numerous and have been excluded.",
    hexcolor(
        "red"
    ): "Histograms are red if the zeros are numerous and have been excluded, "
    "and the displayed range truncates some extreme outliers.",
    hexcolor(
        "forest"
    ): "Histograms are green if the displayed range truncates some extreme outliers.",
    hexcolor("ocean"): None,
    hexcolor(
        "night"
    ): "Histograms are purple if the data is represented as discrete values.",
}


def draw_histogram_figure(
    bin_heights,
    bin_edges,
    zeros_are_dropped,
    range_is_truncated,
    title=None,
    xlabel=None,
    ylabel=None,
    xticks=False,
    yticks=False,
    xticklabels=None,
    yticklabels=False,
    frame_off="all",
    notetaker=None,
    bgcolor=None,
    figwidth=0.75,
    figheight=0.2,
    return_format="svg",
    attach_metadata=True,
    ch_heights=None,
    discrete=False,
    left_thermo=None,
    right_thermo=None,
):
    from .._optional import pyplot as plt

    fig, ax = plt.subplots(nrows=1, ncols=1, sharex=False, sharey=False)
    widths = bin_edges[1:] - bin_edges[:-1]
    if discrete:
        widths = widths * 0.8
        color = hexcolor("night")
    elif zeros_are_dropped:
        if range_is_truncated and not left_thermo and not right_thermo:
            color = hexcolor("red")
        else:
            color = hexcolor("orange")
    elif range_is_truncated and not left_thermo and not right_thermo:
        color = hexcolor("forest")
    else:
        color = hexcolor("ocean")
    # rects1 = ax.bar(
    #     bin_edges[:-1],
    #     bin_heights,
    #     widths,
    #     color=color,
    #     align="center" if discrete else "edge",
    # )

    chart_top = bin_heights.max()

    if ch_heights is not None:
        # rects2 = ax.bar(
        #     bin_edges[:-1],
        #     height=numpy.zeros_like(ch_heights),
        #     width=widths,
        #     bottom=ch_heights,
        #     edgecolor="black",
        #     linewidth=1.0,
        #     align="center" if discrete else "edge",
        #     clip_on=False,
        # )
        chart_top = max(chart_top, ch_heights.max())

    if left_thermo is not None:
        leftmost, rightmost = bin_edges[0], bin_edges[-1]
        full_width = rightmost - leftmost
        ax.bar(
            numpy.asarray(
                [
                    leftmost - full_width * 0.05,
                ]
            ),
            height=chart_top * left_thermo,
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            color="red",
            align="edge",
        )
        ax.bar(
            numpy.asarray(
                [
                    leftmost - full_width * 0.05,
                ]
            ),
            height=chart_top * (1 - left_thermo),
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            bottom=chart_top * left_thermo,
            color="#e9edd3",
            align="edge",
        )
        # Thermo-border
        ax.bar(
            numpy.asarray(
                [
                    leftmost - full_width * 0.05,
                ]
            ),
            height=chart_top,
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            color="#FFFFFF00",
            align="edge",
            edgecolor="black",
            linewidth=0.5,
            clip_on=False,
        )

    if right_thermo is not None:
        leftmost, rightmost = bin_edges[0], bin_edges[-1]
        full_width = rightmost - leftmost
        ax.bar(
            numpy.asarray(
                [
                    rightmost + full_width * 0.015,
                ]
            ),
            height=chart_top * right_thermo,
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            color="red",
            align="edge",
        )
        ax.bar(
            numpy.asarray(
                [
                    rightmost + full_width * 0.015,
                ]
            ),
            height=chart_top * (1 - right_thermo),
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            bottom=chart_top * right_thermo,
            color="#e9edd3",
            align="edge",
        )
        # Thermo-border
        ax.bar(
            numpy.asarray(
                [
                    rightmost + full_width * 0.015,
                ]
            ),
            height=chart_top,
            width=numpy.asarray(
                [
                    full_width * 0.035,
                ]
            ),
            color="#FFFFFF00",
            align="edge",
            edgecolor="black",
            linewidth=0.5,
            clip_on=False,
        )

    if figheight:
        fig.set_figheight(figheight)
    if figwidth:
        fig.set_figwidth(figwidth)
    fig.set_dpi(300)

    if frame_off == "all" or frame_off is True or frame_off == 1:
        frame_off = {"left", "right", "top", "bottom"}
    elif frame_off is None:
        frame_off = set()
    elif not isinstance(frame_off, set):
        print(type(frame_off))
        print(frame_off)
        frame_off = set(frame_off)

    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    if not isinstance(xticks, (list, tuple, numpy.ndarray)) and not xticks:
        ax.get_xaxis().set_ticks([])
    else:
        frame_off.discard("bottom")
        if xticklabels is not None:
            ax.get_xaxis().set_ticklabels(xticklabels)
            ax.get_xaxis().set_ticks(xticks)
    if not isinstance(yticks, (list, tuple, numpy.ndarray)) and not yticks:
        ax.get_yaxis().set_ticks([])
    else:
        frame_off.discard("left")
        if yticklabels is not None:
            ax.get_yaxis().set_ticklabels(yticklabels)
            ax.get_yaxis().set_ticks(yticks)

    if frame_off is not None:
        if isinstance(frame_off, str):
            frame_off = [frame_off]
        from .plotting import adjust_spines

        sides = {"left", "right", "top", "bottom"}
        for side in frame_off:
            sides.remove(side)
        adjust_spines(ax, sides)

    if bgcolor is not None:
        ax.set_axis_bgcolor(bgcolor)

    if notetaker is not None and color != hexcolor("ocean"):
        notetaker.add(_histogram_notes[color])

    if return_format.lower() == "svg":
        from .plotting import plot_as_svg_xhtml

        pad_inches = 0.1
        try:
            if "left" in frame_off and "bottom" in frame_off:
                ax.set_position([0, 0, 1, 1], which="both")
                ax.set_frame_on(False)
                pad_inches = 0
            elif figheight <= 1.0:
                plt.tight_layout(pad=figheight / 2)
                pad_inches = 0
            else:
                plt.tight_layout(pad=0.5)
        except ValueError:
            pass  # tight layout sometimes fails when the figsize is too small
        tooltip = _histogram_notes[color] or ""
        if ch_heights is not None:
            if len(tooltip):
                tooltip += "\n"
            tooltip += (
                "Black marks indicate frequencies considering only chosen alternatives."
            )
        if left_thermo or right_thermo:
            if len(tooltip):
                tooltip += "\n"
            tooltip += (
                "Side thermometers indicate the fraction of all observations "
                "at or beyond the limits, and are scaled independently."
            )

        ret = plot_as_svg_xhtml(plt, tooltip=tooltip, pad_inches=pad_inches)
        plt.clf()
        plt.close()  # do not showing empty windows?
        if attach_metadata:
            ret.metadata = dict(
                bin_heights=bin_heights,
                bin_edges=bin_edges,
                zeros_are_dropped=zeros_are_dropped,
                range_is_truncated=range_is_truncated,
            )
            ret.attrib["zeros_are_dropped"] = str(zeros_are_dropped)
            ret.attrib["range_is_truncated"] = str(range_is_truncated)
        ret.attrib["style"] = f"min-width:{figwidth}in"
        return ret


def histogram_figure(
    data_to_bin,
    bins="sturges",
    weights=None,
    ch_weights=None,
    return_all=False,
    percentiles_to_cut=1,
    maxrange=None,
    fixedrange=None,
    piecerange=None,
    mask=None,
    **kwargs,
):
    density = None if ch_weights is None else True
    discrete = bins == "discrete"

    if mask is not None:
        data_to_bin = data_to_bin[mask]
        if weights is not None:
            weights = weights[mask]
        if ch_weights is not None:
            ch_weights = ch_weights[mask]

    threshold_for_dropping_zeros = 0.35

    if piecerange is not None and piecerange != (None, None):
        if piecerange[0] is None:
            fixedrange = (None, piecerange[1])
        elif piecerange[1] is None:
            fixedrange = (0, data_to_bin.max() + 1e-8)
            threshold_for_dropping_zeros = 100
        else:
            fixedrange = (0, piecerange[1] - piecerange[0])
            threshold_for_dropping_zeros = 100
        percentiles_to_cut = 0
        discrete = False

    if discrete:
        bin_labels, bin_inv = numpy.unique(data_to_bin, return_inverse=True)
        bin_heights = numpy.bincount(bin_inv, weights)
        bin_edges = numpy.arange(len(bin_heights) + 1)
        zeros_are_dropped, range_is_truncated = False, False
        if ch_weights is not None:
            ch_heights = numpy.bincount(bin_inv, ch_weights)
        else:
            ch_heights = None
        if density:
            bin_heights = bin_heights / (bin_heights.sum())
            if ch_heights is not None:
                ch_heights = ch_heights / (ch_heights.sum())
        if (("xticks" not in kwargs) and (kwargs.get("figwidth", 1) >= 2)) or (
            kwargs.get("xticks") is True
        ):
            kwargs["xticklabels"] = bin_labels
            kwargs["xticks"] = bin_edges[:-1]
        pct_under, pct_over = 0, 0
    else:
        (
            bin_heights,
            bin_edges,
            zeros_are_dropped,
            range_is_truncated,
            ch_heights,
            pct_under,
            pct_over,
        ) = histogram_data(
            data_to_bin,
            bins,
            weights,
            ch_weights,
            density=density,
            percentiles_to_cut=percentiles_to_cut,
            maxrange=maxrange,
            fixedrange=fixedrange,
            threshold_for_dropping_zeros=threshold_for_dropping_zeros,
        )
        bin_labels = None

    if piecerange is not None and piecerange[0] is not None:
        bin_edges += piecerange[0]

    try:
        f = draw_histogram_figure(
            bin_heights,
            bin_edges,
            zeros_are_dropped,
            range_is_truncated,
            ch_heights=ch_heights,
            discrete=discrete,
            left_thermo=pct_under or None,
            right_thermo=pct_over or None,
            **kwargs,
        )
    except Exception:
        print("bin_heights\n", bin_heights)
        print("bin_edges\n", bin_edges)
        print("zeros_are_dropped\n", zeros_are_dropped)
        print("range_is_truncated\n", range_is_truncated)
        raise
    if return_all:
        return f, bin_heights, bin_edges, zeros_are_dropped, range_is_truncated
    else:
        return f


def sizable_histogram_figure(
    *args, sizer=1, discrete=None, on_error="ignore", **kwargs
):
    kwargs["figwidth"] = sizer
    kwargs["figheight"] = sizer / 3.5
    if sizer >= 2:
        kwargs["frame_off"] = {"top", "right"}
        kwargs["xticks"] = True
        # kwargs['yticks'] = True
    if sizer < 4:
        kwargs.pop("title", None)
    if sizer < 3:
        kwargs.pop("xlabel", None)
        kwargs.pop("ylabel", None)
    if discrete is not None:
        if discrete:
            kwargs["bins"] = "discrete"
    try:
        return histogram_figure(*args, **kwargs)
    except Exception:
        if on_error == "ignore":
            return None
        else:
            raise


def seems_like_discrete_data(arr, dictionary=None, return_uniques=False):
    if isinstance(arr.dtype, pandas.CategoricalDtype):
        # print('seems_like_discrete_data? YES categorical')
        if return_uniques:
            return (True, list(arr.cat.categories))
        return True
    elif numpy.issubdtype(arr.dtype, numpy.bool_):
        # print('seems_like_discrete_data? YES bool')
        if return_uniques:
            return (True, [False, True])
        return True
    else:
        pass
        # print('seems_like_discrete_data? not bool but',arr.dtype)
    if dictionary is None:
        if len(numpy.unique(arr[:100])) < 6:
            if len(numpy.unique(arr[:1000])) < 6:
                uniques = numpy.unique(arr)
                if len(uniques) < 6:
                    # print('seems_like_discrete_data? YES uniques < 6')
                    if return_uniques:
                        return (True, uniques)
                    return True
        # print('seems_like_discrete_data? too many and no dictionary')
    else:
        uniq = numpy.unique(arr)
        not_in_dict = 0
        for i in uniq:
            if i not in dictionary:
                not_in_dict += 1
        if not_in_dict > 2:
            # print(f'seems_like_discrete_data? dictionary but {not_in_dict} missing keys')
            if return_uniques:
                return False, None
            return False
        else:
            # print(f'seems_like_discrete_data? dictionary with {not_in_dict} missing keys')
            if return_uniques:
                return True, uniq
            return True
    if return_uniques:
        return False, None
    return False
