import os
import matplotlib.pyplot as plt
import numpy as np
from .show_evaluation import fgsm_evaluate, construct_CREI

def show_performance(model_name, batch_size, dataset):
    epsilons = np.arange(0, 0.3, 0.01)  # Create a range of epsilon values
    if dataset == "CIFAR-10":
        epsilons = np.arange(0, 0.3, 0.01)
    elif dataset == "CIFAR-100":
        epsilons = np.arange(0, 0.3, 0.01)
    accuracies = []
    correct_confidences = []
    incorrect_confidences = []
    creis = []
    metrics = []

    for epsilon in epsilons:
        acc_before, conf_correct_before, conf_incorrect_before = fgsm_evaluate(model_name, 0, batch_size, dataset)
        acc_after, conf_correct_after, conf_incorrect_after = fgsm_evaluate(model_name, epsilon, batch_size, dataset)
        crei = construct_CREI(acc_before, acc_after, conf_correct_before, conf_correct_after, conf_incorrect_before, conf_incorrect_after)
        accuracies.append(acc_after)
        correct_confidences.append(conf_correct_after)
        incorrect_confidences.append(conf_incorrect_after)
        creis.append(crei)
        metrics.append([acc_after, conf_correct_after, conf_incorrect_after, crei])

    # Create the first figure for the line plots
    fig1, ax1 = plt.subplots(figsize=(10, 6))
    ax1.plot(epsilons, accuracies, label='$Acc$')
    ax1.plot(epsilons, correct_confidences, label=r'$Conf_+$', linestyle='--')
    ax1.plot(epsilons, incorrect_confidences, label=r'$Conf_-$', linestyle='--')
    ax1.plot(epsilons, creis, label='$CREI$')
    ax1.set_xlabel(r'$\epsilon$', fontsize=20)
    ax1.set_title('Performance over Epsilon for model {}'.format(model_name))
    ax1.legend()
    ax1.grid(True)

    base_name, extension = os.path.splitext(model_name)
    fig1.savefig(f'./attachments/PoE/{dataset}/{base_name}_PoE.png')

    return fig1