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):
    stims = []
    f = nix.File.open(filename, nix.FileMode.ReadOnly)
    b = f.blocks[0]
    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
    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 = tags within this group
    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 stimuli
    stim_extent = stim_tag.extents[:]   # duration of stimuli

    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))
    os.chdir('/home/lisa/PycharmProjects/torus_ephys')
    print(os.getcwd())
    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)+2.5, 'o', color='k')
        plt.xlim([stim_pos[0]-.5, stim_pos[0] + stim_extent[-1] + .5])
        plt.xlabel('time [s]')
        plt.ylabel('voltage [mV]')
        print('saving...')
        for s_i in range(len(b.groups)):
            plt.savefig(filename + stim_type + str(s_i) + '.png')
        # f.close()
        # embed()
        # exit()

        return


def analyze_fis(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 'FICurve' in g.name:     # go through loop, until 'sam' is found
        #     break

        rtag_data = g.tags[0]   # rtag_data = tags within this group
        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 stimuli
        stim_extent = stim_tag.extents[:]   # duration of stimuli

        voltage_trace = rtag_data.references['V-1']
        for i in range(len(voltage_trace)):
            print(voltage_trace[i])
        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))

        for i in range(len(stim_pos)):
            plt.plot(dims[:-1], voltage_trace[:-1], color='#3673A4')
            stim_y = [np.max(voltage_trace)+1.25, np.max(voltage_trace)+1.25]
            plt.plot([stim_pos[i], stim_pos[i] + stim_extent[i]], stim_y, 'k')
            xmin = stim_pos[i] - .1
            xmax = stim_pos[i] + stim_extent[i] + .1
            y = np.ones(len(spike_data[i]))*(np.max(voltage_trace)+2.5)
            plt.plot(spike_data[i], y, 'o', color='#DC143C')
            plt.xlim(xmin, xmax)
            plt.xlabel('time [s]')
            plt.ylabel('voltage [mV]')
            plt.show()

        return

if __name__ == '__main__':
    data_dir = '/home/lisa/data/'
    os.chdir(data_dir)
    data_set = glob.glob('2019*')    # only look at single cells
    # data_set = '2019*'    # only look at single cells
    for i in range(len(data_set)):
        if '08' in data_set[i]:
            continue
    # print(data_dir + data_set + '/' + data_set)
        analyze_sams(data_dir + data_set[i] + '/' + data_set[i] + '.nix')