import copy
from collections import Iterable
import numpy as np
import pandas as pd
import dask
import dask.threaded
from dask import delayed
from dask.compatibility import reraise
from dask.threaded import pack_exception
from sklearn.preprocessing import Imputer

from . import time_series
from .time_series import TimeSeries
from .features import generate_dask_graph

__all__ = ['featurize_time_series', 'featurize_single_ts',
           'featurize_ts_files', 'assemble_featureset']


def featurize_single_ts(ts, features_to_use, custom_script_path=None,
                        custom_functions=None, raise_exceptions=True):
    """Compute feature values for a given single time-series. Data is
    returned as dictionaries/lists of lists.

    Parameters
    ----------
    ts : TimeSeries object
        Single time series to be featurized.
    features_to_use : list of str
        List of feature names to be generated.
    custom_functions : dict, optional
        Dictionary of custom feature functions to be evaluated for the given
        time series, or a dictionary representing a dask graph of function
        evaluations. Dictionaries of functions should have keys `feature_name`
        and values functions that take arguments (t, m, e); in the case of a
        dask graph, these arrays should be referenced as 't', 'm', 'e',
        respectively, and any values with keys present in `features_to_use`
        will be computed.
    raise_exceptions : bool, optional
        If True, exceptions during feature computation are raised immediately;
        if False, exceptions are supressed and `np.nan` is returned for the
        given feature and any dependent features. Defaults to True.

    Returns
    -------
    dict
        Dictionary with feature names as keys, lists of feature values (one per
        channel) as values.
    """
    # Initialize empty feature array for all channels
    feature_values = np.empty((len(features_to_use), ts.n_channels))
    for (t_i, m_i, e_i), i in zip(ts.channels(), range(ts.n_channels)):
        feature_graph = generate_dask_graph(t_i, m_i, e_i)
        feature_graph.update(ts.meta_features)

        if custom_functions:
            # If values in custom_functions are functions, add calls to graph
            if all(hasattr(v, '__call__') for v in custom_functions.values()):
                feature_graph.update({feat: f(t_i, m_i, e_i)
                                      for feat, f in custom_functions.items()})
            # Otherwise, custom_functions is another dask graph
            else:
                feature_graph.update(custom_functions)

        # Do not execute in parallel; parallelization has already taken place
        # at the level of time series, so we compute features for a single time
        # series in serial.
        if raise_exceptions:
            raise_callback = reraise
        else:
            raise_callback = lambda e, tb: None
        dask_values = dask.get(feature_graph, features_to_use,
                               raise_exception=raise_callback,
                               pack_exception=pack_exception)
        feature_values[:, i] = [x if not isinstance(x, Exception) else np.nan
                                for x in dask_values]
    index = pd.MultiIndex.from_product((features_to_use, range(ts.n_channels)),
                                       names=('feature', 'channel'))
    return pd.Series(feature_values.ravel(), index=index)


def assemble_featureset(features_list, time_series=None,
                        meta_features_list=None, names=None):
    """Transforms raw feature data (as returned by `featurize_single_ts`) into
    a pd.DataFrame.

    Parameters
    ----------
    features_list : list of pd.Series
        List of series (one per time series file) with (feature name, channel)
        multiindex.
    time_series : list of TimeSeries
        If provided, the name and metafeatures from the time series objects
        will be used, overriding the `meta_features_list` and `names` values.
    meta_features_list : list of dict
        If provided, the columns of `metadata` will be added to the featureset.
    names : list of str
        If provided, the (row) index of the featureset will be set accordingly.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns containing feature values, indexed by name.
    """
    if time_series is not None:
        meta_features_list, names = zip(*[(ts.meta_features, ts.name)
                                          for ts in time_series])
    if len(features_list) > 0:
        feat_df = pd.concat(features_list, axis=1, ignore_index=True).T
        feat_df.index = names
    else:
        feat_df = pd.DataFrame(index=names)

    if meta_features_list and any(meta_features_list):  # not all empty dicts
        meta_df = pd.DataFrame(list(meta_features_list), index=names)
        meta_df.columns = pd.MultiIndex.from_tuples([(c, '') for c in meta_df],
                                                    names=['feature', 'channel'])
        feat_df = pd.concat((feat_df, meta_df), axis=1)

    return feat_df


# TODO should this be changed to use TimeSeries objects? or maybe an optional
# argument for TimeSeries? some redundancy here...
def featurize_time_series(times, values, errors=None, features_to_use=[],
                          meta_features={}, names=None,
                          custom_script_path=None, custom_functions=None,
                          scheduler=dask.threaded.get, raise_exceptions=True):
    """Versatile feature generation function for one or more time series.

    For a single time series, inputs may have the form:

    - `times`:  (n,) array or (p, n) array (for p channels of measurement)
    - `values`: (n,) array or (p, n) array (for p channels of measurement)
    - `errors`: (n,) array or (p, n) array (for p channels of measurement)

    For multiple time series, inputs may have the form:

    - `times`: list of (n,) arrays, list of (p, n) arrays (for p channels of
      measurement), or list of lists of (n,) arrays (for
      multichannel data with different time values per channel)
    - `values`: list of (n,) arrays, list of (p, n) arrays (for p channels of
      measurement), or list of lists of (n,) arrays (for
      multichannel data with different time values per channel)
    - `errors`: list of (n,) arrays, list of (p, n) arrays (for p channels of
      measurement), or list of lists of (n,) arrays (for
      multichannel data with different time values per channel)

    In the case of multichannel measurements, each channel will be
    featurized separately, and the index of the output featureset will contain
    a `channel` coordinate.

    Parameters
    ----------
    times : array, list of array, or list of lists of array
        Array containing time values for a single time series, or a list of
        arrays each containing time values for a single time series, or a list
        of lists of arrays for multichannel data with different time values per
        channel
    values : array or list of array
        Array containing measurement values for a single time series, or a list
        of arrays each containing (possibly multivariate) measurement values
        for a single time series, or a list of lists of arrays for multichannel
        data with different time values per channel
    errors : array or list/tuple of array, optional
        Array containing measurement error values for a single time series, or
        a list of arrays each containing (possibly multivariate) measurement
        values for a single time series, or a list of lists of arrays for
        multichannel data with different time values per channel
    features_to_use : list of str, optional
        List of feature names to be generated. Defaults to an empty list, which
        will result in only meta_features features being stored.
    meta_features : dict/Pandas.Series or list of dicts/Pandas.DataFrame
        dict/Series (for a single time series) or DataFrame (for multiple time
        series) of metafeature information; features are added to the output
        featureset, and their values are consumable by custom feature scripts.
    names : str or list of str, optional
        Name or list of names for each time series, if applicable; will be
        stored in the (row) index of the featureset.
    custom_script_path : str, optional
        Path to Python script containing function definitions for the
        generation of any custom features. Defaults to None.
    custom_functions : dict, optional
        Dictionary of custom feature functions to be evaluated for the given
        time series, or a dictionary representing a dask graph of function
        evaluations.  Dictionaries of functions should have keys `feature_name`
        and values functions that take arguments (t, m, e); in the case of a
        dask graph, these arrays should be referenced as 't', 'm', 'e',
        respectively, and any values with keys present in `features_to_use`
        will be computed.
    scheduler : function, optional
        `dask` scheduler function used to perform feature extraction
        computation. Defaults to `dask.threaded.get`.
    raise_exceptions : bool, optional
        If True, exceptions during feature computation are raised immediately;
        if False, exceptions are supressed and `np.nan` is returned for the
        given feature and any dependent features. Defaults to True.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns containing feature values, indexed by name.
    """
    if times is None:
        times = copy.deepcopy(values)
        if isinstance(times, np.ndarray) and (times.ndim == 1
                                              or 1 in times.shape):
            times[:] = np.linspace(0., time_series.DEFAULT_MAX_TIME,
                                   times.size)
        else:
            for t in times:
                if isinstance(t, np.ndarray) and (t.ndim == 1 or 1 in t.shape):
                    t[:] = np.linspace(0., time_series.DEFAULT_MAX_TIME,
                                       t.size)
                else:
                    for t_i in t:
                        t_i[:] = np.linspace(0., time_series.DEFAULT_MAX_TIME,
                                             t_i.size)

    if errors is None:
        errors = copy.deepcopy(values)
        if isinstance(errors, np.ndarray) and (errors.ndim == 1
                                               or 1 in errors.shape):
            errors[:] = time_series.DEFAULT_ERROR_VALUE
        else:
            for e in errors:
                if isinstance(e, np.ndarray) and (e.ndim == 1 or 1 in e.shape):
                    e[:] = time_series.DEFAULT_ERROR_VALUE
                else:
                    for e_i in e:
                        e_i[:] = time_series.DEFAULT_ERROR_VALUE

    # One single-channel time series:
    if not isinstance(values[0], Iterable):
        times, values, errors = [times], [values], [errors]
    # One multi-channel time series:
    elif isinstance(values, np.ndarray) and values.ndim == 2:
        times, values, errors = [times], [values], [errors]

    if names is None:
        names = np.arange(len(times))

    if isinstance(meta_features, pd.Series):
        meta_features = meta_features.to_dict()
    meta_features = pd.DataFrame(meta_features, index=names)

    all_time_series = [delayed(TimeSeries(t, m, e,
                                          meta_features=meta_features.loc[name],
                                          name=name), pure=True)
                       for t, m, e, name in zip(times, values, errors, names)]

    all_features = [delayed(featurize_single_ts, pure=True)(ts, features_to_use,
                                                            custom_script_path,
                                                            custom_functions,
                                                            raise_exceptions)
                    for ts in all_time_series]
    result = delayed(assemble_featureset, pure=True)(all_features, all_time_series)
    return result.compute(get=scheduler)


def featurize_ts_files(ts_paths, features_to_use, custom_script_path=None,
                       custom_functions=None, scheduler=dask.threaded.get,
                       raise_exceptions=True):
    """Feature generation function for on-disk time series (.npz) files.

    By default, computes features concurrently using the
    `dask.threaded.get` scheduler. Other possible options include
    `dask.local.get` for synchronous computation (e.g., when debugging),
    or `dask.distributed.Executor.get` for distributed computation.

    In the case of multichannel measurements, each channel will be
    featurized separately, and the index of the output featureset will contain
    a `channel` coordinate.

    Parameters
    ----------
    ts_paths : list of str
        List of paths to time series data, stored in `numpy` .npz format.
        See `time_series.load` for details.
    features_to_use : list of str, optional
        List of feature names to be generated. Defaults to an empty list, which
        will result in only meta_features features being stored.
    custom_script_path : str, optional
        Path to Python script containing function definitions for the
        generation of any custom features. Defaults to None.
    custom_functions : dict, optional
        Dictionary of custom feature functions to be evaluated for the given
        time series, or a dictionary representing a dask graph of function
        evaluations.  Dictionaries of functions should have keys `feature_name`
        and values functions that take arguments (t, m, e); in the case of a
        dask graph, these arrays should be referenced as 't', 'm', 'e',
        respectively, and any values with keys present in `features_to_use`
        will be computed.
    scheduler : function, optional
        `dask` scheduler function used to perform feature extraction
        computation. Defaults to `dask.threaded.get`.
    raise_exceptions : bool, optional
        If True, exceptions during feature computation are raised immediately;
        if False, exceptions are supressed and `np.nan` is returned for the
        given feature and any dependent features. Defaults to True.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns containing feature values, indexed by name.
    """
    all_time_series = [delayed(time_series.load, pure=True)(ts_path)
                       for ts_path in ts_paths]
    all_features = [delayed(featurize_single_ts, pure=True)(ts, features_to_use,
                                                            custom_script_path,
                                                            custom_functions,
                                                            raise_exceptions)
                    for ts in all_time_series]
    names, meta_feats, all_labels = zip(*[(ts.name, ts.meta_features, ts.label)
                                          for ts in all_time_series])
    result = delayed(assemble_featureset, pure=True)(all_features,
                                                     meta_features_list=meta_feats,
                                                     names=names)
    fset, labels = dask.compute(result, all_labels, get=scheduler)