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 os
from IPython import embed
from params import *
import datetime
import itertools
from statisitic_functions import *


if __name__ == '__main__':
    ###################################################################################################################
    # parameter
    # save_path = '../../thesis/Figures/Results/'
    # inch = 2.45
    cat_v1 = [0, 0, 750, 0]
    cat_v2 = [750, 750, 1000, 1000]
    cat_v = [0, 580, 750, 1000]
    cat_n = ['Eigenmannia', 'Apteronotus', 'Apteronotus']
    pos = [1, 2, 3, 4]

    ###################################################################################################################
    # load data
    freq = np.load('../data/f.npy', allow_pickle=True)
    c = np.load('../data/all_changes.npy', allow_pickle=True)
    names = np.load('../data/n.npy', allow_pickle=True)
    stl = np.load('../data/stl.npy', allow_pickle=True)

    ###################################################################################################################
    # figure
    fig = plt.figure(figsize=[16 / inch, 8 / inch])
    spec = gridspec.GridSpec(ncols=3, nrows=1, figure=fig, hspace=0.3, wspace=0.2,
                             width_ratios=[4,1,2], left=0.08, bottom=0.15, right=0.95, top=0.90)
    ax = fig.add_subplot(spec[0, 0])
    ax1 = fig.add_subplot(spec[0, 2])
    ax2 = fig.add_subplot(spec[0, 1])

    ###############################################################################################################
    # analysis: separates the activity and frequencies by the species and sex
    activitys = []
    freqs = []
    for p in range(3):
        f = freq[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p])]
        a = c[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p])]

        # f = freq[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & (stl==False)]
        # a = c[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & (stl==False)]
        #
        # f_m = freq[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & stl]
        # a_m = c[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & stl]

        activitys.append(a)
        freqs.append(f)

        ###############################################################################################################
        # point plot of freq vs activity
        ax.plot(f, a, 'o', alpha=0.7, color=color_efm[p], label=labels[p])
        # ax.plot(f_m, a_m, 'o', alpha=0.7, color=color_efm[p], markeredgecolor='k')
        ax.set_xlim(400, 1000)
        ax.set_xticks([400, 600, 800, 1000])
        ax.set_xticklabels([400, 600, 800, 1000])
        ax.set_xlabel('EODf [Hz]', fontsize=fs)
        ax.set_ylabel('Activity [Changes/min]', fontsize=fs)
        ax.set_ylim(-0.2, 6)

    ###############################################################################################################
    # boxplot
    bp = ax1.boxplot(activitys, widths=0.6, showfliers=False)
    for element in ['boxes', 'medians']:
        for idx in range(3):
            bp[element][idx].set(linewidth=2, color=color_efm[idx])

    for element in ['whiskers', 'caps']:
        for idx, num in enumerate([0,2,4]):
            for i in [0,1]:
                bp[element][num+i].set(linewidth=2, color=color_efm[idx])
                if element == 'whiskers' and i == 1:
                    ax1.text(idx+1, bp['whiskers'][num+i].get_ydata()[1]+0.3, str(len(activitys[idx])),
                             horizontalalignment='center')

    ###############################################################################################################
    # histogram: activity distribution
    for j in [1,2,0]:
        n, bins = np.histogram(activitys[j], bins=np.arange(0,8.2,0.2))
        ax2.hist(activitys[j], bins=np.arange(0,8.2,0.25), orientation='horizontal', histtype='step',
                 color=color_efm[j], linewidth=2)

    ################################################################################################################
    # statisitics: cohens d' and mann whitney u

    for subset in itertools.combinations([0,1,2], 2):
        print(labels[subset[0]], labels[subset[1]])
        U, p = stats.mannwhitneyu(activitys[subset[0]], activitys[subset[1]])
        print('U = ', U, 'p = ', p*3)
        r = 1-(2*U)/(len(activitys[subset[0]])*len(activitys[subset[1]]))
        print('r = ',r)
        print('d = ',cohen_d(activitys[subset[0]], activitys[subset[1]]))
        significance_bar(ax1, p*3, None, subset[0]+1, subset[1]+1, 3.8)

    #################################################################################################################
    # nice axis
    tagx = [-0.15, -0.4, -0.17]
    tagy = [1.05,1.05,1.05]
    for ax_idx, axis in enumerate([ax,ax2,ax1]):
        axis.text(tagx[ax_idx], tagy[ax_idx], chr(ord('A') + ax_idx), transform=axis.transAxes, fontsize='large')
        axis.make_nice_ax()

    ax1.set_xlim(0.5,3.5)
    ax1.set_xticks([1, 2, 3])
    ax1.set_xticklabels(['$\it{E}$', '$\it{A}\u2640$', '$\it{A}\u2642$'])
    ax1.set_xlabel('Species', fontsize=fs)

    ax1.set_ylim(ax.get_ylim())
    ax1.spines["left"].set_visible(False)
    ax1.get_yaxis().set_visible(False)

    ax2.set_ylim(ax.get_ylim())
    ax2.set_yticklabels('')
    ax2.set_xlabel('n', fontsize=fs)

    fig.align_xlabels([ax,ax1,ax2])
    fig.legend(loc='upper right', bbox_to_anchor=(0.9, 0.85, 0.1, 0.1))

    # fig.savefig(save_path_pres+'activ_freq.pdf')
    # fig.savefig(save_path+'activ_freq.pdf')
    plt.show()
    embed()