import copy
from collections import Iterable
import numpy as np


__all__ = ['load', 'TimeSeries', 'DEFAULT_MAX_TIME',
           'DEFAULT_ERROR_VALUE']


DEFAULT_MAX_TIME = 1.0
DEFAULT_ERROR_VALUE = 1e-4


def _ndim(x):
    """Return number of dimensions for a (possibly ragged) array."""
    n = 0
    while isinstance(x, Iterable):
        x = x[0]
        n += 1
    return n


def _compatible_shapes(x, y):
    """Check recursively that iterables x and y (and each iterable contained
    within, if applicable), have compatible sizes.
    """
    if hasattr(x, 'shape') and hasattr(y, 'shape'):
        return x.shape == y.shape
    else:
        return (len(x) == len(y) and all(np.shape(x_i) == np.shape(y_i)
                                         for x_i, y_i in zip(x, y)))


def _default_values_like(old_values, value=None, upper=None):
    """Creates a range of default values with the same shape as the input
    `old_values`. If `value` is provided then each entry will equal `value`;
    if `upper` is provided then the values will be linearly-spaced from 0 to
    `upper`.

    Parameters
    ----------
    old_values : (n,) or (p,n) array or list of (n,) arrays
        Input array(s), typically time series measurements for which default
        time or error values need to be inferred.
    value : float, optional
        Value that each output entry will be set to (omitted if `upper` is
        provided).
    upper : float, optional
        Upper bound of range of linearly-spaced output entries (omitted if
        `value` is provided).
    """
    if value and upper:
        raise ValueError("Only one of `value` or `upper` may be proivded.")
    elif value is not None:
        lower = value
        upper = value
    elif upper is not None:
        lower = 0.
    else:
        raise ValueError("Either `value` or `upper` must be provided.")

    new_values = copy.deepcopy(old_values)
    if _ndim(old_values) == 1 or (isinstance(old_values, np.ndarray) and 1 in
                                  old_values.shape):
        new_values[:] = np.linspace(lower, upper, len(new_values))
    else:
        for new_array in new_values:
            new_array[:] = np.linspace(lower, upper, len(new_array))

    return new_values


def _make_array_if_possible(x):
    """Helper function to cast (1, n) arrays to (n,) arrrays, or uniform lists
    of arrays to (p, n) arrays.
    """
    try:
        x = np.asfarray(x).squeeze()
    except ValueError:
        pass
    return x


def load(ts_path):
    """Load serialized TimeSeries from .npz file."""
    with np.load(ts_path) as npz_file:
        data = dict(npz_file)

    for key in ['time', 'measurement', 'error']:
        if key not in data:  # combine channel arrays into list
            n_channels = sum(1 for c in data.keys() if key in c)  # time0, ...
            data[key] = [data[key + str(i)] for i in range(n_channels)]

    # Convert 0d arrays to single values
    if 'name' in data:
        data['name'] = data['name'].item()
    if 'label' in data:
        data['label'] = data['label'].item()

    return TimeSeries(t=data.get('time'), m=data.get('measurement'),
                      e=data.get('error'),
                      meta_features=dict(zip(data['meta_feat_names'],
                                             data['meta_feat_values'])),
                      name=data.get('name'), label=data.get('label'))


class TimeSeries(object):
    """Class representing a single time series of measurements and metadata.

    A `TimeSeries` object encapsulates a single set of time-domain
    measurements, along with any metadata describing the observation.
    Typically the observations will consist of times, measurements, and
    (optionally) measurement errors. The measurements can be scalar- or
    vector-valued (i.e., "multichannel"); for multichannel measurements, the
    times and errors can also be vector-valued, or they can be shared across
    all channels of measurement.

    Attributes
    ----------
    time : (n,) or (p, n) array or list of (n,) arrays
        Array(s) of times corresponding to measurement values. If `measurement`
        is two-dimensional, this can be one-dimensional (same times for each
        channel) or two-dimensional (different times for each channel). If
        `time` is one-dimensional then it will be broadcast to match
        `measurement.shape`.
    measurement : (n,) or (p, n) array or list of (n,) arrays
        Array(s) of measurement values; can be two-dimensional for
        multichannel data. In the case of multichannel data with different
        numbers of measurements for each channel, `measurement` will be a list
        of arrays instead of a single two-dimensional array.
    error : (n,) or (p, n) array or list of (n,) arrays
        Array(s) of measurement errors for each value. If `measurement` is
        two-dimensional, this can be one-dimensional (same times for each
        channel) or two-dimensional (different times for each channel).
        If `error` is one-dimensional then it will be broadcast match
        `measurement.shape`.
    label : str, float, or None
        Class label or regression target for the given time series (if
        applicable).
    meta_features : dict
        Dictionary of feature names/values specified independently of the
        featurization process in `featurize`.
    name : str or None
        Identifying name for the given time series (if applicable).
        Typically the name of the raw data file from which the time series was
        created.
    path : str or None
        Path to the file where the time series is stored on disk (if
        applicable).
    channel_names : list of str
        List of names of channels of measurement; by default these are simply
        `channel_{i}`, but can be arbitrary depending on the nature of the
        different measurement channels.
    """
    def __init__(self, t=None, m=None, e=None, label=None, meta_features={},
                 name=None, path=None, channel_names=None):
        """Create a `TimeSeries` object from measurement values/metadata.

        See `TimeSeries` documentation for parameter values.
        """
        if t is None and m is None:
            raise ValueError("Either times or measurements must be provided.")
        elif m is None:
            m = _default_values_like(t, value=np.nan)

        # If m is 1-dimensional, so are t and e
        if _ndim(m) == 1:
            self.n_channels = 1
            if t is None:
                t = _default_values_like(m, upper=DEFAULT_MAX_TIME)
            if e is None:
                e = _default_values_like(m, value=DEFAULT_ERROR_VALUE)
        # If m is 2-dimensional, t and e could be 1d or 2d; default is 1d
        elif isinstance(m, np.ndarray) and m.ndim == 2:
            self.n_channels = len(m)
            if t is None:
                t = _default_values_like(m[0], upper=DEFAULT_MAX_TIME)
            if e is None:
                e = _default_values_like(m[0], value=DEFAULT_ERROR_VALUE)
        # If m is ragged (list of 1d arrays), t and e should also be ragged
        elif _ndim(m) == 2:
            self.n_channels = len(m)
            if t is None:
                t = _default_values_like(m, upper=DEFAULT_MAX_TIME)
            if e is None:
                e = _default_values_like(m, value=DEFAULT_ERROR_VALUE)
        else:
            raise ValueError("m must be a 1D or 2D array, or a 2D list of"
                             " arrays.")

        self.time = _make_array_if_possible(t)
        self.measurement = _make_array_if_possible(m)
        self.error = _make_array_if_possible(e)
        self.sort()  # re-order by time before broadcasting

        if _ndim(self.time) == 1 and _ndim(self.measurement) == 2:
            if isinstance(self.measurement, np.ndarray):
                self.time = np.broadcast_to(self.time, self.measurement.shape)
            else:
                raise ValueError("Times for each channel must be provided if m"
                                 " is a ragged array.")

        if _ndim(self.error) == 1 and _ndim(self.measurement) == 2:
            if isinstance(self.measurement, np.ndarray):
                self.error = np.broadcast_to(self.error, self.measurement.shape)
            else:
                raise ValueError("Errors for each channel must be provided if"
                                 " m is a ragged array.")

        if not (_compatible_shapes(self.measurement, self.time) and
                _compatible_shapes(self.measurement, self.error)):
            raise ValueError("times, values, errors are not of compatible"
                             " types/sizes. Please refer to the docstring"
                             " for list of allowed input types.")

        self.label = label
        self.meta_features = dict(meta_features)
        self.name = name
        self.path = path
        if channel_names is None:
            self.channel_names = ["channel_{}".format(i)
                                  for i in range(self.n_channels)]
        else:
            self.channel_names = channel_names

    def channels(self):
        """Iterates over measurement channels (whether one or multiple)."""
        t_channels = self.time
        m_channels = self.measurement
        e_channels = self.error
        if isinstance(self.time, np.ndarray) and self.time.ndim == 1:
            t_channels = np.broadcast_to(self.time,
                                         (self.n_channels, len(self.time)))
        if (isinstance(self.measurement, np.ndarray) and
                self.measurement.ndim == 1):
            m_channels = np.broadcast_to(self.measurement,
                                         (self.n_channels, len(self.measurement)))
        if isinstance(self.error, np.ndarray) and self.error.ndim == 1:
            e_channels = np.broadcast_to(self.error,
                                         (self.n_channels, len(self.error)))
        return zip(t_channels, m_channels, e_channels)

    def sort(self):
        """Sort times, measurements, and errors by time."""
        if _ndim(self.time) == 1:
            inds = np.argsort(self.time)
            self.time = self.time[inds]
            if _ndim(self.measurement) == 1:
                self.measurement = self.measurement[inds]
            else:
                for i in range(len(self.measurement)):
                   self.measurement[i] = self.measurement[i][inds]
            if _ndim(self.error) == 1:
                self.error = self.error[inds]
            else:
                for i in range(len(self.error)):
                   self.error[i] = self.error[i][inds]
        else:  # if time is 2d, so are measurement and error
            for i in range(len(self.time)):
                inds = np.argsort(self.time[i])
                self.time[i] = self.time[i][inds]
                self.measurement[i] = self.measurement[i][inds]
                self.error[i] = self.error[i][inds]


    def save(self, path=None):
        """Store TimeSeries object as a single .npz file.

        Attributes are stored in the following arrays:
            - time
            - measurement
            - error
            - meta_feat_names
            - meta_feat_values
            - name
            - label

        If `path` is omitted then the `path` attribute from the TimeSeries
        object is used.
        """
        if path is None:
            path = self.path

        data = {'meta_feat_names': list(self.meta_features.keys()),
                'meta_feat_values': list(self.meta_features.values())}

        for key in ['time', 'measurement', 'error']:
            value = getattr(self, key)
            if isinstance(value, np.ndarray):
                data[key] = value
            else:  # list of arrays -> save each channel separately
                for i, value_i in enumerate(value):
                    data[key + str(i)] = value_i

        if self.name:
            data['name'] = self.name
        if self.label:
            data['label'] = self.label
        np.savez(path, **data)
