from inferelator.distributed.inferelator_mp import MPControl
from inferelator.regression import base_regression
from inferelator import utils
import copy

import numpy as np
import scipy.sparse as sps
from dask import distributed

"""
This package contains the dask-specific multiprocessing functions (these are used in place of map calls to allow the
more advanced memory and task tools of dask to be used)
"""

DASK_SCATTER_TIMEOUT = 120


def amusr_regress_dask(X, Y, priors, prior_weight, n_tasks, genes, tfs, G, remove_autoregulation=True,
                       lambda_Bs=None, lambda_Ss=None, Cs=None, Ss=None, regression_function=None):
    """
    Execute multitask (AMUSR)

    :return: list
        Returns a list of regression results that the amusr_regression pileup_data can process
    """

    assert MPControl.is_dask()

    from inferelator.regression.amusr_regression import format_prior, run_regression_EBIC
    DaskController = MPControl.client

    # Allows injecting a regression function for testing
    regression_function = run_regression_EBIC if regression_function is None else regression_function

    # Gets genes, n_tasks, prior_weight, and remove_autoregulation from regress_dask()
    # Other arguments are passed in
    def regression_maker(j, x_df, y_list, prior, tf):
        level = 0 if j % 100 == 0 else 2
        utils.Debug.allprint(base_regression.PROGRESS_STR.format(gn=genes[j], i=j, total=G),
                             level=level)

        gene = genes[j]
        x, y, tasks = [], [], []

        if remove_autoregulation:
            tf = [t for t in tf if t != gene]
        else:
            pass

        for k, y_data in y_list:
            x.append(x_df[k].get_gene_data(tf))  # list([N, K])
            y.append(y_data)
            tasks.append(k)  # [T,]

        prior = format_prior(prior, gene, tasks, prior_weight)
        return j, regression_function(x, y, tf, tasks, gene, prior,
                                      lambda_Bs=lambda_Bs, lambda_Ss=lambda_Ss, Cs=Cs, Ss=Ss)

    def response_maker(y_df, i):
        y = []
        gene = genes[i]
        for k in range(n_tasks):
            if gene in y_df[k].gene_names:
                y.append((k, y_df[k].get_gene_data(gene, force_dense=True).reshape(-1, 1)))
        return y

    # Scatter common data to workers
    [scatter_x] = DaskController.client.scatter([X], broadcast=True, hash=False)
    [scatter_priors] = DaskController.client.scatter([priors], broadcast=True, hash=False)

    # Wait for scattering to finish before creating futures
    distributed.wait(scatter_x, timeout=DASK_SCATTER_TIMEOUT)
    distributed.wait(scatter_priors, timeout=DASK_SCATTER_TIMEOUT)

    future_list = [DaskController.client.submit(regression_maker, i, scatter_x, response_maker(Y, i), scatter_priors,
                                                tfs)
                   for i in range(G)]

    # Collect results as they finish instead of waiting for all workers to be done
    result_list = process_futures_into_list(future_list)

    DaskController.client.cancel(scatter_x)
    DaskController.client.cancel(scatter_priors)

    return result_list


def bbsr_regress_dask(X, Y, pp_mat, weights_mat, G, genes, nS):
    """
    Execute regression (BBSR)

    :return: list
        Returns a list of regression results that the pileup_data can process
    """
    assert MPControl.is_dask()

    from inferelator.regression import bayes_stats
    DaskController = MPControl.client

    def regression_maker(j, x, y, pp, weights):
        level = 0 if j % 100 == 0 else 2
        utils.Debug.allprint(base_regression.PROGRESS_STR.format(gn=genes[j], i=j, total=G), level=level)
        data = bayes_stats.bbsr(x, utils.scale_vector(y), pp[j, :].flatten(), weights[j, :].flatten(), nS)
        data['ind'] = j
        return j, data

    # Scatter common data to workers
    [scatter_x] = DaskController.client.scatter([X.values], broadcast=True, hash=False)
    [scatter_pp] = DaskController.client.scatter([pp_mat.values], broadcast=True, hash=False)
    [scatter_weights] = DaskController.client.scatter([weights_mat.values], broadcast=True, hash=False)

    # Wait for scattering to finish before creating futures
    distributed.wait(scatter_x, timeout=DASK_SCATTER_TIMEOUT)
    distributed.wait(scatter_pp, timeout=DASK_SCATTER_TIMEOUT)
    distributed.wait(scatter_weights, timeout=DASK_SCATTER_TIMEOUT)

    future_list = [DaskController.client.submit(regression_maker, i, scatter_x,
                                                Y.get_gene_data(i, force_dense=True).flatten(),
                                                scatter_pp, scatter_weights)
                   for i in range(G)]

    # Collect results as they finish instead of waiting for all workers to be done
    result_list = process_futures_into_list(future_list)

    DaskController.client.cancel(scatter_x)
    DaskController.client.cancel(scatter_pp)
    DaskController.client.cancel(scatter_weights)

    return result_list


def sklearn_regress_dask(X, Y, model, G, genes, min_coef):
    """
    Execute regression (SKLearn)

    :return: list
        Returns a list of regression results that the pileup_data can process
    """
    assert MPControl.is_dask()

    from inferelator.regression import sklearn_regression
    DaskController = MPControl.client

    def regression_maker(j, x, y):
        level = 0 if j % 100 == 0 else 2
        utils.Debug.allprint(base_regression.PROGRESS_STR.format(gn=genes[j], i=j, total=G), level=level)
        data = sklearn_regression.sklearn_gene(x, utils.scale_vector(y), copy.copy(model))
        data['ind'] = j
        return j, data

    # Scatter common data to workers
    [scatter_x] = DaskController.client.scatter([X.values], broadcast=True, hash=False)

    # Wait for scattering to finish before creating futures
    distributed.wait(scatter_x, timeout=DASK_SCATTER_TIMEOUT)

    future_list = [DaskController.client.submit(regression_maker, i, scatter_x,
                                                Y.get_gene_data(i, force_dense=True).flatten())
                   for i in range(G)]

    # Collect results as they finish instead of waiting for all workers to be done
    result_list = process_futures_into_list(future_list)

    DaskController.client.cancel(scatter_x)

    return result_list


def lasso_stars_regress_dask(X, Y, alphas, num_subsamples, random_seed, method, params, G, genes):
    """
    Execute regression (LASSO-StARS)

    :return: list
        Returns a list of regression results that the pileup_data can process
    """
    assert MPControl.is_dask()

    from inferelator.regression import stability_selection
    DaskController = MPControl.client

    def regression_maker(j, x, y):
        level = 0 if j % 100 == 0 else 2
        utils.Debug.allprint(base_regression.PROGRESS_STR.format(gn=genes[j], i=j, total=G), level=level)
        data = stability_selection.stars_model_select(x, utils.scale_vector(y), alphas, num_subsamples=num_subsamples,
                                                      method=method, random_seed=random_seed, **params)
        data['ind'] = j
        return j, data

    # Scatter common data to workers
    [scatter_x] = DaskController.client.scatter([X.values], broadcast=True, hash=False)

    # Wait for scattering to finish before creating futures
    distributed.wait(scatter_x, timeout=DASK_SCATTER_TIMEOUT)

    future_list = [DaskController.client.submit(regression_maker, i, scatter_x,
                                                Y.get_gene_data(i, force_dense=True).flatten())
                   for i in range(G)]

    # Collect results as they finish instead of waiting for all workers to be done
    result_list = process_futures_into_list(future_list)

    DaskController.client.cancel(scatter_x)

    return result_list


def build_mi_array_dask(X, Y, bins, logtype):
    """
    Calculate MI into an array with dask (the naive map is very inefficient)

    :param X: np.ndarray (n x m1)
        Discrete array of bins
    :param Y: np.ndarray (n x m2)
        Discrete array of bins
    :param bins: int
        The total number of bins that were used to make the arrays discrete
    :param logtype: np.log func
        Which log function to use (log2 gives bits, ln gives nats)
    :return mi: np.ndarray (m1 x m2)
        Returns the mutual information array
    """

    assert MPControl.is_dask()

    from inferelator.regression.mi import _calc_mi, _make_table, _make_discrete

    # Get a reference to the Dask controller
    DaskController = MPControl.client

    m1, m2 = X.shape[1], Y.shape[1]

    def mi_make(i, x, y):
        x = _make_discrete(x, bins)
        return i, [_calc_mi(_make_table(x, y[:, j], bins), logtype=logtype) for j in range(m2)]

    # Scatter Y to workers and keep track as Futures
    [scatter_y] = DaskController.client.scatter([Y], broadcast=True, hash=False)

    # Wait for scattering to finish before creating futures
    distributed.wait(scatter_y, timeout=DASK_SCATTER_TIMEOUT)

    # Build an asynchronous list of Futures for each calculation of mi_make
    future_list = [DaskController.client.submit(mi_make, i,
                                                X[:, i].A.flatten() if sps.isspmatrix(X) else X[:, i].flatten(),
                                                scatter_y)
                   for i in range(m1)]

    # Collect results as they finish instead of waiting for all workers to be done
    mi_list = process_futures_into_list(future_list)

    # Convert the list of lists to an array
    mi = np.array(mi_list)
    assert (m1, m2) == mi.shape, "Array {sh} produced [({m1}, {m2}) expected]".format(sh=mi.shape, m1=m1, m2=m2)

    DaskController.client.cancel(scatter_y)

    return mi


def dask_map(func, *args, **kwargs):
    """
    Dask map

    :param func: function to map
    :type func: callable
    :param args: positional arguments for func
    :type args: iterable
    :param kwargs: keyword (non-iterable) arguments for func. Keywords will be passed to dask client.map.
    :return: List of mapped results
    :rtype: list
    """

    assert MPControl.is_dask()

    def _func_caller(f, i, *a, **k):
        return i, f(*a, **k)

    return process_futures_into_list([MPControl.client.client.submit(_func_caller, func, i, *za, **kwargs)
                                      for i, za in enumerate(zip(*args))])


def process_futures_into_list(future_list, raise_on_error=True, check_results=True):
    """
    Take a list of futures and turn them into a list of results
    Results must be of the form i, data (where i is the output order)

    :param future_list: A list of executing futures
    :type future_list: list
    :param raise_on_error: Should an error be raised if a job can't be restarted or just move on from it.
    Defaults to True
    :type raise_on_error: bool
    :param check_results: Should the result object be checked (and restarted if there's a problem)
    If False, this will raise an error with the result of a failed future is retrieved.
    Defaults to True.
    :type check_results: bool
    :return output_list: A list of results from the completed futures
    :rtype: list
    """

    DaskController = MPControl.client
    output_list = [None] * len(future_list)
    complete_gen = distributed.as_completed(future_list)

    for finished_future in complete_gen:

        DaskController.check_cluster_state()

        # Jobs can be cancelled in certain situations
        if check_results and (finished_future.cancelled() or (finished_future.status == "erred")):
            error = finished_future.exception()
            utils.Debug.vprint("Restarting job (Error: {er})".format(er=error), level=0)

            # Restart cancelled futures and put them back into the work pile
            try:
                DaskController.client.retry(finished_future)
                complete_gen.update([finished_future])
            except KeyError:
                if raise_on_error:
                    raise

        # In the event of success, get the data
        else:
            i, result_data = finished_future.result()
            output_list[i] = result_data
            finished_future.cancel()

    return output_list

