from sfransen.utils_quintin import *
import matplotlib.pyplot as plt
import argparse
import matplotlib.ticker as tkr
from umcglib.froc.p_auc import partial_auc
import numpy as np

parser = argparse.ArgumentParser(
    description='Visualise froc results')
parser.add_argument('-saveas', 
    help='')
parser.add_argument('-comparison', 
    help='')
parser.add_argument('--experiment', '-s', 
    metavar='[series_name]', required=True, nargs='+',
    help='List of series to include, must correspond with' +
        "path files in ./data/")
parser.add_argument('-yaml_metric',
    help='List of series to include, must correspond with' +
        "path files in ./data/")
args = parser.parse_args()

if args.comparison:
    colors = ['r','r','b','b','g','g','y','y']
    plot_type = ['-','--','-','--','-','--','-','--']
else: 
    colors = ['r','b','g','k','y','c']
    plot_type = ['-','-','-','-','-','-']

yaml_metric = args.yaml_metric
experiments = args.experiment
print(experiments)
experiment_path = []
auroc = []
paufroc = []
False_possitives = []
sensitivity = []

fig = plt.figure(1)
ax = fig.add_subplot(111)
for idx in range(len(args.experiment)):

    False_possitives_mean = np.linspace(0, 2.5, 200)

    for fold in range(5):
        print('fold:',fold)
        experiment_metrics = {}
        experiment_path = f'./../train_output/{experiments[idx]}_{fold}/froc_metrics_{yaml_metric}.yml'
        experiment_metrics = read_yaml_to_dict(experiment_path)
        pfroc = partial_auc(experiment_metrics["sensitivity"],experiment_metrics["FP_per_case"],low=0.1, high=2.5)
        paufroc.append(round(pfroc,2))
        False_possitives.append(experiment_metrics["FP_per_case"])
        sensitivity_ = np.interp(False_possitives_mean,experiment_metrics["FP_per_case"],experiment_metrics["sensitivity"])
        sensitivity.append(sensitivity_)

    # calculate mean and std
    sensitivity_mean = np.squeeze(np.mean(sensitivity,axis=0))
    sensitivity_std = np.multiply(np.squeeze(np.std(sensitivity,axis=0)),2)

    plt.plot(False_possitives_mean, sensitivity_mean,color=colors[idx],linestyle=plot_type[idx])
    plt.fill_between(False_possitives_mean, np.subtract(sensitivity_mean,sensitivity_std), np.add(sensitivity_mean,sensitivity_std))
    ax.set(xscale="log")
    ax.axes.xaxis.set_minor_locator(tkr.LogLocator(base=10, subs='all'))
    ax.axes.xaxis.set_minor_formatter(tkr.NullFormatter())
    ax.axes.xaxis.set_major_formatter(tkr.ScalarFormatter())
    ax.axes.grid(True, which="both", ls="--", c='#d3d3d3')
    ax.axes.set_xlim(left=0, right=2.5)
    ax.axes.xaxis.set_major_locator(tkr.FixedLocator([0,0.1,0.5,1,2.5]))

fpr = []
tpr = []
for idx in range(len(args.experiment)):
    experiment_path = f'./../train_output/{experiments[idx]}/froc_metrics_{yaml_metric}.yml'
    experiment_metrics = read_yaml_to_dict(experiment_path)
    auroc.append(round(experiment_metrics['auroc'],3))

    fpr_mean = np.linspace(0, 1, 200)

    for fold in range(5):
        print('fold:',fold)
        experiment_metrics = {}
        experiment_path = f'./../train_output/{experiments[idx]}_{fold}/froc_metrics_{yaml_metric}.yml'
        experiment_metrics = read_yaml_to_dict(experiment_path)
        # pfroc = partial_auc(experiment_metrics["tpr"],experiment_metrics["fpr"],low=0.1, high=2.5)
        paufroc.append(round(pfroc,2))
        fpr.append(experiment_metrics["fpr"])
        tpr_ = np.interp(fpr_mean,experiment_metrics["fpr"],experiment_metrics["tpr"])
        tpr.append(tpr_)

    tpr_mean = np.squeeze(np.mean(tpr,axis=0))
    tpr_std = np.multiply(np.squeeze(np.std(tpr,axis=0)),2)

    plt.figure(2)
    plt.plot(fpr_mean, tpr_mean,color=colors[idx],linestyle=plot_type[idx])
    plt.fill_between(fpr_mean, np.subtract(tpr_mean,tpr_std), np.add(tpr_mean,tpr_std))

print(auroc)
experiments = [exp.replace('train_10h_', '') for exp in experiments] 
experiments = [exp.replace('train_n0.001_', '') for exp in experiments] 
experiments = [exp.replace('_', ' ') for exp in experiments] 
# experiments = ['10% noise','1% noise','0.1% noise','0.05% noise']

concat_func = lambda x,y: x + " (" + str(y) + ")"
experiments_paufroc = list(map(concat_func,experiments,paufroc)) # list the map function

plt.figure(1)
plt.title('fROC curve')
plt.xlabel('False positive per case')
plt.ylabel('Sensitivity')
plt.legend(experiments_paufroc,loc='lower right')
# plt.xlim([0,50])
plt.grid()
plt.ylim([0,1])
plt.yticks([0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1])
plt.savefig(f"./../train_output/fROC_{args.saveas}.png", dpi=300)

concat_func = lambda x,y: x + " (" + str(y) + ")"
experiments_auroc = list(map(concat_func,experiments,auroc)) # list the map function

plt.figure(2)
plt.title('ROC curve')
plt.legend(experiments_auroc,loc='lower right')
plt.xlabel('False positive rate')
plt.ylabel('True positive rate')
plt.grid()
plt.savefig(f"./../train_output/ROC_{args.saveas}.png", dpi=300)