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


def stimulus_collector(filename):
    try:
        stims = []
        f = nix.File.open(filename, nix.FileMode.ReadOnly)
        b = f.blocks[0]
        for g in b.groups:
            # embed()
            # exit()
            stims.append(g.multi_tags[0].name)
    except KeyError:
        embed()
        exit()
    return stims


def analyze_sams(filename):
    f = nix.File.open(filename, nix.FileMode.ReadOnly)
    b = f.blocks[0]
    b.metadata.pprint(max_depth=-1) # max_depth=-1: alles rausschreiben
    for g in b.groups:

        if 'sam' in g.name.lower(): # go through loop, until 'sam' is found
             break
    # embed()
    # exit()
    rtag_data = g.tags[0]
    rtag_data.metadata.pprint(max_depth=-1)

    print(40*'*')

    stim_tag = g.multi_tags[0]
    stim_type = g.multi_tags[0].name
    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 = voltage_trace.dimensions[0].axis(len(voltage_trace))
    plt.plot(dims[:-1], voltage_trace[:-1])
    for i in range(len(stim_pos)):
        stim_y = [np.max(voltage_trace)+10, np.max(voltage_trace)+10]
        plt.plot([stim_pos[i], stim_pos[i] + stim_extent[i]], stim_y, 'k')

    for j in range(len(spike_data)):
        for k in range(len(spike_data[j])):
            plt.plot(spike_data[j][k], np.max(voltage_trace)+5, 'o', color='k')
        plt.xlim([stim_pos[0]-1, stim_extent[-1]+1])
        plt.title(filename[-26:-13] + ': ' + stim_type)
        # plt.show()
        print('saving...')
        os.chdir('/home/lisa/Dropbox/Masterarbeit/')
        plt.savefig(filename[-26:-13] + ': ' + stim_type + '.png')
        plt.close()
        # f.close()
        # embed()
        # exit()

        return


if __name__ == '__main__':
    data_dir = '/home/lisa/data/'
    os.chdir(data_dir)
    data_set = glob.glob('2019-*')
    for i in range(len(data_set)):
        # data_set = '2019-08-22-aa-invivo-1'
        print(data_dir + data_set[i] + '/' + data_set[i])
        analyze_sams(data_dir + data_set[i] + '/' + data_set[i] + '.nix')
        # analyze_sams(data_dir + data_set + '/' + data_set + '.nix')