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
import datetime

if __name__ == '__main__':
    ###################################################################################################################
    # parameter and variables
    # plot params
    inch = 2.45
    save_path = '../../thesis/Figures/Results/'
    kernel_size = 100

    ###################################################################################################################
    # load all the data of one day
    # for filename_idx in [1, 4, 6]:
    for filename_idx in [1]:
        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)
        ###############################################################################################################
        # get fish
        # for fish_number in range(len(power_means)):
        for fish_number in [14]:
            if 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
                kernel = np.ones(kernel_size) / kernel_size
                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
                fig1, ax1 = plt.subplots(1, 1, figsize=(13 / inch, 8 / inch))
                fig1.subplots_adjust(left=0.12, bottom=0.15, right=0.99, top=0.99)

                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.set_title('freq: %.1f, power: %.1f' %(freq[:,2][fish_number], power_means[fish_number]), fontsize=fs + 2)
                # ax1.set_title('freq: %.1f, Nr: %.1f' %(freq[:,2][fish_number], fish_number), fontsize=fs + 2)
                ax1.set_xticks(smooth_x[::350])
                ax1.beautimechannelaxis()
                ax1.timeaxis()
                # fig1.autofmt_xdate()

                fig1.savefig(save_path + 'trajectory_'+str(fish_number)+'.pdf')
                # fig1.savefig('../../../goettingen2021_poster/pictures/trajectory_'+ str(fish_number)+'.pdf')
                print(fish_number, freq[fish_number,2])
                plt.show()
                embed()