from __future__ import annotations

import enum

import numpy
import numpy.ma as ma
import pandas

from .. import warning
from . import dictx
from .arraytools import scalarize
from .data_expansion import parse_piece
from .histograms import seems_like_discrete_data, sizable_histogram_figure


def statistics_for_array(
    arr,
    histogram=True,
    varname=None,
    ch_weights=None,
    avail=None,
    dictionary=None,
    flatten=False,
    on_error="ignore",
    **kwargs,
):
    """
    Generate a statistical analysis of an array.

    Parameters
    ----------
    arr : ndarray
            The array to analyze.
    histogram : bool or numeric
            If this evaluates as True, a histogram for the data array
            will be generated. If given as a number, this indicates the
            size and level of detail in the histogram.
    varname : str, optional
            A label for the x-axis of the histogram.  Only appears if
            the size is 3 or greater.  The varname is also used to
            detect piecewise linear segments, when given as
            "piece(label, lowerbound, upperbound)". Note this function
            does not *process* the array to extract a linear segment,
            this just flags the histogram generator to know that this
            is a linear segment.
    ch_weights : ndarray, optional
            If given, this indicates the choice-weights of the observations.
            Must be the same shape as `arr`. Useful to analyze whether the
            distribution of the data appears different for chosen
            alternatives than for all alternatives.
    avail : ndarray, optional
            If given, this indicates the availability of the observations.
            Must be the same shape as `arr`. If missing, all finite values
            in `arr` are considered available.
    dictionary : dict, optional
            A dictionary of {value: str} pairs, which give labels to use on
            the histogram instead of the code values stored in `arr`.
    flatten : bool, default False
            Indicates if `arr` (and `ch_weights` if given) should be
            flattened before processing.
    on_error : {'ignore', 'raise'}
            How to handle exceptions encountered while generating statistics.

    Other keyword parameters are passed through to
    :meth:`sizable_histogram_figure`.

    Returns
    -------
    dictx

    """
    a = arr

    if flatten:
        a = a.reshape(-1)
        if ch_weights is not None:
            ch_weights = ch_weights.reshape(-1)

    stats = dictx()

    if varname is not None:
        stats.description = varname

    base_def, lowbound, highbound, piece_exactly = parse_piece(varname)

    try:
        _nans = numpy.isnan(a)
    except TypeError:
        can_nan = False
    else:
        can_nan = True

    a_masked_is_category = None

    if avail is None and can_nan:
        a_masked = ma.masked_array(a, mask=~numpy.isfinite(a))
        stats.n = scalarize(numpy.isfinite(a).sum(axis=0))
    elif avail is not None:
        a_masked = ma.masked_array(a, mask=~avail)
        stats.n = scalarize((avail != 0).sum(axis=0))
    elif isinstance(a, pandas.Series) and isinstance(a.dtype, pandas.CategoricalDtype):
        a_masked = ma.masked_array(a, mask=pandas.isnull(a))
        a_masked_is_category = True
        stats.n = scalarize((~pandas.isnull(a)).sum(axis=0))
    else:
        a_masked = a
        stats.n = scalarize(a.shape[0])

    if can_nan:
        stats.minimum = scalarize(numpy.nanmin(a_masked, axis=0))
        stats.maximum = scalarize(numpy.nanmax(a_masked, axis=0))
        with warning.ignore_warnings():
            if ma.is_masked:
                stats.median = scalarize(ma.median(a_masked, axis=0))
            else:
                stats.median = scalarize(numpy.nanmedian(a_masked, axis=0))

    if ch_weights is not None:
        a_ch_weighted = ma.masked_array(a, mask=(ch_weights <= 0))
        stats.n_chosen = scalarize(a_ch_weighted.count(axis=0))
        if can_nan:
            stats.minimum_chosen = scalarize(numpy.nanmin(a_ch_weighted, axis=0))
            stats.maximum_chosen = scalarize(numpy.nanmax(a_ch_weighted, axis=0))
            with warning.ignore_warnings():
                stats.median_chosen = scalarize(numpy.nanmedian(a_ch_weighted, axis=0))

    if ch_weights is not None:
        if avail is None and can_nan:
            ch_weightsx = ma.masked_array(ch_weights, mask=~numpy.isfinite(a))
        elif avail is not None:
            ch_weightsx = ma.masked_array(ch_weights, mask=~avail)
        elif isinstance(a, pandas.Series) and isinstance(
            a.dtype, pandas.CategoricalDtype
        ):
            ch_weightsx = ma.masked_array(ch_weights, mask=pandas.isnull(a))
        else:
            ch_weightsx = ch_weights
    else:
        ch_weightsx = None

    if a_masked_is_category is None:
        try:
            a_masked_is_category = isinstance(a_masked.dtype, pandas.CategoricalDtype)
        except Exception:
            a_masked_is_category = False

    if (can_nan and histogram) or a_masked_is_category:
        with warning.ignore_warnings():
            discrete_ = kwargs.pop("discrete", None)
            try:
                if a_masked_is_category:
                    discrete_ = True
            except TypeError:
                pass
            if discrete_ is None:
                discrete_ = seems_like_discrete_data(a_masked, dictionary)

            if lowbound is not None or highbound is not None:
                stats.histogram = sizable_histogram_figure(
                    a_masked,
                    sizer=histogram,
                    title=None,
                    xlabel=base_def,
                    ylabel="Frequency",
                    ch_weights=ch_weightsx,
                    piecerange=(lowbound, highbound),
                    **kwargs,
                )
            else:
                stats.histogram = sizable_histogram_figure(
                    a_masked,
                    sizer=histogram,
                    title=None,
                    xlabel=varname,
                    ylabel="Frequency",
                    ch_weights=ch_weightsx,
                    discrete=discrete_,
                    on_error=on_error,
                    **kwargs,
                )
    if can_nan:
        stats.mean = scalarize(numpy.mean(a_masked, axis=0))
        stats.stdev = scalarize(numpy.std(a_masked, axis=0))
        stats.zeros = scalarize(numpy.sum(numpy.logical_not(a_masked), axis=0))
        with warning.ignore_warnings():
            stats.positives = scalarize(numpy.sum(a_masked > 0, axis=0))
            stats.negatives = scalarize(numpy.sum(a_masked < 0, axis=0))
        n_nans = scalarize(numpy.sum(numpy.isnan(a), axis=0))
        if numpy.any(n_nans):
            stats.nans = n_nans
        n_infs = scalarize(numpy.sum(numpy.isinf(a), axis=0))
        if numpy.any(n_infs):
            stats.infs = n_infs
        a_masked.mask |= a_masked == 0
        stats.nonzero_minimum = scalarize(numpy.min(a_masked, axis=0))
        stats.nonzero_maximum = scalarize(numpy.max(a_masked, axis=0))
        stats.nonzero_mean = scalarize(numpy.mean(a_masked, axis=0))
        stats.nonzero_stdev = scalarize(numpy.std(a_masked, axis=0))
    elif a_masked_is_category:
        stats.nulls = pandas.isnull(a).sum()
        the_mode = tuple(a.mode().values)
        if len(the_mode) == 1:
            the_mode = the_mode[0]
        stats.mode = the_mode

    if dictionary is not None:
        stats.dictionary = dictionary

    return stats


def statistics_for_array5(arr, histogram=5, *args, **kwargs):
    """
    Generate a statistical analysis of an array.

    Parameters
    ----------
    arr : ndarray
            The array to analyze.
    histogram : bool or numeric, default 5
            If this evaluates as True, a histogram for the data array
            will be generated. If given as a number, this indicates
            the size and level of detail in the histogram.
    varname : str, optional
            A label for the x-axis of the histogram.  Only appears if
            the size is 3 or greater.  The varname is also used to
            detect piecewise linear segments, when given as
            "piece(label, lowerbound, upperbound)". Note this function
            does not *process* the array to extract a linear segment,
            this just flags the histogram generator to know that this
            is a linear segment.
    ch_weights : ndarray, optional
            If given, this indicates the choice-weights of the
            observations.  Must be the same shape as `arr`. Useful to
            analyze whether the distribution of the data appears
            different for chosen alternatives than for all alternatives.
    avail : ndarray, optional
            If given, this indicates the availability of the observations.
            Must be the same shape as `arr`. If missing, all finite values
            in `arr` are considered available.
    dictionary : dict, optional
            A dictionary of {value: str} pairs, which give labels to use
            on the histogram instead of the code values stored in `arr`.
    flatten : bool, default False
            Indicates if `arr` (and `ch_weights` if given) should be
            flattened before processing.

    Other keyword parameters are passed through to
    :meth:`sizable_histogram_figure`.

    Returns
    -------
    dictx

    """
    return statistics_for_array(arr, *args, histogram=histogram, **kwargs)


def statistics_for_dataframe(df, histogram=True, ch_weights=None, avail=None, **kwargs):
    s = {}
    for col in df.columns:
        if isinstance(df.dtypes[col], pandas.CategoricalDtype):
            s[col] = statistics_for_array(
                df[col],
                varname=col,
                histogram=histogram,
                ch_weights=ch_weights,
                avail=avail,
                **kwargs,
            )
        elif numpy.issubdtype(df.dtypes[col], numpy.floating):
            s[col] = statistics_for_array(
                df[col],
                varname=col,
                histogram=histogram,
                ch_weights=ch_weights,
                avail=avail,
                **kwargs,
            )
        elif numpy.issubdtype(df.dtypes[col], numpy.integer):
            s[col] = statistics_for_array(
                df[col],
                varname=col,
                histogram=histogram,
                ch_weights=ch_weights,
                avail=avail,
                **kwargs,
            )
        else:
            s[col] = statistics_for_array(
                df[col],
                varname=col,
                histogram=histogram,
                ch_weights=ch_weights,
                avail=avail,
                **kwargs,
            )

    from .dataframe import DataFrameViewer

    result = DataFrameViewer.from_dict(s, orient="index")
    if "nans" in result.columns:
        result["nans"].fillna(0, inplace=True)
    try:
        result = result.drop("description", axis=1)
    except Exception:
        pass
    return result


def uniques(s, dictionary=None):
    """
    Count the frequency for each unique value in a pandas.Series.

    Parameters
    ----------
    s : pandas.Series
            The values to count.
    dictionary : Mapping or Enum, optional
            A mapping converting the actual values in the series
            to a more meaningful value (e.g. changes code numbers
            to names).  If an Enum class is passed, the values
            are treated as keys, and the names as values.

    Returns
    -------
    pandas.Series
    """
    action = s
    len_action = len(action)
    try:
        action = action[~numpy.isnan(action)]
    except TypeError:
        num_nan = 0
    else:
        num_nan = len_action - len(action)
    x = numpy.unique(action, return_counts=True)
    if dictionary is None:
        y = pandas.Series(x[1], x[0])
    else:
        if isinstance(dictionary, enum.EnumMeta):
            dictionary = {v: k for k, v in dictionary.__members__.items()}
        y = pandas.Series(x[1], [dictionary.get(j, j) for j in x[0]])
    if num_nan:
        y[numpy.nan] = num_nan
    return y


def invmap(s, mapping):
    """
    Apply an inverse map.

    When the values in `s` are the values of a Mapping,
    but you want them to be the keys instead.  Note
    that values need not be unique, so an arbitrary
    key with the correct value will result for non-unique
    inverse mappings.

    Parameters
    ----------
    s : array-like
    mapping : Mapping or Enum, optional
            A mapping. If an Enum class is passed, the
            __members__ is used.

    Returns
    -------
    array-like
    """
    if isinstance(mapping, enum.EnumMeta):
        mapping = mapping.__members__
    inv_map = {v: k for k, v in mapping.items()}
    return s.map(inv_map)
