from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np

from oneibl.one import ONE
import alf.io as ioalf

from brainbox.processing import bincount2D


class Bunch(dict):
    """A subclass of dictionary with an additional dot syntax."""
    def __init__(self, *args, **kwargs):
        super(Bunch, self).__init__(*args, **kwargs)
        self.__dict__ = self

    def copy(self):
        """Return a new Bunch instance which is a copy of the current Bunch instance."""
        return Bunch(super(Bunch, self).copy())


def one_to_bunch(spikes_obj):
    """Convert a ONE spikes object into a Bunch of spike times and spike clusters."""
    return Bunch(spike_times=spikes_obj['times'], spike_clusters=spikes['clusters'])


def firing_rates(spike_times, spike_clusters, bin_size):
    """Return the time-dependent firing rate of a population of neurons.

    :param spike_times: the spike times of all neurons, in seconds
    :param spike_clusters: the cluster numbers of all spikes
    :param bin_size: the bin size, in seconds
    :return: a (n_clusters, n_samples) array with the firing rate of every cluster

    """
    rates, times, clusters = bincount2D(spike_times, spike_clusters, bin_size)
    return rates


def xcorr(x, y, maxlags=None):
    """Cross-correlation between two 1D signals of the same length."""
    ns = len(x)
    if len(y) != ns:
        raise ValueError("x and y should have the same length.")
    maxlags = maxlags or ns - 1
    return np.correlate(x, y, mode='full')[ns - 1 - maxlags:ns + maxlags]


T_BIN = 0.01  # seconds
CORR_LEN = 1  # seconds
CORR_BINS = int(CORR_LEN / T_BIN)  # bins

# get the data from flatiron and the current folder
one = ONE()
eid = one.search(subject='ZM_1150', date='2019-05-07', number=1)
D = one.load(eid[0], clobber=False, download_only=True)
session_path = Path(D.local_path[0]).parent

# load objects
spikes = ioalf.load_object(session_path, 'spikes')

# Get a Bunch instance.
b = one_to_bunch(spikes)

# Compute the firing rates.
rates = firing_rates(b.spike_times, b.spike_clusters, T_BIN)
# Note: I would rather just use spikes['times'] and spikes['clusters'] instead of going
# via a Bunch or DataFrame or similar...

# Compute the cross-correlation between the firing rate of two neurons.
c = xcorr(rates[0], rates[1], CORR_BINS)

# Plot it.
lags = np.linspace(-CORR_LEN, +CORR_LEN, len(c))
plt.plot(c)
plt.show()
