import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.gridspec as gridspec

from IPython import embed
import helper_functions as hf
from params import *
import os

if __name__ == '__main__':
    ###################################################################################################################
    # parameter and variables
    kernel_size = 100
    kernel = np.ones(kernel_size) / kernel_size

    fig1 = plt.figure(constrained_layout=True, figsize=[15 / inch, 15 / inch])
    gs = gridspec.GridSpec(ncols=2, nrows=3, figure=fig1, hspace=0.05, wspace=0.05,
                           width_ratios=[1, 1], height_ratios=[1, 1, 1], left=0.1, bottom=0.15, right=0.95, top=0.98)
    c = 0
    ax_counter = 0
    ###################################################################################################################
    # load all the data of one day

    for filename_idx in [0, 1, 2, 3]:
        filename = sorted(os.listdir('../data/'))[filename_idx]

        all_max_ch_means = np.load('../data/' + filename + '/all_max_ch.npy', allow_pickle=True)
        all_xticks = np.load('../data/' + filename + '/all_xtickses.npy', allow_pickle=True)
        all_ipp = np.load('../data/' + filename + '/all_ipp.npy', allow_pickle=True)
        power_means = np.load('../data/' + filename + '/power_means.npy', allow_pickle=True)
        freq = np.load('../data/' + filename + '/fish_freq_q10.npy', allow_pickle=True)
        names = np.load('../data/' + filename + '/fish_species.npy', allow_pickle=True)

        ###############################################################################################################
        # get fish
        print(filename)

        for fish_number in range(len(power_means)):
            if names[fish_number] == 'Eigenmannia' and power_means[fish_number] >= -90.0:
                ipp = all_ipp[fish_number]
                x_tickses = all_xticks[fish_number]
                max_ch_mean = all_max_ch_means[fish_number]

                # smoothing of max channel mean
                smooth_mcm = np.convolve(max_ch_mean, kernel, 'valid')

                try:
                    smooth_x = x_tickses[int(np.ceil(kernel_size / 2)):-int(np.floor(kernel_size / 2))]
                except:
                    embed()
                    quit()

                #####################################################################################################
                # plot traces
                ax1 = fig1.add_subplot(gs[c])
                ax1.imshow(ipp[::20].T[::-1], vmin=-100, vmax=-50, aspect='auto', interpolation='gaussian',
                           extent=[x_tickses[0], x_tickses[-1], -0.5, 15.5])

                ax1.plot(smooth_x[::20], smooth_mcm[::20], '.', color=color2[4])

                ax1.make_nice_ax()
                ax1.axhline(7.5, xmin=0, xmax=15, color='white', lw=4)
                ax1.set_yticks([0, 1, 2, 3, 4, 5, 6, 7, 7.5, 8, 9, 10, 11, 12, 13, 14, 15])
                ax1.set_yticklabels(['1', '', '3', '', '5', '', '7', '', 'g', '', '10', '', '12', '', '14', '', '16'],
                                    fontsize=9)

                print(ax_counter)
                ax1.text(-0.17, 1, chr(ord('A') + ax_counter), transform=ax1.transAxes, fontsize='large')
                ax_counter += 1
                ax1.invert_yaxis()
                ax1.xaxis_date()
                date_format = mdates.DateFormatter('%H:%M')
                ax1.xaxis.set_major_formatter(date_format)


                if c in [3, 4]:
                    ax1.set_xlabel('Time', fontsize=11)
                if c in [1, 2]:
                    ax1.set_xticks(ax1.get_xticks()[::2])

                if c in [0, 2, 4]:
                    ax1.set_ylabel('Electrode', fontsize=11)

                c += 1

    ###################################################################################################################
    # save plot
    fig1.savefig(save_path + 'eigen_trajectories.pdf')
    plt.show()