import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.colors as mcolors
import matplotlib.gridspec as gridspec
from IPython import embed
from scipy import stats
import math
import os
from IPython import embed
import helper_functions as hf
from params import *

if __name__ == '__main__':

    ###################################################################################################################
    # load data
    ###################################################################################################################
    # load all the data of one day
    filename = sorted(os.listdir('../../../data/'))[1]
    aifl = np.load('../data/'+filename+'/aifl2.npy', allow_pickle=True)

    # all_max_ch_means = np.load('../data/' + filename + '/all_max_ch_means.npy', allow_pickle=True)
    all_xticks = np.load('../data/' + filename + '/all_xtickses.npy', allow_pickle=True)
    all_Ctime_v = np.load('../data/' + filename + '/all_Ctime_v.npy', allow_pickle=True)
    ipp = np.load('../data/' + filename + '/all_ipp.npy', allow_pickle=True)
    power_means = np.load('../data/' + filename + '/power_means.npy', allow_pickle=True)

    ###################################################################################################################
    # parameter and variables
    ###################################################################################################################
    # plot params
    # variables
    sampling_rate = 1/np.diff(all_Ctime_v[0])[0]  # in sec

    # lists
    fish_in_aifl = list(np.unique(np.where(~np.isnan(aifl[:, :, 0]))[1]))
    thresholds = np.full(3, np.nan)
    all_run_std = []
    ###################################################################################################################
    # analysis
    ###################################################################################################################
    fig = plt.figure(figsize=(16, 14))
    spec = gridspec.GridSpec(ncols=2, nrows=2, figure=fig)
    ax1 = fig.add_subplot(spec[0, 0])
    ax2 = fig.add_subplot(spec[0, 1])
    ax3 = fig.add_subplot(spec[1, 0])

    for fish in range(len(fish_in_aifl)):

        run_mean, run_std = hf.running_3binsizes(ipp[fish], sampling_rate)
        all_run_std.append(run_std)

        if len(np.arange(len(run_mean[2]))[~np.isnan(run_mean[2])]) < 10:
            continue

        ###########################################################################################################
        # plot
        # hf.plot_running(all_xticks, run_mean, run_std, threshold, fish, color_vec, fs, fs_ticks, lw)
        # plt.show()

        for ax_cntr ,ax in enumerate([ax1, ax2, ax3]):
            non_nan = np.arange(len(run_mean[ax_cntr]))[~np.isnan(run_mean[ax_cntr])]
            ax.plot(all_xticks[fish][non_nan], run_std[ax_cntr][non_nan], '.', color=color_vec[fish])
            ax.timeaxis()
            ax.make_nice_ax()

        ax1.set_xlabel('Time [h]', fontsize=fs)

    rs = np.array(all_run_std)

    for idx in range(3):
        rs1 = np.hstack(rs[:, idx])
        rs1 = rs1[~np.isnan(rs1)]
        rs1 = rs1[rs1>0]
        print(np.percentile(rs1, [5, 95, 50]))
        thresholds[idx] = np.percentile(rs1, 50)

    if filename not in os.listdir('../data/'):
        os.mkdir('../data/'+filename)

    np.save('../data/' + filename + '/thresholds.npy', thresholds)

    fig2, ax2 = plt.subplots()

    ax2.hist(np.hstack(rs[:,0]), bins=200)
    ax2.hist(np.hstack(rs[:,1]), bins=200)
    ax2.hist(np.hstack(rs[:,2]), bins=200)
    plt.show()

    embed()