import numpy as np
import matplotlib.pyplot as plt

#File for plotting data generated by model

def plot_epoch_loss(epoch_loss, epochs, title):
    """
    Plotting epoch loss of trained model. Saves image in current directory.

    Args:
        epoch_loss (list): Total loss from training
        epochs (int): Epochs trained for
        title (string): File title

    No returns
    """
    
    plt.semilogy(np.linspace(1, epochs, epochs),epoch_loss)
    plt.title("Epoch loss")
    plt.grid()
    plt.xlabel('epochs')
    plt.ylabel('loss')

    plt.savefig(title)
    plt.clf()
    return

def plot_ivp_loss(ivp_loss, epochs, title):
    """
    Plotting just initial value loss of trained model. Saves image in current directory.

    Args:
        ivp_loss (list): ivp loss from training
        epochs (int): Epochs trained for
        title (string): File title

    No returns
    """

    plt.semilogy(np.linspace(1, epochs, epochs),ivp_loss)
    plt.grid()
    plt.xlabel('epochs')
    plt.ylabel('IVP loss')

    plt.savefig(title)
    plt.clf()
    return

def plot_de_loss(de_loss, epochs, title):
    """
    Plotting just differential equation loss of trained model. Saves image in current directory.

    Args:
        de_loss (list): de loss from training
        epochs (int): Epochs trained for
        title (string): File title

    No returns
    """

    plt.semilogy(np.linspace(1, epochs, epochs),de_loss)
    plt.grid()
    plt.xlabel('epochs')
    plt.ylabel('DE loss')

    plt.savefig(title)
    plt.clf()
    return

def plot_solution_prediction(t, solPred, title):
    """
    Plotting predicted solution of trained model. Saves image in current directory.

    Args:
        t (list): Equally spaced points along t for evaluating model
        solPred (list): Solution prediction of model along t
        title (string): File title

    No returns
    """

    plt.plot(t, solPred)
    plt.title("Neural network solution")
    plt.grid()
    plt.xlabel('t')
    plt.ylabel('Predicted u')

    plt.savefig(title)
    plt.clf()
    return

def plot_all_losses(epoch_loss, de_loss, ivp_loss, epochs, title):
    """
    Plotting all different losses of trained model. Saves image in current directory.

    Args:
        epoch_loss (list): Total loss from training
        de_loss (list): de loss from training
        ivp_loss (list): ivp loss from training
        epochs (int): Epochs trained for
        title (string): File title

    No returns
    """

    plt.semilogy(np.linspace(1, epochs, epochs), ivp_loss, label='IVP Loss')
    plt.semilogy(np.linspace(1, epochs, epochs), de_loss, label='DE Loss')
    plt.semilogy(np.linspace(1, epochs, epochs), epoch_loss, label='Total Loss')
    plt.grid()
    plt.xlabel('Epochs')   
    plt.ylabel('Loss')
    plt.legend()

    plt.savefig(title)
    plt.clf()
    return

def plot_predicted_exact(t, solPred, exact, title):
    """
    Plotting predicted solution of model vs input function of solution. Saves image in current directory.

    Args:
        t (list): Equally spaced points along t for evaluating model
        solPred (list): Solution prediction of model along t
        exact (list): list of exact solution evaluated along t
        title (string): File title

    No returns
    """
    
    plt.plot(t, solPred)
    plt.plot(t, exact, "--")
    plt.title("Neural network solution vs Exact solution")
    plt.legend(["u(t)", "Exact u(t)"])
    plt.grid()
    plt.xlabel("t")

    plt.savefig(title)
    plt.clf()

    return

def plot_solution_prediction_system(t, solPred, title):
    """
    Plotting predicted solutions of trained model. Saves image in current directory.

    Args:
        t (list): Equally spaced points along t for evaluating model
        solPred (list): List of lists of solution predictions of model along t
        title (string): File title

    No returns
    """

    if len(solPred) == 2:
        plt.plot(t, solPred[0])
        plt.plot(t, solPred[1])
        plt.title("Neural network solution")
        plt.grid()
        plt.legend(["u(t)", "x(t)"])
        plt.xlabel('t')
        plt.ylabel('u, x')

        plt.savefig(title)
        plt.clf()

    if len(solPred) == 3:
        plt.plot(t, solPred[0])
        plt.plot(t, solPred[1])
        plt.plot(t, solPred[2])
        plt.title("Neural network solution")
        plt.grid()
        plt.legend(["u(t)", "x(t)", "y(t)"])
        plt.xlabel('t')
        plt.ylabel('u, x, y')

        plt.savefig(title)
        plt.clf()

    return

def plot_predicted_exact_system(t, solPred, exact, title):
    """
    Plotting predicted solutions of model vs input functions of solution. Saves image in current directory.

    Args:
        t (list): Equally spaced points along t for evaluating model
        solPred (list): List of lists of solution predictions of model along t
        exact (list): lList of lists of exact solutions evaluated along t
        title (string): File title

    No returns
    """
    
    if len(solPred) == 2:
        plt.plot(t, solPred[0])
        plt.plot(t, solPred[1])
        plt.plot(t, exact[0], "--")
        plt.plot(t, exact[1], "--")
        plt.title("Neural network against exact solution")
        plt.legend(["u(t)", "x(t)", "Exact u(t)", "Exact x(t)"])
        plt.grid()
        plt.xlabel("t")
        plt.ylabel('u, x')

        plt.savefig(title)
        plt.clf()

    if len(solPred) == 3:
        plt.plot(t, solPred[0])
        plt.plot(t, solPred[1])
        plt.plot(t, solPred[2])
        plt.plot(t, exact[0], "--")
        plt.plot(t, exact[1], "--")
        plt.plot(t, exact[2], "--")
        plt.title("Neural network against exact solution")
        plt.legend(["u(t)", "x(t)", "y(t)", "Exact u(t)", "Exact x(t)", "Exact y(t)"])
        plt.grid()
        plt.xlabel("t")
        plt.ylabel('u, x, y')

        plt.savefig(title)
        plt.clf()

    return