import os
import numpy as np
import matplotlib.pyplot as plt
import nixio as nix
from IPython import embed

def analyze_sams(filename):
    # print(filename)
    f = nix.File.open(filename, nix.FileMode.ReadOnly)
    b = f.blocks[0]
    b.metadata.pprint(max_depth=-1) # max_depth=-1: alles rausschreiben
    # print(b.groups)
    for g in b.groups:
        if 'sam' in g.name.lower(): # go through loop, until 'sam' is found
            break
    rtag_data = g.tags[0]
    rtag_data.metadata.pprint(max_depth=-1)
    # print(40*'*')
    stim_tag = g.multi_tags[0]
    stim_pos = stim_tag.positions[:] # beginnings of stimulations
    stim_extent = stim_tag.extents[:] # duration of stimulations

    for r in rtag_data.references:
        print(r.name, r.type)
    # embed()
    # exit()
    voltage_trace = rtag_data.references['V-1']

    spike_data = []

    for idx in range(len(stim_pos)):
        spike_data.append(stim_tag.retrieve_data(idx, 'Spikes-1')[:])

    # embed()
    # exit()
    dims = voltage_trace.dimensions[0].axis(len(voltage_trace))

    # # f.close()
    # embed()
    # exit()
    return stim_pos, stim_extent, spike_data, voltage_trace, dims


if __name__ == '__main__':
    data_dir = '/home/lisa/data'
    data_set = '2019-06-24-aa-invivo-1'
    stim_pos, stim_extent, spike_data, voltage_trace, dims = analyze_sams(os.path.join(data_dir, data_set,
                                                                                       data_set + '.nix'))
    print(len(spike_data))
    stims = []
    for i in range(len(stim_pos)):
        stim_y = [-38, -38]
        plt.plot([stim_pos[i], stim_pos[i] + stim_extent[i]], stim_y, 'k')
    plt.plot(dims[:-1], voltage_trace[:-1])
    spike_y = np.ones(len(spike_data))*80

    for j in range(len(spike_data)):
        print('length of current spike train: ', len(spike_data))
        for k in range(len(spike_data[j])):
            print('current spike time is: ', spike_data[j][k])
            plt.plot(spike_data[j][k], -40, 'o', color='k')
    plt.ylim([-100, -30])
    plt.show()