# Runs after normalization and per_person_ratio_and_factor and pre_plot_aggregation.
import shutil
from pathlib import Path
import itertools
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import collections
from scipy.optimize import minimize_scalar

cols1 = ['F1_' + str(i) for i in range(3, 20, 2)]
cols2 = ['F2_' + str(i) for i in range(3, 20, 2)]
kCols = cols1 + cols2

def PlotWithSlices(df, data_name, output_dir):
    for group_name in ['Gender', 'AgeGroup', 'Family1', 'Family2', 'Family3', 'Family4', 'Education1', 'Career1', 'Career2', 'Language1', 'Word']:
        grouped_df = df.groupby([group_name])[kCols].mean()
        # grouped_df.to_csv(output_dir / (data_name + '_' + group_name + '_raw.csv'), index=True)
        full_group_name = '@'.join([data_name, group_name])
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        print(group_name)
        z_label = grouped_df.index.to_numpy().tolist()
        print(z_label)
        cmap = plt.get_cmap('viridis')
        colors = cmap(np.linspace(0, 1, len(z_label)))
        for key in z_label:
            x = np.arange(0, 9)
            color = colors[z_label.index(key)]
            z = z_label.index(key)
            mdf = grouped_df.loc[key]
            y1 = mdf[cols1].to_numpy(dtype='float')
            y2 = mdf[cols2].to_numpy(dtype='float')
            coeff1 = np.polyfit(x, y1, 4)
            coeff2 = np.polyfit(x, y2, 4)
            line1 = np.poly1d(coeff1)
            line2 = np.poly1d(coeff2)
            line1dd = np.polyder(line1, 2)
            line2dd = np.polyder(line2, 2)
            line1dd_max = minimize_scalar(-line1dd,
                                          bounds=(0, 8), method='bounded')
            line2dd_max = minimize_scalar(-line2dd,
                                          bounds=(0, 8), method='bounded')
            inflection1 = line1dd_max.x
            inflection2 = line2dd_max.x
            inflection1y = line1(inflection1)
            inflection2y = line2(inflection2)
            ax.plot(x, y1, zs=z, zdir='x', c=color, label='F1', linewidth=3.0)
            ax.plot(x, y2, zs=z, zdir='x', c=color, label='F2')
            ax.plot([inflection1, inflection1], [inflection1y-100, inflection1y+100], zs=z, zdir='x', c='black')
            ax.plot([inflection2, inflection2], [inflection2y-100, inflection2y+100], zs=z, zdir='x', c='black')

        ax.set(xticks=range(len(z_label)), xticklabels=z_label)
        plt.title(full_group_name)
        plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
        plt.savefig(output_dir / (full_group_name + '.png'),
                    bbox_inches="tight")
        plt.clf()
        plt.cla()


def PlotNoSlice(df, full_group_name, output_dir):
    x = np.arange(0, 9)
    y1 = df[cols1].to_numpy(dtype='float')
    y2 = df[cols2].to_numpy(dtype='float')
    coeff1 = np.polyfit(x, y1, 4)
    coeff2 = np.polyfit(x, y2, 4)
    line1 = np.poly1d(coeff1)
    line2 = np.poly1d(coeff2)
    line1dd = np.polyder(line1, 2)
    line2dd = np.polyder(line2, 2)
    line1dd_max = minimize_scalar(-line1dd,
                                  bounds=(0, 8), method='bounded')
    line2dd_max = minimize_scalar(-line2dd,
                                  bounds=(0, 8), method='bounded')
    inflection1 = line1dd_max.x
    inflection2 = line2dd_max.x

    # Plot f1/f2
    plt.plot(x, y1, 'o')
    plt.plot(x, y2, 'x')
    plt.plot(x, line1(x), label='F1 fitted')
    plt.plot(x, line2(x), label='F2 fitted')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.title(full_group_name)
    plt.savefig(output_dir / (full_group_name + '.fitted.png'),
                bbox_inches="tight")
    plt.clf()
    plt.cla()
    # Plot deriv and inflection
    plt.plot(x, line1dd(x), label='F1 2nd deriv')
    plt.plot(x, line2dd(x), label='F2 2nd deriv')
    plt.axvline(x=inflection1, linestyle=':', label='F1 break')
    plt.axvline(x=inflection2, linestyle='-.', label='F2 break')
    plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
    plt.title(full_group_name)
    plt.savefig(output_dir / (full_group_name + '.break.png'),
                bbox_inches="tight")
    plt.clf()
    plt.cla()


def SlicePlotDataS(df, output_dir):
    cols1 = ['F1_' + str(i) for i in range(3, 20, 2)]
    cols2 = ['F2_' + str(i) for i in range(3, 20, 2)]
    kCols = cols1 + cols2
    matched_rows = []
    sa_a1_sb_a1 = df[df['IsSba2'] == 'No']
    #sa_a1_sb_a1.to_csv(output_dir / 'sa_a1_sb_a1_raw.csv', index=False)
    sa_a1_sb_a1_mean = sa_a1_sb_a1.groupby(['Pos'])[kCols].mean()
    #sa_a1_sb_a1_mean.to_csv(output_dir / 'sa_a1_sb_a1_mean.csv', index=True)
    PlotNoSlice(sa_a1_sb_a1_mean.iloc[0], 'sa_a1_sb_a1_a', output_dir)
    PlotNoSlice(sa_a1_sb_a1_mean.iloc[1], 'sa_a1_sb_a1_b', output_dir)
    sa_a1_sb_a2 = df[df['IsSba2'] == 'Yes']
    #sa_a1_sb_a2.to_csv(output_dir / 'sa_a1_sb_a2_raw.csv', index=False)
    sa_a1_sb_a2_mean = sa_a1_sb_a2.groupby(['Pos'])[kCols].mean()
    #sa_a1_sb_a2_mean.to_csv(output_dir / 'sa_a1_sb_a2_mean.csv', index=True)
    PlotNoSlice(sa_a1_sb_a2_mean.iloc[0], 'sa_a1_sb_a2_a', output_dir)
    PlotNoSlice(sa_a1_sb_a2_mean.iloc[1], 'sa_a1_sb_a2_b', output_dir)

    matched_rows = []
    for _, row in df.iterrows():
        comps = row['Filename'].split('_')
        lang = comps[0]
        pos = comps[4]
        if lang == 'S' and pos == 'b' and row['Annotation'] == 'a2':
            matched_rows.append(row)
    input_df = pd.DataFrame(matched_rows)
    PlotWithSlices(input_df, 'all_s_sb_a2', output_dir)


def SlicePlotDataM(df, output_dir):
    m_sb_a1 = df[df['IsSba2'] == 'No']
    PlotWithSlices(m_sb_a1, 'm_sb_a1', output_dir)
    m_sb_a1_mean = m_sb_a1.groupby(['IsSba2'])[kCols].mean()
    PlotNoSlice(m_sb_a1_mean.iloc[0], 'm_sb_a1', output_dir)

    m_sb_a2 = df[df['IsSba2'] == 'Yes']
    PlotWithSlices(m_sb_a2, 'm_sb_a2', output_dir)
    m_sb_a2_mean = m_sb_a2.groupby(['IsSba2'])[kCols].mean()
    PlotNoSlice(m_sb_a2_mean.iloc[0], 'm_sb_a2', output_dir)


input_base_dir = Path('./analysis/output/')
output_base_dir = Path('./analysis/output/break/')
shutil.rmtree(output_base_dir, ignore_errors=True)
output_base_dir.mkdir(parents=True, exist_ok=True)

df = pd.read_csv(input_base_dir / 'S_all_plot_raw_data.csv')
SlicePlotDataS(df, output_base_dir)

df = pd.read_csv(input_base_dir / 'M_all_plot_raw_data.csv')
SlicePlotDataM(df, output_base_dir)