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 mpl_toolkits.axes_grid1.inset_locator import inset_axes
from IPython import embed
from scipy import stats, optimize
import pandas as pd
import math
import os
from IPython import embed

from eventdetection import threshold_crossings, merge_events
import helper_functions as hf
from params import *
from statisitic_functions import significance_bar, cohen_d
import itertools


def get_recording_number_in_time_bins(time_bins):
    """
    Calculates the number of the recordings in the time bins

    :param time_bins: numpy array with borders of the time bins
    :return: time_bins_recording: numpy array with the number of recordings to that specific time bin
    """
    # variables
    time_bins_recordings = np.zeros(len(time_bins) - 1)

    # load data
    for index, filename_idx in enumerate([0, 1, 2, 3]):
        filename = sorted(os.listdir('../data/'))[filename_idx]
        time_points = np.load('../data/' + filename + '/all_hms.npy', allow_pickle=True)

        # in which bins is this recording, fill time_bins_recordings
        unique_time_points = np.unique(np.hstack(time_points))
        for idx, tb in enumerate(time_bins[:-1]):
            if np.any((unique_time_points >= tb) & (unique_time_points <= time_bins[idx + 1])):
                time_bins_recordings[idx] += 1

    return time_bins_recordings


def func(x, a, tau, c):
    return a * np.exp(-x / tau) + c


def calc_movement(cbf, i):
    movement = cbf[0, :, i] + cbf[1, :, i]
    movement[np.isnan(movement)] = 0
    re_mov = cbf[0, :, i] - cbf[1, :, i]
    re_mov[np.isnan(re_mov)] = 0

    return movement, re_mov


if __name__ == '__main__':
    ###################################################################################################################
    # parameter and variables
    # plot params
    inch = 2.45

    c = 0
    cat_v1 = [0, 0, 750, 0]
    cat_v2 = [750, 750, 1000, 1000]
    cat_n = ['Eigenmannia', 'Apteronotus', 'Apteronotus']

    ###################################################################################################################
    # load data
    ###################################################################################################################
    # load all the data of one day
    cbf2 = np.load('../data/cbf15.npy', allow_pickle=True)
    stl = np.load('../data/stl.npy', allow_pickle=True)
    freq = np.load('../data/f.npy', allow_pickle=True)
    names = np.load('../data/n.npy', allow_pickle=True)
    speed_average = np.load('../data/speed.npy', allow_pickle=True)

    ###################################################################################################################
    # pie chart
    ###################################################################################################################

    label = ['Transit', 'Stationary',
             '$\it{Eigenmannia\,sp.}$', '$\it{A.\,macrostomus}\,\u2640$', '$\it{A.\,macrostomus}\,\u2642$']
    size = 0.3

    true = freq[(names != 'unknown') & (stl)]
    false = freq[(names != 'unknown') & (stl == False)]

    true_per = len(true)
    false_per = len(false)

    fracs = [true_per, false_per]

    fracs2 = []
    fracs3 = []

    for p in range(3):
        true = freq[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & (stl)]
        false = freq[(names == cat_n[p]) & (freq >= cat_v1[p]) & (freq < cat_v2[p]) & (stl == False)]

        fracs2.append(len(true))
        fracs3.append(len(false))

    ###################################################################################################################
    # figure 2
    # fig2, ax2 = plt.subplots(figsize=[9 / inch, 7.8 / inch], subplot_kw=dict(aspect="equal"))
    #
    # vals = np.array([fracs2, fracs3])
    #
    # ax2.pie(vals.sum(axis=0), colors=['#2e4053', '#ab1717', '#004d8d'], labels=label[2:],
    #         startangle=90,  radius=0.7, wedgeprops=dict(width=size, edgecolor='w', linewidth=3),
    #         labeldistance=1.6, textprops=dict(ha='center'))
    #
    # wedges = ax2.pie(vals.flatten('F'), colors=['#425364', '#818c97', '#b32e2e', '#cc7373', '#195e98', '#6694ba'],
    #                  startangle=90, radius=0.7+size, wedgeprops=dict(width=size, edgecolor='w', alpha=0.82))
    #
    # ax2.pie(vals.sum(axis=0), colors=['#2e4053', '#ab1717', '#004d8d'],
    #         startangle=90, radius=0.7+size, wedgeprops=dict(width=size+size, edgecolor='w', linewidth=3, fill=False))
    #
    # ax2.legend(wedges[0][:2], label[:2],
    #           loc="center left",
    #           bbox_to_anchor=(0.64, -0.46, 0.5, 1))
    #
    # centre_circle = plt.Circle((0, 0), 0.2, color='black', fc='white', linewidth=0)
    # fig2.gca().add_artist(centre_circle)
    #
    # plt.axis('equal')
    # plt.tight_layout()
    # fig2.savefig(save_path + 'pie.pdf')
    # # fig2.savefig('../../../goettingen2021_poster/pictures/pie.pdf')
    #
    # plt.show()
    #
    # embed()
    # quit()
    ###################################################################################################################
    # figure 2
    fig2, ax2 = plt.subplots(figsize=[9 / inch, 7.8 / inch], subplot_kw=dict(aspect="equal"))

    vals = np.array([fracs2, fracs3])

    ax2.pie(vals.sum(axis=1), colors=['#1b2631', '#5d6d7e'], labels=label[:2],
            startangle=180, radius=0.7, wedgeprops=dict(width=size, edgecolor='w', linewidth=3),
            labeldistance=1.9, textprops=dict(ha='center'))

    wedges = ax2.pie(vals.flatten(), colors=['#3B4A5A', '#b32e2e', '#195e98',
                                             '#818c97', '#cc7373', '#6694ba'],
                     startangle=180, radius=0.7 + size, wedgeprops=dict(width=size, edgecolor='w', alpha=0.82))

    ax2.pie(vals.sum(axis=1), colors=color_diffdays[:2],
            startangle=180, radius=0.7 + size,
            wedgeprops=dict(width=size + size, edgecolor='w', linewidth=3, fill=False))

    ax2.legend(wedges[0][:3], label[2:],
               loc="center left",
               bbox_to_anchor=(0.5, -0.35, 0.5, 0.5))

    centre_circle = plt.Circle((0, 0), 0.2, color='black', fc='white', linewidth=0)
    fig2.gca().add_artist(centre_circle)

    plt.axis('equal')
    plt.tight_layout()
    fig2.savefig(save_path + 'pie.pdf')
    # fig2.savefig('../../../goettingen2021_poster/pictures/pie.pdf')

    plt.show()
    embed()