'''
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
'''
import pandas as pd 
import matplotlib.pyplot as plt
import numpy as np 
import os 
import glob
import sys
config_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
sys.path.append(config_dir)
from config import RESULTS_FOLDER

# %%
def plot_train_logs(train_fpaths, valid_fpaths, network_names):
    train_dfs = [pd.read_csv(path) for path in train_fpaths]
    valid_dfs = [pd.read_csv(path) for path in valid_fpaths]

    train_losses = [df['Loss'].values for df in train_dfs]
    valid_metrics = [df['Metric'].values for df in valid_dfs]
    train_epochs = [np.arange(len(train_loss))+1 for train_loss in train_losses]
    valid_epochs = [2*(np.arange(len(valid_metric))+1) for valid_metric in valid_metrics]
    min_losses = [np.min(train_loss) for train_loss in train_losses]
    min_losses_epoch = [np.argmin(train_loss) + 1 for train_loss in train_losses]
    max_dscs = [np.max(valid_metric) for valid_metric in valid_metrics]
    max_dscs_epoch = [2*(np.argmax(valid_metric)+1) for valid_metric in valid_metrics]
    fig, ax = plt.subplots(1,2, figsize=(20,10))
    fig.patch.set_facecolor('white')
    fig.patch.set_alpha(1)

    for i in range(len(train_losses)):
        ax[0].plot(train_epochs[i], train_losses[i])
        ax[1].plot(valid_epochs[i], valid_metrics[i])
        ax[0].plot(min_losses_epoch[i], min_losses[i], '-o', color='red')
        ax[1].plot(max_dscs_epoch[i], max_dscs[i], '-o', color='red')
        
        ax[0].text(np.min(train_epochs[i]), np.min(train_losses[i]), f'Total epochs: {len(train_epochs[i])}', fontsize=15)
        
    legend_labels_trainloss = [f"{network_names[i]}; Min loss: {round(min_losses[i], 4)} ({len(train_epochs[i])})" for i in range(len(network_names))]
    legend_labels_validdice = [f"{network_names[i]}; Max DSC: {round(max_dscs[i], 4)} ({len(valid_epochs[i])})" for i in range(len(network_names))]

    ax[0].legend(legend_labels_trainloss, fontsize=16)
    ax[1].legend(legend_labels_validdice, fontsize=16)
    ax[0].set_title('Train loss', fontsize=25)
    ax[1].set_title('Valid DSC', fontsize=25)
    ax[0].set_ylabel('Dice loss', fontsize=20)
    ax[1].set_ylabel('Dice score', fontsize=20)
    ax[0].grid(True)
    ax[1].grid(True)
    plt.show()

#%%
fold = 0
network = ['unet']
inputsize = [192, 192, 160, 128]
p = 2
inputsize_dict = {
    'unet': 192,
    'attentionunet': 192,
    'segresnet': 192,
    'dynunet': 160,
    'unetr': 160,
    'swinunetr': 128
}

experiment_code = [f"{network[i]}_fold{fold}_randcrop{inputsize[i]}" for i in range(len(network))]
save_logs_dir = os.path.join(RESULTS_FOLDER, 'logs')
save_logs_folders = [os.path.join(save_logs_dir, 'fold'+str(fold), network[i], experiment_code[i]) for i in range(len(experiment_code))]
train_fpaths = [os.path.join(save_logs_folders[i], 'trainlog_gpu0.csv') for i in range(len(save_logs_folders))]
valid_fpaths = [os.path.join(save_logs_folders[i], 'validlog_gpu0.csv') for i in range(len(save_logs_folders))]
legend_lbls = [f'{network[i]}, N = {inputsize[i]}' for i in range(len(network))]
plot_train_logs(train_fpaths, valid_fpaths, legend_lbls)