from pathlib import Path
import numpy as np
from brainbox.processing import bincount2D
import matplotlib.pyplot as plt
from oneibl.one import ONE
import alf.io
plt.ion()


def raster_complete(R, times, Clusters):
    '''
    Plot a rasterplot for the complete recording
    (might be slow, restrict R if so),
    ordered by insertion depth
    '''

    plt.imshow(R, aspect='auto', cmap='binary', vmax=T_BIN / 0.001 / 4,
               origin='lower', extent=np.r_[times[[0, -1]], Clusters[[0, -1]]])

    plt.xlabel('Time (s)')
    plt.ylabel('Cluster #; ordered by depth')
    plt.show()

    # plt.savefig('/home/mic/Rasters/%s.svg' %(trial_number))
    # plt.close('all')
    plt.tight_layout()


if __name__ == '__main__':

    # get data
    one = ONE()
    eid = one.search(lab='wittenlab', date='2019-08-04')
    D = one.load(eid[0], clobber=False, download_only=True)
    alf_path = Path(D.local_path[0]).parent
    spikes = alf.io.load_object(alf_path, 'spikes')

    # bin activity
    T_BIN = 0.01  # [sec]
    R, times, Clusters = bincount2D(spikes['times'], spikes['clusters'], T_BIN)

    # Order activity by anatomical depth of neurons
    d = dict(zip(spikes['clusters'], spikes['depths']))
    y = sorted([[i, d[i]] for i in d])
    isort = np.argsort([x[1] for x in y])
    R = R[isort, :]

    # Check the number of clusters x number of time bins
    print(R.shape, '#clusters x #timebins')

    # get a raster plot for the complete recording
    raster_complete(R, times, Clusters)
