"""
Module for statistics and plotting.
"""

import matplotlib.pyplot as plt

from pumaguard.traininghistory import (
    TrainingHistory,
)


def plot_training_progress(filename, full_history):
    """
    Plot the training progress and store in file.
    """
    plt.figure(figsize=(18, 10))
    plt.subplot(1, 2, 1)
    plt.plot(full_history.history["accuracy"], label="Training Accuracy")
    plt.plot(full_history.history["val_accuracy"], label="Validation Accuracy")
    plt.legend(loc="lower right")
    plt.ylabel("Accuracy")
    plt.ylim([min(plt.ylim()), 1])
    plt.title("Training and Validation Accuracy")

    plt.subplot(1, 2, 2)
    plt.plot(full_history.history["loss"], label="Training Loss")
    plt.plot(full_history.history["val_loss"], label="Validation Loss")
    plt.legend(loc="upper right")
    plt.ylabel("Cross Entropy")
    plt.ylim([0, 1.0])
    plt.title("Training and Validation Loss")

    print("Created plot of learning history")
    plt.savefig(filename)


def print_training_stats(full_history: TrainingHistory):
    """
    Print some stats of training.
    """
    print(
        f'Total time {sum(full_history.history["duration"])} for '
        f'{len(full_history.history["accuracy"])} epochs'
    )

    best_accuracy, best_val_accuracy, best_loss, best_val_loss, best_epoch = (
        full_history.get_best_epoch("accuracy")
    )
    print(
        f"Best accuracy - epoch {best_epoch} - accuracy: "
        f"{best_accuracy:.4f} - val_accuracy: {best_val_accuracy:.4f} - "
        f"loss: {best_loss:.4f} - val_loss: {best_val_loss:.4f}"
    )

    best_accuracy, best_val_accuracy, best_loss, best_val_loss, best_epoch = (
        full_history.get_best_epoch("val_accuracy")
    )
    print(
        f"Best val_accuracy - epoch {best_epoch} - accuracy: "
        f"{best_accuracy:.4f} - val_accuracy: {best_val_accuracy:.4f} "
        f"- loss: {best_loss:.4f} - val_loss: {best_val_loss:.4f}"
    )

    plot_training_progress("training-progress.png", full_history=full_history)
