# Utility methods for dealing with Dataset objects

import logging
import warnings
from typing import Any, Dict, List, Union

import numpy as np
import pandas as pd
from scipy import stats

logger = logging.getLogger(__name__)

try:
    import sqlalchemy
    from sqlalchemy.engine.default import DefaultDialect
    from sqlalchemy.sql.elements import WithinGroup
except ImportError:
    logger.debug("Unable to load SqlAlchemy or one of its subclasses.")
    DefaultDialect = None
    WithinGroup = None


SCHEMAS = {
    "api_np": {"NegativeInfinity": -np.inf, "PositiveInfinity": np.inf,},
    "api_cast": {"NegativeInfinity": -float("inf"), "PositiveInfinity": float("inf"),},
    "mysql": {"NegativeInfinity": -1.79e308, "PositiveInfinity": 1.79e308,},
    "mssql": {"NegativeInfinity": -1.79e308, "PositiveInfinity": 1.79e308,},
}


def get_sql_dialect_floating_point_infinity_value(
    schema: str, negative: bool = False
) -> float:
    res: Union[Dict, None] = SCHEMAS.get(schema)
    if res is None:
        if negative:
            return -np.inf
        else:
            return np.inf
    else:
        if negative:
            return res["NegativeInfinity"]
        else:
            return res["PositiveInfinity"]


def is_valid_partition_object(partition_object):
    """Tests whether a given object is a valid continuous or categorical partition object.
    :param partition_object: The partition_object to evaluate
    :return: Boolean
    """
    return is_valid_continuous_partition_object(
        partition_object
    ) or is_valid_categorical_partition_object(partition_object)


def is_valid_categorical_partition_object(partition_object):
    """Tests whether a given object is a valid categorical partition object.
    :param partition_object: The partition_object to evaluate
    :return: Boolean
    """
    if (
        partition_object is None
        or ("weights" not in partition_object)
        or ("values" not in partition_object)
    ):
        return False
    # Expect the same number of values as weights; weights should sum to one
    return len(partition_object["values"]) == len(
        partition_object["weights"]
    ) and np.allclose(np.sum(partition_object["weights"]), 1)


def is_valid_continuous_partition_object(partition_object):
    """Tests whether a given object is a valid continuous partition object. See :ref:`partition_object`.

    :param partition_object: The partition_object to evaluate
    :return: Boolean
    """
    if (
        (partition_object is None)
        or ("weights" not in partition_object)
        or ("bins" not in partition_object)
    ):
        return False

    if "tail_weights" in partition_object:
        if len(partition_object["tail_weights"]) != 2:
            return False
        comb_weights = partition_object["tail_weights"] + partition_object["weights"]
    else:
        comb_weights = partition_object["weights"]

    ## TODO: Consider adding this check to migrate to the tail_weights structure of partition objects
    # if (partition_object['bins'][0] == -np.inf) or (partition_object['bins'][-1] == np.inf):
    #     return False

    # Expect one more bin edge than weight; all bin edges should be monotonically increasing; weights should sum to one
    return (
        (len(partition_object["bins"]) == (len(partition_object["weights"]) + 1))
        and np.all(np.diff(partition_object["bins"]) > 0)
        and np.allclose(np.sum(comb_weights), 1.0)
    )


def categorical_partition_data(data):
    """Convenience method for creating weights from categorical data.

    Args:
        data (list-like): The data from which to construct the estimate.

    Returns:
        A new partition object::

            {
                "values": (list) The categorical values present in the data
                "weights": (list) The weights of the values in the partition.
            }

        See :ref:`partition_object`.
    """

    # Make dropna explicit (even though it defaults to true)
    series = pd.Series(data)
    value_counts = series.value_counts(dropna=True)

    # Compute weights using denominator only of nonnull values
    null_indexes = series.isnull()
    nonnull_count = (null_indexes == False).sum()

    weights = value_counts.values / nonnull_count
    return {"values": value_counts.index.tolist(), "weights": weights}


def kde_partition_data(data, estimate_tails=True):
    """Convenience method for building a partition and weights using a gaussian Kernel Density Estimate and default bandwidth.

    Args:
        data (list-like): The data from which to construct the estimate
        estimate_tails (bool): Whether to estimate the tails of the distribution to keep the partition object finite

    Returns:
        A new partition_object::

        {
            "bins": (list) The endpoints of the partial partition of reals,
            "weights": (list) The densities of the bins implied by the partition.
        }

        See :ref:`partition_object`.
    """
    kde = stats.kde.gaussian_kde(data)
    evaluation_bins = np.linspace(
        start=np.min(data) - (kde.covariance_factor() / 2),
        stop=np.max(data) + (kde.covariance_factor() / 2),
        num=np.floor(
            ((np.max(data) - np.min(data)) / kde.covariance_factor()) + 1
        ).astype(int),
    )
    cdf_vals = [kde.integrate_box_1d(-np.inf, x) for x in evaluation_bins]
    evaluation_weights = np.diff(cdf_vals)

    if estimate_tails:
        bins = np.concatenate(
            (
                [np.min(data) - (1.5 * kde.covariance_factor())],
                evaluation_bins,
                [np.max(data) + (1.5 * kde.covariance_factor())],
            )
        )
    else:
        bins = np.concatenate(([-np.inf], evaluation_bins, [np.inf]))

    weights = np.concatenate(([cdf_vals[0]], evaluation_weights, [1 - cdf_vals[-1]]))

    return {"bins": bins, "weights": weights}


def partition_data(data, bins="auto", n_bins=10):
    warnings.warn(
        "partition_data is deprecated and will be removed. Use either continuous_partition_data or \
                    categorical_partition_data instead.",
        DeprecationWarning,
    )
    return continuous_partition_data(data, bins, n_bins)


def continuous_partition_data(data, bins="auto", n_bins=10, **kwargs):
    """Convenience method for building a partition object on continuous data

    Args:
        data (list-like): The data from which to construct the estimate.
        bins (string): One of 'uniform' (for uniformly spaced bins), 'ntile' (for percentile-spaced bins), or 'auto'
            (for automatically spaced bins)
        n_bins (int): Ignored if bins is auto.
        kwargs (mapping): Additional keyword arguments to be passed to numpy histogram

    Returns:
        A new partition_object::

        {
            "bins": (list) The endpoints of the partial partition of reals,
            "weights": (list) The densities of the bins implied by the partition.
        }
        See :ref:`partition_object`.
    """
    if bins == "uniform":
        bins = np.linspace(start=np.min(data), stop=np.max(data), num=n_bins + 1)
    elif bins == "ntile":
        bins = np.percentile(data, np.linspace(start=0, stop=100, num=n_bins + 1))
    elif bins != "auto":
        raise ValueError("Invalid parameter for bins argument")

    try:
        hist, bin_edges = np.histogram(data, bins, density=False, **kwargs)
    except ValueError as e:
        raise ValueError(
            "Unable to compute histogram. Did you know you can pass additional kwargs to numpy histogram,"
            "such as a range? Numpy error was: " + str(e)
        )
    except TypeError as e:
        raise TypeError(
            "Unable to compute histogram. numpy histogram raised error: " + str(e)
        )

    return {"bins": bin_edges, "weights": hist / len(data)}


def build_continuous_partition_object(
    dataset, column, bins="auto", n_bins=10, allow_relative_error=False
):
    """Convenience method for building a partition object on continuous data from a dataset and column

    Args:
        dataset (GE Dataset): the dataset for which to compute the partition
        column (string): The name of the column for which to construct the estimate.
        bins (string): One of 'uniform' (for uniformly spaced bins), 'ntile' (for percentile-spaced bins), or 'auto'
            (for automatically spaced bins)
        n_bins (int): Ignored if bins is auto.
        allow_relative_error: passed to get_column_quantiles, set to False for only precise
            values, True to allow approximate values on systems with only binary choice (e.g. Redshift), and to a
            value between zero and one for systems that allow specification of relative error (e.g.
            SparkDFDataset).

    Returns:

        A new partition_object::

            {
                "bins": (list) The endpoints of the partial partition of reals,
                "weights": (list) The densities of the bins implied by the partition.
            }

            See :ref:`partition_object`.
    """
    bins = dataset.get_column_partition(column, bins, n_bins, allow_relative_error)
    if isinstance(bins, np.ndarray):
        bins = bins.tolist()
    else:
        bins = list(bins)
    weights = list(
        np.array(dataset.get_column_hist(column, tuple(bins)))
        / dataset.get_column_nonnull_count(column)
    )
    tail_weights = (1 - sum(weights)) / 2
    partition_object = {
        "bins": bins,
        "weights": weights,
        "tail_weights": [tail_weights, tail_weights],
    }
    return partition_object


def build_categorical_partition_object(dataset, column, sort="value"):
    """Convenience method for building a partition object on categorical data from a dataset and column

    Args:
        dataset (GE Dataset): the dataset for which to compute the partition
        column (string): The name of the column for which to construct the estimate.
        sort (string): must be one of "value", "count", or "none".
            - if "value" then values in the resulting partition object will be sorted lexigraphically
            - if "count" then values will be sorted according to descending count (frequency)
            - if "none" then values will not be sorted

    Returns:
        A new partition_object::

        {
            "values": (list) the categorical values for which each weight applies,
            "weights": (list) The densities of the values implied by the partition.
        }
        See :ref:`partition_object`.
    """
    counts = dataset.get_column_value_counts(column, sort)
    return {
        "values": list(counts.index),
        "weights": list(np.array(counts) / dataset.get_column_nonnull_count(column)),
    }


def infer_distribution_parameters(data, distribution, params=None):
    """Convenience method for determining the shape parameters of a given distribution

    Args:
        data (list-like): The data to build shape parameters from.
        distribution (string): Scipy distribution, determines which parameters to build.
        params (dict or None): The known parameters. Parameters given here will not be altered. \
                               Keep as None to infer all necessary parameters from the data data.

    Returns:
        A dictionary of named parameters::

        {
            "mean": (float),
            "std_dev": (float),
            "loc": (float),
            "scale": (float),
            "alpha": (float),
            "beta": (float),
            "min": (float),
            "max": (float),
            "df": (float)
        }

        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.kstest.html#scipy.stats.kstest
    """

    if params is None:
        params = dict()
    elif not isinstance(params, dict):
        raise TypeError(
            "params must be a dictionary object, see great_expectations documentation"
        )

    if "mean" not in params.keys():
        params["mean"] = data.mean()

    if "std_dev" not in params.keys():
        params["std_dev"] = data.std()

    if distribution == "beta":
        # scipy cdf(x, a, b, loc=0, scale=1)
        if "alpha" not in params.keys():
            # from https://stats.stackexchange.com/questions/12232/calculating-the-parameters-of-a-beta-distribution-using-the-mean-and-variance
            params["alpha"] = (params["mean"] ** 2) * (
                ((1 - params["mean"]) / params["std_dev"] ** 2) - (1 / params["mean"])
            )
        if "beta" not in params.keys():
            params["beta"] = params["alpha"] * ((1 / params["mean"]) - 1)

    elif distribution == "gamma":
        # scipy cdf(x, a, loc=0, scale=1)
        if "alpha" not in params.keys():
            # Using https://en.wikipedia.org/wiki/Gamma_distribution
            params["alpha"] = params["mean"] / params.get("scale", 1)

    # elif distribution == 'poisson':
    #    if 'lambda' not in params.keys():
    #       params['lambda'] = params['mean']

    elif distribution == "uniform":
        # scipy cdf(x, loc=0, scale=1)
        if "min" not in params.keys():
            if "loc" in params.keys():
                params["min"] = params["loc"]
            else:
                params["min"] = min(data)
        if "max" not in params.keys():
            if "scale" in params.keys():
                params["max"] = params["scale"]
            else:
                params["max"] = max(data) - params["min"]

    elif distribution == "chi2":
        # scipy cdf(x, df, loc=0, scale=1)
        if "df" not in params.keys():
            # from https://en.wikipedia.org/wiki/Chi-squared_distribution
            params["df"] = params["mean"]

    #  Expon only uses loc and scale, use default
    # elif distribution == 'expon':
    # scipy cdf(x, loc=0, scale=1)
    #    if 'lambda' in params.keys():
    # Lambda is optional
    #        params['scale'] = 1 / params['lambda']
    elif distribution != "norm":
        raise AttributeError(
            "Unsupported distribution type. Please refer to Great Expectations Documentation"
        )

    params["loc"] = params.get("loc", 0)
    params["scale"] = params.get("scale", 1)

    return params


def _scipy_distribution_positional_args_from_dict(distribution, params):
    """Helper function that returns positional arguments for a scipy distribution using a dict of parameters.

       See the `cdf()` function here https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.beta.html#Methods\
       to see an example of scipy's positional arguments. This function returns the arguments specified by the \
       scipy.stat.distribution.cdf() for tha distribution.

       Args:
           distribution (string): \
               The scipy distribution name.
           params (dict): \
               A dict of named parameters.

       Raises:
           AttributeError: \
               If an unsupported distribution is provided.
    """

    params["loc"] = params.get("loc", 0)
    if "scale" not in params:
        params["scale"] = 1

    if distribution == "norm":
        return params["mean"], params["std_dev"]
    elif distribution == "beta":
        return params["alpha"], params["beta"], params["loc"], params["scale"]
    elif distribution == "gamma":
        return params["alpha"], params["loc"], params["scale"]
    # elif distribution == 'poisson':
    #    return params['lambda'], params['loc']
    elif distribution == "uniform":
        return params["min"], params["max"]
    elif distribution == "chi2":
        return params["df"], params["loc"], params["scale"]
    elif distribution == "expon":
        return params["loc"], params["scale"]


def validate_distribution_parameters(distribution, params):
    """Ensures that necessary parameters for a distribution are present and that all parameters are sensical.

       If parameters necessary to construct a distribution are missing or invalid, this function raises ValueError\
       with an informative description. Note that 'loc' and 'scale' are optional arguments, and that 'scale'\
       must be positive.

       Args:
           distribution (string): \
               The scipy distribution name, e.g. normal distribution is 'norm'.
           params (dict or list): \
               The distribution shape parameters in a named dictionary or positional list form following the scipy \
               cdf argument scheme.

               params={'mean': 40, 'std_dev': 5} or params=[40, 5]

       Exceptions:
           ValueError: \
               With an informative description, usually when necessary parameters are omitted or are invalid.

    """

    norm_msg = (
        "norm distributions require 0 parameters and optionally 'mean', 'std_dev'."
    )
    beta_msg = "beta distributions require 2 positive parameters 'alpha', 'beta' and optionally 'loc', 'scale'."
    gamma_msg = "gamma distributions require 1 positive parameter 'alpha' and optionally 'loc','scale'."
    # poisson_msg = "poisson distributions require 1 positive parameter 'lambda' and optionally 'loc'."
    uniform_msg = (
        "uniform distributions require 0 parameters and optionally 'loc', 'scale'."
    )
    chi2_msg = "chi2 distributions require 1 positive parameter 'df' and optionally 'loc', 'scale'."
    expon_msg = (
        "expon distributions require 0 parameters and optionally 'loc', 'scale'."
    )

    if distribution not in [
        "norm",
        "beta",
        "gamma",
        "poisson",
        "uniform",
        "chi2",
        "expon",
    ]:
        raise AttributeError("Unsupported  distribution provided: %s" % distribution)

    if isinstance(params, dict):
        # `params` is a dictionary
        if params.get("std_dev", 1) <= 0 or params.get("scale", 1) <= 0:
            raise ValueError("std_dev and scale must be positive.")

        # alpha and beta are required and positive
        if distribution == "beta" and (
            params.get("alpha", -1) <= 0 or params.get("beta", -1) <= 0
        ):
            raise ValueError("Invalid parameters: %s" % beta_msg)

        # alpha is required and positive
        elif distribution == "gamma" and params.get("alpha", -1) <= 0:
            raise ValueError("Invalid parameters: %s" % gamma_msg)

        # lambda is a required and positive
        # elif distribution == 'poisson' and params.get('lambda', -1) <= 0:
        #    raise ValueError("Invalid parameters: %s" %poisson_msg)

        # df is necessary and required to be positve
        elif distribution == "chi2" and params.get("df", -1) <= 0:
            raise ValueError("Invalid parameters: %s:" % chi2_msg)

    elif isinstance(params, tuple) or isinstance(params, list):
        scale = None

        # `params` is a tuple or a list
        if distribution == "beta":
            if len(params) < 2:
                raise ValueError("Missing required parameters: %s" % beta_msg)
            if params[0] <= 0 or params[1] <= 0:
                raise ValueError("Invalid parameters: %s" % beta_msg)
            if len(params) == 4:
                scale = params[3]
            elif len(params) > 4:
                raise ValueError("Too many parameters provided: %s" % beta_msg)

        elif distribution == "norm":
            if len(params) > 2:
                raise ValueError("Too many parameters provided: %s" % norm_msg)
            if len(params) == 2:
                scale = params[1]

        elif distribution == "gamma":
            if len(params) < 1:
                raise ValueError("Missing required parameters: %s" % gamma_msg)
            if len(params) == 3:
                scale = params[2]
            if len(params) > 3:
                raise ValueError("Too many parameters provided: %s" % gamma_msg)
            elif params[0] <= 0:
                raise ValueError("Invalid parameters: %s" % gamma_msg)

        # elif distribution == 'poisson':
        #    if len(params) < 1:
        #        raise ValueError("Missing required parameters: %s" %poisson_msg)
        #   if len(params) > 2:
        #        raise ValueError("Too many parameters provided: %s" %poisson_msg)
        #    elif params[0] <= 0:
        #        raise ValueError("Invalid parameters: %s" %poisson_msg)

        elif distribution == "uniform":
            if len(params) == 2:
                scale = params[1]
            if len(params) > 2:
                raise ValueError("Too many arguments provided: %s" % uniform_msg)

        elif distribution == "chi2":
            if len(params) < 1:
                raise ValueError("Missing required parameters: %s" % chi2_msg)
            elif len(params) == 3:
                scale = params[2]
            elif len(params) > 3:
                raise ValueError("Too many arguments provided: %s" % chi2_msg)
            if params[0] <= 0:
                raise ValueError("Invalid parameters: %s" % chi2_msg)

        elif distribution == "expon":

            if len(params) == 2:
                scale = params[1]
            if len(params) > 2:
                raise ValueError("Too many arguments provided: %s" % expon_msg)

        if scale is not None and scale <= 0:
            raise ValueError("std_dev and scale must be positive.")

    else:
        raise ValueError(
            "params must be a dict or list, or use ge.dataset.util.infer_distribution_parameters(data, distribution)"
        )

    return


def create_multiple_expectations(df, columns, expectation_type, *args, **kwargs):
    """Creates an identical expectation for each of the given columns with the specified arguments, if any.

    Args:
        df (great_expectations.dataset): A great expectations dataset object.
        columns (list): A list of column names represented as strings.
        expectation_type (string): The expectation type.

    Raises:
        KeyError if the provided column does not exist.
        AttributeError if the provided expectation type does not exist or df is not a valid great expectations dataset.

    Returns:
        A list of expectation results.


    """
    expectation = getattr(df, expectation_type)
    results = list()

    for column in columns:
        results.append(expectation(column, *args, **kwargs))

    return results


def get_approximate_percentile_disc_sql(selects: List, sql_engine_dialect: Any) -> str:
    return ", ".join(
        [
            "approximate "
            + str(
                stmt.compile(
                    dialect=sql_engine_dialect, compile_kwargs={"literal_binds": True}
                )
            )
            for stmt in selects
        ]
    )


def check_sql_engine_dialect(
    actual_sql_engine_dialect: Any, candidate_sql_engine_dialect: Any,
) -> bool:
    try:
        # noinspection PyTypeChecker
        return isinstance(actual_sql_engine_dialect, candidate_sql_engine_dialect)
    except (AttributeError, TypeError):
        return False
