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():
            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)

    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')[:])


    # dims = [[]] * len(spike_data)
    # for j in range(len(spike_data)):
    #     dims[j] = spike_data.dimensions[j]
    # print(spike_data, len(spike_data))
    # embed()
    # exit()
    f.close()
    return stim_pos, stim_extent, spike_data, voltage_trace


if __name__ == '__main__':
    data_dir = '/home/lisa/data'
    data_set = '2019-06-24-aa-invivo-1'
    stim_pos, stim_extent, spike_data, voltage_trace = analyze_sams(os.path.join(data_dir, data_set, data_set + '.nix'))
    for i in range(len(stim_pos)):
        plt.scatter(voltage_trace.dimensions[i].axis(1000), spike_data[i][:1000])
        plt.show()
        # spikes = stim_tag.retrieve_data(i, 'Spikes-1')[:]
        # print(spikes)
        # embed()
        # exit()