# -*- coding: utf-8 -*-

from collections.abc import Iterable

import numpy as np

from stingray import Lightcurve
from stingray.events import EventList
import stingray.utils as utils

__all__ = ['Covariancespectrum', 'AveragedCovariancespectrum']


class Covariancespectrum(object):
    """
          Compute a covariance spectrum for the data. The input data can be
          either in event data or pre-made light curves. Event data can either
          be in the form of a ``numpy.ndarray`` with ``(time stamp, energy)`` pairs or
          a :class:`stingray.events.EventList` object. If light curves are formed ahead
          of time, then a list of :class:`stingray.Lightcurve` objects should be passed to the
          object, ideally one light curve for each band of interest.

          For the case where the data is input as a list of :class:`stingray.Lightcurve` objects,
          the reference band(s) should either be

          1. a single :class:`stingray.Lightcurve` object,
          2. a list of :class:`stingray.Lightcurve` objects with the reference band for each band
             of interest pre-made, or
          3. ``None``, in which case reference bands will
             formed by combining all light curves *except* for the band of interest.

          In the case of event data, ``band_interest`` and ``ref_band_interest`` can
          be (multiple) pairs of energies, and the light curves for the bands of
          interest and reference bands will be produced dynamically.


          Parameters
          ----------
          data : {``numpy.ndarray`` | :class:`stingray.events.EventList` object | list of :class:`stingray.Lightcurve` objects}
              ``data`` contains the time series data, either in the form of a
              2-D array of ``(time stamp, energy)`` pairs for event data, or as a
              list of light curves.
              Note : The event list must be in sorted order with respect to the
              times of arrivals.

          dt : float
              The time resolution of the :class:`stingray.Lightcurve` formed from the energy bin.
              Only used if ``data`` is an event list.

          band_interest : {``None``, iterable of tuples}
              If ``None``, all possible energy values will be assumed to be of
              interest, and a covariance spectrum in the highest resolution
              will be produced.
              Note: if the input is a list of :class:`stingray.Lightcurve` objects, then the user may
              supply their energy values here, for construction of a
              reference band.

          ref_band_interest : {``None``, tuple, :class:`stingray.Lightcurve`, list of :class:`stingray.Lightcurve` objects}
              Defines the reference band to be used for comparison with the
              bands of interest. If ``None``, all bands *except* the band of
              interest will be used for each band of interest, respectively.
              Alternatively, a tuple can be given for event list data, which will
              extract the reference band (always excluding the band of interest),
              or one may put in a single :class:`stingray.Lightcurve` object to be used (the same
              for each band of interest) or a list of :class:`stingray.Lightcurve` objects, one for
              each band of interest.

          std : float or np.array or list of numbers
              The term ``std`` is used to calculate the excess variance of a band.
              If ``std`` is set to ``None``, default Poisson case is taken and the
              std is calculated as ``mean(lc)**0.5``. In the case of a single
              float as input, the same is used as the standard deviation which
              is also used as the std. And if the std is an iterable of
              numbers, their mean is used for the same purpose.

          Attributes
          ----------
          unnorm_covar : np.ndarray
              An array of arrays with mid point ``band_interest`` and their
              covariance. It is the array-form of the dictionary ``energy_covar``.
              The covariance values are unnormalized.

          covar : np.ndarray
              Normalized covariance spectrum.

          covar_error : np.ndarray
              Errors of the normalized covariance spectrum.

          References
          ----------
          [1] Wilkinson, T. and Uttley, P. (2009), Accretion disc variability\
              in the hard state of black hole X-ray binaries. Monthly Notices\
              of the Royal Astronomical Society, 397: 666–676.\
              doi: 10.1111/j.1365-2966.2009.15008.x

          Examples
          --------
          See the `notebooks repository <https://github.com/StingraySoftware/notebooks>`_ for
          detailed notebooks on the code.

    """

    def __init__(self, data, dt=None, band_interest=None,
                 ref_band_interest=None, std=None):

        self.dt = dt
        self.std = std

        # check whether data is an EventList object:
        if isinstance(data, EventList):
            data = np.vstack([data.time, data.energy]).T

        # check whether the data contains a list of Lightcurve objects
        if isinstance(data[0], Lightcurve):
            self.use_lc = True
            self.lcs = data
        else:
            self.use_lc = False

        # if band_interest is None, extract the energy bins and make an array
        # with the lower and upper bounds of the energy bins
        if band_interest is None:
            if not self.use_lc:
                self._create_band_interest(data)
            else:
                self.band_interest = np.vstack([np.arange(len(data)),
                                                np.arange(1, len(data)+1, 1)]).T
        else:
            if np.size(band_interest) < 2:
                raise ValueError('band_interest must contain at least 2 values '
                                 '(minimum and maximum values for each band) '
                                 'and be a 2D array!')

            self.band_interest = np.atleast_2d(band_interest)

        if self.use_lc is False and not dt:
            raise ValueError("If the input data is event data, the dt keyword "
                             "must be set and supply a time resolution for "
                             "creating light curves!")

        # if we don't have light curves already, make them:
        if not self.use_lc:
            if not np.all(np.diff(data, axis=0).T[0] >= 0):
                utils.simon("The event list must be sorted with respect to "
                            "times of arrivals.")
                data = data[data[:, 0].argsort()]

            self.lcs = self._make_lightcurves(data)

        # check whether band of interest contains a Lightcurve object:
        if np.size(ref_band_interest) == 1  or isinstance(ref_band_interest,
                                                          Lightcurve):
            if isinstance(ref_band_interest, Lightcurve):
                self.ref_band_lcs = ref_band_interest
            # ref_band_interest must either be a Lightcurve, or must have
            # multiple entries

            elif ref_band_interest is None:
                if self.use_lc:
                    self.ref_band_lcs = \
                        self._make_reference_bands_from_lightcurves(ref_band_interest)
                else:
                    self.ref_band_lcs = \
                        self._make_reference_bands_from_event_data(data)
            else:
                raise ValueError("ref_band_interest must contain either "
                                 "a Lightcurve object, a list of Lightcurve "
                                 "objects or a tuple of length 2.")
        else:
            # check whether ref_band_interest is a list of light curves
            if isinstance(ref_band_interest[0], Lightcurve):
                self.ref_band_lcs = ref_band_interest
                assert len(ref_band_interest) == len(self.lcs), "The list of " \
                                                                "reference light " \
                                                                "curves must have " \
                                                                "the same length as " \
                                                                "the list of light curves" \
                                                                "of interest."
            # if not, it must be a tuple, so we're going to make a list of light
            # curves
            else:
                if self.use_lc:
                    self.ref_band_lcs = \
                        self._make_reference_bands_from_lightcurves(bounds=
                                                                    ref_band_interest)
                else:
                    self.ref_band_lcs = \
                        self._make_reference_bands_from_event_data(data)

        self._construct_covar()

    def _make_reference_bands_from_event_data(self, data, bounds=None):
        """
        Helper method constructing reference bands for each band of interest, and constructing
        light curves from these reference bands. This operates only if the data given to
        :class:`Covariancespectrum` is event list data (i.e. photon arrival times and energies).

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape ``(N, 2)``, where N is the number of photons. First column contains the
            times of arrivals, second column the corresponding photon energies.

        bounds : iterable
            The energy bounds to use for the reference band. Must be of type ``(elow, ehigh)``.

        Returns
        -------

        lc_all: list of :class:`stingray.Lightcurve` objects.
            The list of `:class:`stingray.Lightcurve` objects containing all reference
            bands, between the values given in ``bounds``.

        """

        if not bounds:
            bounds = [np.min(data[:, 1]), np.max(data[:, 1])]

        if bounds[1] <= np.min(self.band_interest[:, 0]) or \
           bounds[0] >= np.max(self.band_interest[:, 1]):
            elow = bounds[0]
            ehigh = bounds[1]

            toa = data[np.logical_and(
                data[:, 1] >= elow,
                data[:, 1] <= ehigh)]

            lc_all = Lightcurve.make_lightcurve(toa, self.dt,
                                                tstart=self.tstart,
                                                tseg=self.tseg)

        else:

            lc_all = []
            for i, b in enumerate(self.band_interest):
                elow = b[0]
                ehigh = b[1]

                emask1 = data[np.logical_and(
                    data[:, 1] <= elow,
                    data[:, 1] >= bounds[0])]

                emask2 = data[np.logical_and(
                    data[:, 1] <= bounds[1],
                    data[:, 1] >= ehigh)]

                toa = np.vstack([emask1, emask2])
                lc = Lightcurve.make_lightcurve(toa, self.dt,
                                                tstart=self.tstart,
                                                tseg=self.tseg)
                lc_all.append(lc)

        return lc_all

    def _make_reference_bands_from_lightcurves(self, bounds=None):
        '''
        Helper class to construct reference bands for all light curves in ``band_interest``, assuming the
        data is given to the class :class:`Covariancespectrum` as a (set of) lightcurve(s). Generally
        sums up all other light curves within ``bounds`` that are *not* the band of interest.

        Parameters
        ----------
        bounds : iterable
            The energy bounds to use for the reference band. Must be of type ``(elow, ehigh)``.

        Returns
        -------
        lc_all: list of :class:`stingray.Lightcurve` objects.
            The list of :class:`stingray.Lightcurve` objects containing all reference bands,
            between the values given in ``bounds``.

        '''

        if not bounds:
            bounds_idx = [0, len(self.band_interest)]

        else:
            low_bound = self.band_interest.searchsorted(bounds[0])
            high_bound = self.band_interest.searchsorted(bounds[1])

            bounds_idx = [low_bound, high_bound]

        lc_all = []
        for i, b in enumerate(self.band_interest):

            # initialize empty counts array
            counts = np.zeros_like(self.lcs[0].counts)
            for j in range(bounds_idx[0], bounds_idx[1], 1):
                if i == j:
                    continue
                else:
                    counts += self.lcs[j].counts

            # make a combined light curve
            lc = Lightcurve(self.lcs[0].time, counts, skip_checks=True)

            # add to list of reference light curves
            lc_all.append(lc)

        return lc_all

    def _construct_covar(self):
        """
        Helper method to construct the covariance attribute and fill it with values.
        """

        self.avg_covar = False
        covar = np.zeros(len(self.lcs))
        covar_err = np.zeros(len(self.lcs))
        xs_var = np.zeros(len(self.lcs))

        for i in range(len(self.lcs)):
            lc = self.lcs[i]

            if np.size(self.ref_band_lcs) == 1 or isinstance(self.ref_band_lcs,
                                                             Lightcurve):
                lc_ref = self.ref_band_lcs
            else:
                lc_ref = self.ref_band_lcs[i]

            cv = self._compute_covariance(lc, lc_ref)
            cv_err = self._calculate_covariance_error(lc, lc_ref)

            covar[i] = cv
            covar_err[i] = cv_err

            xs = self._calculate_excess_variance(lc_ref)
            if not xs > 0:
                utils.simon("The excess variance in the reference band is "
                            "negative. This implies that the reference "
                            "band was badly chosen. Beware that the "
                            "covariance spectra will have NaNs!")

            xs_var[i] = xs

        self.unnorm_covar = covar
        energy_covar = covar / xs_var**0.5

        self.covar = energy_covar

        self.covar_error = covar_err

        return

    def _make_lightcurves(self, data):
        """
        Create light curves for all bands of interest from ``data``. Takes the information the
        ``band_interest`` attribute and event data in ``data``, and produces a list of
        :class:`stingray.Lightcurve` objects.

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape ``(N, 2)``, where ``N`` is the number of photons. First column contains the
            times of arrivals, second column the corresponding photon energies.

        Returns
        -------
        lc_all : iterable of :class:`stingray.Lightcurve` objects
            A list of :class:`stingray.Lightcurve` objects of all bands of interest.
        """

        self.tstart = np.min(data[:, 0])
        self.tend = np.max(data[:, 0])

        self.tseg = self.tend - self.tstart

        lc_all = []

        for i, b in enumerate(self.band_interest):
            elow = b[0]
            ehigh = b[1]

            toa = data[np.logical_and(
                data[:, 1] >= elow,
                data[:, 1] <= ehigh)]

            lc = Lightcurve.make_lightcurve(toa, self.dt, tstart=self.tstart,
                                            tseg=self.tseg)
            lc_all.append(lc)

        return lc_all

    def _create_band_interest(self, data):
        """
        If no bands of interest are given, but event data is, create bands of interest for each
        discrete enery value in the second column of ``data``.

        Parameters
        ----------
        data : numpy.ndarray
            Array of shape (N, 2), where N is the number of photons. First column contains the
            times of arrivals, second column the corresponding photon energies.

        """

        unique_energy = np.unique(data[:, 1])
        energ_diff = np.diff(unique_energy)

        energy_low = np.zeros_like(unique_energy)
        energy_high = np.zeros_like(unique_energy)

        energy_low[:-1] = unique_energy[:-1] - 0.5 * energ_diff
        energy_high[:-1] = unique_energy[:-1] + 0.5 * energ_diff

        energy_low[-1] = unique_energy[-1] - 0.5 * energ_diff[-1]
        energy_high[-1] = unique_energy[-1] + 0.5 * energ_diff[-1]

        energy_list = np.vstack([energy_low, energy_high]).T

        self.band_interest = energy_list

    def _calculate_excess_variance(self, lc):
        """Calculate excess variance in a band with the standard deviation."""
        std = self._calculate_std(lc)
        return np.var(lc) - std**2

    def _calculate_std(self, lc):
        """Return std calculated for the possible types of `std`"""
        if self.std is None:
            std = np.mean(lc)**0.5
        elif isinstance(self.std, Iterable):
            std = np.mean(self.std)  # Iterable of numbers
        else:  # Single float number
            std = self.std

        return std

    def _compute_covariance(self, lc1, lc2):
        """Calculate and return the covariance between two time series."""
        return np.cov(lc1.counts, lc2.counts)[0][1]

    def _calculate_covariance_error(self, lc_x, lc_y):
        """Calculate the error of the normalized covariance spectrum."""
        # Excess Variance of reference band
        xs_x = self._calculate_excess_variance(lc_x)
        # Standard deviation of light curve
        err_y = self._calculate_std(lc_y)
        # Excess Variance of reference band
        xs_y = self._calculate_excess_variance(lc_y)
        # Standard deviation of light curve
        err_x = self._calculate_std(lc_x)
        # Number of time bins in lightcurve
        nn = lc_x.n
        # Number of segments averaged
        if not self.avg_covar:
            mm = 1
        else:
            mm = self.nbins

        num = xs_x*err_y + xs_y*err_x + err_x*err_y
        denom = nn * mm * xs_y

        return (num / denom)**0.5


class AveragedCovariancespectrum(Covariancespectrum):
    """
    Compute a covariance spectrum for the data, defined in [covar spectrum]_ Equation 15.

    Parameters
    ----------
    data : {numpy.ndarray | list of :class:`stingray.Lightcurve` objects}
        ``data`` contains the time series data, either in the form of a
        2-D array of ``(time stamp, energy)`` pairs for event data, or as a
        list of :class:`stingray.Lightcurve` objects.
        Note : The event list must be in sorted order with respect to the
        times of arrivals.

    segment_size : float
        The length of each segment in the averaged covariance spectrum.
        The number of segments will be calculated automatically using the
        total length of the data set and the segment_size defined here.

    dt : float
        The time resolution of the :class:`stingray.Lightcurve` formed
        from the energy bin. Only used if `data` is an event list.

    band_interest : {``None``, iterable of tuples}
        If ``None``, all possible energy values will be assumed to be of
        interest, and a covariance spectrum in the highest resolution
        will be produced.
        Note: if the input is a list of :class:`stingray.Lightcurve` objects,
        then the user may supply their energy values here, for construction of a
        reference band.

    ref_band_interest : {None, tuple, :class:`stingray.Lightcurve`, list of :class:`stingray.Lightcurve` objects}
        Defines the reference band to be used for comparison with the
        bands of interest. If None, all bands *except* the band of
        interest will be used for each band of interest, respectively.
        Alternatively, a tuple can be given for event list data, which will
        extract the reference band (always excluding the band of interest),
        or one may put in a single :class:`stingray.Lightcurve` object to be used (the same
        for each band of interest) or a list of :class:`stingray.Lightcurve` objects, one for
        each band of interest.

    std : float or np.array or list of numbers
        The term ``std`` is used to calculate the excess variance of a band.
        If ``std`` is set to ``None``, default Poisson case is taken and the
        ``std`` is calculated as ``mean(lc)**0.5``. In the case of a single
        float as input, the same is used as the standard deviation which
        is also used as the std. And if the std is an iterable of
        numbers, their mean is used for the same purpose.

    Attributes
    ----------
    unnorm_covar : np.ndarray
        An array of arrays with mid point band_interest and their
        covariance. It is the array-form of the dictionary ``energy_covar``.
        The covariance values are unnormalized.

    covar : np.ndarray
        Normalized covariance spectrum.

    covar_error : np.ndarray
        Errors of the normalized covariance spectrum.

    References
    ----------
    .. [covar spectrum] http://arxiv.org/pdf/1405.6575v2.pdf
    """

    def __init__(self, data, segment_size, dt=None, band_interest=None,
                 ref_band_interest=None, std=None):

        self.segment_size = segment_size

        Covariancespectrum.__init__(self, data, dt=dt,
                                    band_interest=band_interest,
                                    ref_band_interest=ref_band_interest,
                                    std=std)

    def _construct_covar(self):
        """
        Helper method to construct the covariance attribute and fill it with values.
        """
        self.avg_covar = True

        start_time = self.lcs[0].time[0]

        covar = np.zeros(len(self.lcs))
        covar_err = np.zeros(len(self.lcs))
        xs_var = np.zeros(len(self.lcs))

        for i in range(len(self.lcs)):
            lc = self.lcs[i]

            if np.size(self.ref_band_lcs) == 1:
                lc_ref = self.ref_band_lcs
            else:
                lc_ref = self.ref_band_lcs[i]

            tstart = start_time
            tend = start_time + self.segment_size
            cv = 0.0
            cv_err = 0.0
            xs = 0.0

            self.nbins = int((tend - tstart)/self.segment_size)
            for k in range(self.nbins):
                start_ind = lc.time.searchsorted(tstart)
                end_ind = lc.time.searchsorted(tend)

                lc_seg = lc.truncate(start=start_ind, stop=end_ind)
                lc_ref_seg = lc_ref.truncate(start=start_ind, stop=end_ind)

                cv += self._compute_covariance(lc_seg, lc_ref_seg)
                cv_err += self._calculate_covariance_error(lc_seg, lc_ref_seg)
                xs += self._calculate_excess_variance(lc_ref_seg)
                if not xs > 0:
                    utils.simon("The excess variance in the reference band is "
                                "negative. This implies that the reference "
                                "band was badly chosen. Beware that the "
                                "covariance spectra will have NaNs!")

                tstart += self.segment_size
                tend += self.segment_size


            covar[i] = cv/self.nbins
            covar_err[i] = cv_err/self.nbins
            xs_var[i] = xs/self.nbins

        self.unnorm_covar = covar
        energy_covar = covar / xs_var**0.5

        self.covar = energy_covar

        self.covar_error = covar_err

        return
