"""
**hep_ml.commonutils** contains some helpful functions and classes
which are often used (by other modules)
"""

from __future__ import print_function, division, absolute_import

from multiprocessing.pool import ThreadPool
import numbers
import itertools

import numpy
import pandas
from scipy.special import expit
try:
    import sklearn.model_selection as sklearn_cross_validation
except ImportError:
    # The cross_validation module is deprecated as of sklearn 0.18, in favour
    # of the model_selection module
    import sklearn.cross_validation as sklearn_cross_validation
from sklearn.neighbors.unsupervised import NearestNeighbors

__author__ = "Alex Rogozhnikov"


def _threads_wrapper(func_and_args):
    func = func_and_args[0]
    args = func_and_args[1:]
    return func(*args)


def map_on_cluster(parallel_profile, *args, **kw_args):
    """
    The same as map, but the first argument is ipc_profile. Distributes the task over IPython cluster.

    :param parallel_profile: the IPython cluster profile to use.
    :type parallel_profile: None or str
    :param list args: function, arguments
    :param dict kw_args: kwargs for LoadBalacedView.map_sync

    (function copied from REP)

    :return: the result of mapping
    """
    if parallel_profile is None:
        return map(*args)
    elif str.startswith(parallel_profile, 'threads-'):
        n_threads = int(parallel_profile[len('threads-'):])
        pool = ThreadPool(processes=n_threads)
        func, params = args[0], args[1:]
        return pool.map(_threads_wrapper, zip(itertools.cycle([func]), *params))
    else:
        from IPython.parallel import Client
        return Client(profile=parallel_profile).load_balanced_view().map_sync(*args, **kw_args)


def sigmoid_function(x, width):
    """ Sigmoid function is smoothing of Heaviside function,
    the less width, the closer we are to Heaviside function
    :type x: array-like with floats, arbitrary shape
    :type width: float, if width == 0, this is simply Heaviside function
    """
    assert width >= 0, 'the width should be non-negative'
    if width > 0.0001:
        return expit(x / width)
    else:
        return (x > 0) * 1.0


def generate_sample(n_samples, n_features, distance=2.0):
    """Generates some test distribution,
    signal and background distributions are gaussian with same dispersion and different centers,
    all variables are independent (gaussian correlation matrix is identity).

    This function is frequently used in tests. """
    from sklearn.datasets import make_blobs

    centers = numpy.zeros((2, n_features))
    centers[0, :] = - distance / 2
    centers[1, :] = distance / 2

    X, y = make_blobs(n_samples=n_samples, n_features=n_features, centers=centers)
    columns = ["column" + str(x) for x in range(n_features)]
    X = pandas.DataFrame(X, columns=columns)
    return X, y


def check_uniform_label(uniform_label):
    """Convert uniform label to numpy.array

    :param uniform_label: label or list of labels (examples: 0, 1, [0], [1], [0, 1])
    :return: numpy.array (with [0], [1] or [0, 1])
    """
    if isinstance(uniform_label, numbers.Number):
        return numpy.array([uniform_label])
    else:
        return numpy.array(uniform_label)


def train_test_split(*arrays, **kw_args):
    """Does the same thing as train_test_split, but preserves columns in DataFrames.
    Uses the same parameters: test_size, train_size, random_state, and has the same interface

    :type list[numpy.array|pandas.DataFrame] arrays: arrays to split
    """
    assert len(arrays) > 0, "at least one array should be passed"
    length = len(arrays[0])
    for array in arrays:
        assert len(array) == length, "different size"
    train_indices, test_indices = sklearn_cross_validation.train_test_split(range(length), **kw_args)
    result = []
    for array in arrays:
        if isinstance(array, pandas.DataFrame):
            result.append(array.iloc[train_indices, :])
            result.append(array.iloc[test_indices, :])
        else:
            result.append(array[train_indices])
            result.append(array[test_indices])
    return result


def weighted_quantile(array, quantiles, sample_weight=None, array_sorted=False, old_style=False):
    """ Very close to numpy.percentile, but supports weights.
    NOTE: quantiles should be in [0, 1]!
    :param array: numpy.array with data
    :param quantiles: array-like with many percentiles
    :param sample_weight: array-like of the same length as `array`
    :param array_sorted: bool, if True, then will avoid sorting
    :param old_style: if True, will correct output to be consistent with numpy.percentile.
    :return: numpy.array with computed percentiles.
    """
    array = numpy.array(array)
    quantiles = numpy.array(quantiles)
    sample_weight = check_sample_weight(array, sample_weight)
    assert numpy.all(quantiles >= 0) and numpy.all(quantiles <= 1), 'Percentiles should be in [0, 1]'

    if not array_sorted:
        sorter = numpy.argsort(array)
        array, sample_weight = array[sorter], sample_weight[sorter]

    weighted_quantiles = numpy.cumsum(sample_weight) - 0.5 * sample_weight
    if old_style:
        # To be convenient with numpy.percentile
        weighted_quantiles -= weighted_quantiles[0]
        weighted_quantiles /= weighted_quantiles[-1]
    else:
        weighted_quantiles /= numpy.sum(sample_weight)
    return numpy.interp(quantiles, weighted_quantiles, array)


def build_normalizer(signal, sample_weight=None):
    """Prepares normalization function for some set of values
    transforms it to uniform distribution from [0, 1]. Example of usage:
    >>>normalizer = build_normalizer(signal)
    >>>pylab.hist(normalizer(background))
    >>># this one should be uniform in [0,1]
    >>>pylab.hist(normalizer(signal))

    :param numpy.array signal: shape = [n_samples] with floats
    :param numpy.array sample_weight: shape = [n_samples], non-negative weights associated to events.
    """
    sample_weight = check_sample_weight(signal, sample_weight)
    assert numpy.all(sample_weight >= 0.), 'sample weight must be non-negative'
    sorter = numpy.argsort(signal)
    signal, sample_weight = signal[sorter], sample_weight[sorter]
    predictions = numpy.cumsum(sample_weight) / numpy.sum(sample_weight)

    def normalizing_function(data):
        return numpy.interp(data, signal, predictions)

    return normalizing_function


def compute_cut_for_efficiency(efficiency, mask, y_pred, sample_weight=None):
    """ Computes such cut(s), that provide given target global efficiency(ies).
    Example:
    >>> p = classifier.predict_proba(X)
    >>> threshold = compute_cut_for_efficiency(0.5, mask=y == 1, y_pred=p[:, 1])

    :type efficiency: float or numpy.array with target efficiencies, shape = [n_effs]
    :type mask: array-like, shape = [n_samples], True for needed samples
    :type y_pred: array-like, shape = [n_samples], predictions or scores (float)
    :type sample_weight: None | array-like, shape = [n_samples]
    :return: float or numpy.array, shape = [n_effs]
    """
    sample_weight = check_sample_weight(mask, sample_weight)
    assert len(mask) == len(y_pred), 'lengths are different'
    efficiency = numpy.array(efficiency)
    is_signal = mask > 0.5
    y_pred, sample_weight = y_pred[is_signal], sample_weight[is_signal]
    return weighted_quantile(y_pred, 1. - efficiency, sample_weight=sample_weight)


# region Knn-related functions


def compute_knn_indices_of_signal(X, is_signal, n_neighbours=50):
    """For each event returns the knn closest signal(!) events. No matter of what class the event is.

    :type X: numpy.array, shape = [n_samples, n_features] the distance is measured over these variables
    :type is_signal: numpy.array, shape = [n_samples] with booleans
    :rtype numpy.array, shape [len(dataframe), knn], each row contains indices of closest signal events
    """
    assert len(X) == len(is_signal), "Different lengths"
    signal_indices = numpy.where(is_signal)[0]
    X_signal = numpy.array(X)[numpy.array(is_signal)]
    neighbours = NearestNeighbors(n_neighbors=n_neighbours, algorithm='kd_tree').fit(X_signal)
    _, knn_signal_indices = neighbours.kneighbors(X)
    return numpy.take(signal_indices, knn_signal_indices)


def compute_knn_indices_of_same_class(X, y, n_neighbours=50):
    """