"""
Load a nix data file from the data set of Henke et al. and make a basic
plot of its contents.

The nix files contain the spike times of a cell recorded in a recording session, stimulus
onset and durations, reproduction onset times and durations, reward information,
feedback range widths, and distance covered by the animal throughout the session.

Use with Python 3.
"""
__author__ = ("Kay Thurley")
__version__ = "1.0, November 2021"


import sys
import nixio
import numpy as np


###################################################### commandline parameters

file_name = sys.argv[1]				# second argument should be the name of the file to load

###################################################### load data

file = nixio.File.open(file_name, nixio.FileMode.ReadOnly)
session = file.blocks[0]
metadata = file.sections[0]
for key, m in metadata.items():
    print(key+':', metadata[key])


###################################################### plotting initialization

import matplotlib as mpl
import matplotlib.pyplot as plt

fontsize = 12.0
markersize = 6

mpl.rcParams['lines.linewidth'] = 2
mpl.rcParams['lines.markersize'] = markersize
mpl.rcParams['font.size'] = fontsize
mpl.rcParams['xtick.direction'] = 'out'
mpl.rcParams['ytick.direction'] = 'out'


###################################################### plotting

x_axis = session.data_arrays['spike times'].dimensions[0]
spikes = session.data_arrays['spike times'].data[:]

fig = plt.figure(figsize=(11, 3))
ax = fig.add_axes([0.075, 0.25, 0.85, 0.625])

# -- plot spikes
ax.eventplot(spikes, colors=[[0, 0, 0]], lineoffsets=1, linelengths=.25, linewidth=.5, alpha=.5)


# -- mark stimuli / measurement phase
stimuli = session.multi_tags['stimuli']
stimulus_onsets = stimuli.positions[:]
stimulus_durations = stimuli.extents[:]
for i, onset in enumerate(stimulus_onsets):
    ax.axvspan(onset, onset+stimulus_durations[i],
                          color='silver', alpha=0.5, zorder=0, label="measurement phase")
    ax.text(onset, 1.2, stimuli.features[0].data[i][0], color="slategray", fontsize=8)


# -- mark reproduction phase
reproductions = session.multi_tags['reproductions']
reproduction_onsets = reproductions.positions[:]
reproduction_durations = reproductions.extents[:]
rewards = session.data_arrays['rewards'].data[:]
for i, onset in enumerate(reproduction_onsets):
    # green color for reward, red for no reward
    if rewards[i]:
        col = 'green'
    else:
        col = 'red'
    ax.axvspan(onset, onset + reproduction_durations[i],
               color=col, alpha=0.5, zorder=0, label="reproduction phase")


# -- add speed information
ax2 = ax.twinx()
distances = session.data_arrays['distance covered']
x_axis = distances.dimensions[0]
x = x_axis.axis(distances.data.shape[0])
speed = np.diff(distances.data[:])/np.diff(x)
line_speed, = ax2.plot(x[:-1], speed, '.-', linewidth=1, alpha=0.5, color=np.ones(3)*.5,
                  markersize=1)


# -- annotate plot
ax.set_yticks([])
ax.set_xlabel(x_axis.label.capitalize() + " (" + x_axis.unit + ")")
ax2.set_ylabel('Speed (' + distances.unit + '/' + x_axis.unit + ')')


###################################################### finishing

file.close()
plt.show()