"""
__author__ = Hagai Hargil
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from scipy import stats
from typing import Dict, List
import gspread
from oauth2client.service_account import ServiceAccountCredentials
import warnings
from collections import deque
import pandas as pd
import holoviews as hv

from plot_and_analyze_hyper_hypo import process_ep_output, plot_all_ca_traces
from readout_ep_output import walk_ep_dirs_old, extract_data_from_ep_dirs


def multipage(filename, figs=None, dpi=200):
    pp = PdfPages(filename)
    if figs is None:
        figs = [plt.figure(n) for n in plt.get_fignums()]
    for fig in figs:
        fig.savefig(pp, format='pdf')
    pp.close()


def add_column_headers(sheet, heads: List):
    """ Take a sheet and add the header row """

    sheet.insert_row(heads, 1)


def add_row_headers(sheet, heads: List):
    """ Take a sheet and add the first column """

    cell_list = sheet.range(f'A2:A{2 + len(heads) - 1}')
    for idx, cell in enumerate(cell_list):
        cell.value = heads[idx]

    sheet.update_cells(cell_list)


def write_dict_to_gsheet(data: Dict, sheet, row: str):
    """
    Write the data dict into a Google sheet
    :return: None
    """
    headers: List = sheet.row_values(1)
    if [] == headers:
        raise UserWarning("Missing header row")

    try:
        row_idx = sheet.find(row)
    except CellNotFound:
        raise UserWarning(f"Missing row key {row}.")

    for key, item in data.items():
        print(key, item)
        try:
            col_idx = sheet.find(key)
        except CellNotFound:
            warnings.warn(f"Key {key} was missing.")
        sheet.update_cell(row_idx.row, col_idx.col, item)


def write_all_data_to_gsheet(data: Dict, sheetname="Sheet1",
                             scope='https://spreadsheets.google.com/feeds',
                             json_filename=r'C:\Users\Hagai\Documents\GitHub\python-gsheet\client_secret.json',
                             spreadsheet_name="Hypo_Hyper_from_EP"):
    """
    Write the data dict into a Google sheet
    :return: None
    """
    creds = ServiceAccountCredentials.from_json_keyfile_name(json_filename, scope)
    client = gspread.authorize(creds)
    try:  # if the worksheet already exists
        sheet = client.open(spreadsheet_name).add_worksheet(sheetname, rows=100, cols=30)
    except:
        sheet = client.open(spreadsheet_name).worksheet(sheetname)

    col_heads = deque(list(data.values())[0].keys())
    col_heads.appendleft(0)
    row_heads = list(data.keys())
    add_column_headers(sheet=sheet, heads=col_heads)
    add_row_headers(sheet=sheet, heads=row_heads)

    for key, dict_data in data.items():
        write_dict_to_gsheet(data=dict_data, sheet=sheet, row=key)


def process_parsed_data_with_hypo_hyper(dict_of_data, list_of_wanted):
    """ Take the extracted data and find the relevant EP variables """

    dict_of_std_hypo = {key: np.array([]) for key in list_of_wanted}
    dict_of_averages_hypo = {key: np.array([]) for key in list_of_wanted}
    dict_of_std_hyper = {key: np.array([]) for key in list_of_wanted}
    dict_of_averages_hyper = {key: np.array([]) for key in list_of_wanted}

    for key in sorted(dict_of_data):
        for key2 in sorted(dict_of_averages_hypo):
            if key2 in key:
                average = np.mean(dict_of_data[key])
                cur_std = np.std(dict_of_data[key])
                if "HYPO" in key:
                    dict_of_averages_hypo[key2] = np.append(dict_of_averages_hypo[key2], average)
                    dict_of_std_hypo[key2] = np.append(dict_of_std_hypo[key2], cur_std)
                else:
                    dict_of_averages_hyper[key2] = np.append(dict_of_averages_hyper[key2], average)
                    dict_of_std_hyper[key2] = np.append(dict_of_std_hyper[key2], cur_std)

    return dict_of_std_hypo, dict_of_averages_hypo, dict_of_std_hyper, dict_of_averages_hyper


def plot_and_export_results(dict_of_std_hypo, dict_of_averages_hypo, dict_of_std_hyper, dict_of_averages_hyper):
    hypo_mean: Dict = {}
    hypo_std: Dict = {}
    hyper_mean: Dict = {}
    hyper_std: Dict = {}

    for key in sorted(dict_of_averages_hyper):
        plt.figure()
        last_ind = len(dict_of_averages_hyper[key]) if len(dict_of_averages_hyper[key]) < len(dict_of_averages_hypo[key]) \
            else len(dict_of_averages_hypo[key])
        _, pval = stats.ttest_ind(dict_of_averages_hypo[key], dict_of_averages_hyper[key],
                                  equal_var=False)
        plt.title(f"{key}, p-value (Welch): {pval:.3f}")
        ###
        np.save(arr=dict_of_averages_hyper, file='avg_hyper.npy')
        np.save(arr=dict_of_averages_hypo, file='avg_hypo.npy')
        ###

        hypo_std[key] = np.std(dict_of_averages_hypo[key])
        hypo_mean[key] = np.mean(dict_of_averages_hypo[key])
        hyper_std[key] = np.std(dict_of_averages_hyper[key])
        hyper_mean[key] = np.mean(dict_of_averages_hyper[key])
        plt.errorbar(np.arange(2), [hypo_mean[key], hyper_mean[key]],
                     yerr=[hypo_std[key], hyper_std[key]], fmt='o')
        labels = ['Hypo', 'Hyper']
        plt.xticks(np.arange(2), labels)

    # Write to a Google sheet - only means of each mouse
    # dict_of_data = {'hypo_mean': hypo_mean,
    #                 'hypo_std': hypo_std,
    #                 'hyper_mean': hyper_mean,
    #                 'hyper_std': hyper_std}
    # Write to a Google sheet - all data points with their STD
    dict_of_data = {'hypo_mean': dict_of_averages_hypo,
                    'hypo_std': dict_of_std_hypo,
                    'hyper_mean': dict_of_averages_hyper,
                    'hyper_std': dict_of_std_hyper}
    # write_all_data_to_gsheet(data=dict_of_data, sheetname="Rotated Mice2")


def process_reg_and_rotated_mice(dict_of_data, list_of_wanted):
    """ Used when we measured EP for the same mouse in two different orientations """

    cols = ['Measure', 'FOV', 'Hemisphere', 'Orientation', 'Data']
    df = pd.DataFrame([], columns=cols)
    keys = np.array(list(dict_of_data.keys()), dtype=str)

    for key in list_of_wanted:
        indices = np.core.defchararray.find(keys, key)
        rel_keys = keys[indices >= 0]
        for rel_key in rel_keys:
            # Find the values of the specific entry
            cur_fov = 1 if 'FOV_1' in rel_key else 2
            hemi = 'Hyper' if 'HYPER' in rel_key else 'Hypo'
            ori = 'Regular' if 'regular' in rel_key else 'Rotated'
            data = np.mean(dict_of_data[rel_key])
            helper = pd.DataFrame([[key, cur_fov, hemi, ori, data]], columns=cols)
            df = df.append(helper, ignore_index=True)

    df.set_index(['Measure', 'FOV', 'Hemisphere', 'Orientation'], drop=True, inplace=True)
    return df


def plot_rotated_data(df):
    for measure in df.index.levels[0]:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        plt.title(measure)
        all_hypo = df.loc[measure].xs('Hypo', level='Hemisphere')
        all_hyper = df.loc[measure].xs('Hyper', level='Hemisphere')

        data_to_plot = np.zeros((2, 4), dtype=np.float64)
        data_to_plot[:, 0] = np.squeeze(all_hypo.xs('Regular', level='Orientation').values)
        data_to_plot[:, 1] = np.squeeze(all_hypo.xs('Rotated', level='Orientation').values)
        data_to_plot[:, 2] = np.squeeze(all_hyper.xs('Regular', level='Orientation').values)
        data_to_plot[:, 3] = np.squeeze(all_hyper.xs('Rotated', level='Orientation').values)
        x_data = np.array([[np.arange(4)], [np.arange(4)]])
        ax.set_xticks((0, 1, 2, 3))
        ax.set_xticklabels(('Hypo-Regular', 'Hypo-Rotated', 'Hyper-Regular', 'Hyper-Rotated'))

        ax.scatter(x_data, data_to_plot)


if __name__ == '__main__':
    directory = r'X:\David\THY_1_GCaMP_BEFOREAFTER_TAC_290517\results'
    # list_of_wanted = ['S_or_run_stim', 'S_or_run_spont', 'S_or_run_juxta',
    #                   'S_or_stand_stim', 'S_or_stand_spont', 'S_or_stand_juxta',
    #                   'C_df_run_stim', 'C_df_run_spont', 'C_df_run_juxta',
    #                   'C_df_stand_stim', 'C_df_stand_spont', 'C_df_stand_juxta']
    list_of_wanted = ['C_df']
    # For normal TAC analysis
    # dict_of_std_hypo, dict_of_averages_hypo, \
    # dict_of_std_hyper, dict_of_averages_hyper = process_parsed_data_with_hypo_hyper(dict_of_data, list_of_wanted)
    # plot_and_export_results(dict_of_std_hypo, dict_of_averages_hypo, dict_of_std_hyper, dict_of_averages_hyper)

    # For regular-rotated analysis
    # dict_of_data = walk_ep_dirs_old(directory=directory, file_end=r"compiled_r.+\.mat$",
    #                                 list_of_wanted=list_of_wanted)
    # df = process_reg_and_rotated_mice(dict_of_data, np.array(list_of_wanted))
    # plot_rotated_data(df)

    # Alternative to normal TAC
    df = extract_data_from_ep_dirs(directory=directory, file_end=r"*compiled.mat",
                                   list_of_wanted=list_of_wanted)
    # df = plot_all_ca_traces(df)
    # Save figures to a single PDF
    # multipage('hypo_hyper_rotated_and_regular_same_mouse.pdf')
