# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_Experiment.ipynb (unless otherwise specified).

__all__ = ['HephaestusModelEvaluation', 'annotateBarPlot', 'plotBar', 'plotTrainingAccuracies',
           'plotPerfectPredictionAccuracies', 'plotAllPerfectPredictionAccuracies', 'plotAvgEditDistDecreases',
           'plotAllAvgEditDistDecreases', 'plotFailedPredictionRates']

# Cell
#export
from typing import Dict, Any, List
import pandas as pd
import os
from copy import deepcopy

import sys
sys.path.append("..")

from .AbstractMethod import *
from .IOUtils import *
from .DatasetConstruction import *
from .HephaestusModel import *
from .TrainModels import *

# Cell
class HephaestusModelEvaluation:
    """
    Helper class to centralize computations when evaluating the effectiveness of a `HephaestusModel`.

    Required Arguments:

    - `model`: The `HephaestusModel` to evaluate
    - `testSourceMethods`: A list of AbstractMethods representing the buggy abstract methods in the testing set. These
      methods should not appear at all in the training or validation sets that were used to train the model.
    - `testTargetMethods`: A list of AbstractMethods representing the fixed abstract methods in the testing set. These
      methods should not appear at all in the training or validation sets that were used to train the model.

    Optional Arguments:

    - `isControl`: Set to `True` if the model represents the control, i.e. the model was trained purely with
      AbstractMethods and not with EditOperations. Defaults to `False`.

    Once created, the `HephaestusModelEvaluation` will contain the following attributes which can be freely accessed:

    - `outputMethods`: A list of AbstractMethods which were translated from the given buggy (source) AbstractMethods.
      These represent the model's predictions of what it thinks the fixed methods are. Some of these entries may be
      `None`, representing failed predictions.
    - `numPerfectPredictions`: The number of `outputMethods` whose tokens exactly match the actual fixed method according
      to the provided `testTargetMethods`.
    - `perfectPredictionRatio`: The ratio of perfect predictions to total predictions, in the range [0, 1].
    - `numFailedPredictions`: The number of `outputMethods` which were not able to be predicted due to ill-formed
      EditOperations or other factors. These are represented by `None` values in `outputMethods`.
    - `failedPredictionRatio`: The ratio of failed predictions to total predictions, in the range [0, 1].
    - `avgEditDistDecrease`: A value representing how much the model "helped" in reducing the Levenshtein edit distance
      to the `testTargetMethods`. E.g. a value of 3 means that on average, the edit distance from the output methods to
      the target methods was 3 less than the edit distance from the source methods to the target methods. Negative values
      mean that the model made the output methods further away from the target methods than they originally were as given
      by the source methods.
    - `trainingStats`: A Pandas DataFrame representing model information while it was training, as in
      `HephaestusModel.getTrainingStats`.
    """

    def __init__(self,
        model: HephaestusModel,
        testSourceMethods: List[AbstractMethod],
        testTargetMethods: List[AbstractMethod],
        isControl: bool = False
    ) -> None:

        # use the model to determine output methods
        self.outputMethods = model.translate(testSourceMethods, applyEditOperations = not isControl)

        # sanity check -- make sure that the lengths of all lists of methods are equal
        if not len(testSourceMethods) == len(testTargetMethods) == len(self.outputMethods):
            raise RuntimeError("HephaestusModelEvaluation: method number mismatch")

        # determine number of perfect predictions and failed predictions, and edit distance stuff
        self.numPerfectPredictions = 0
        self.numFailedPredictions = 0
        editDistDecreases = []
        for outputMethod, sourceMethod, targetMethod in zip(self.outputMethods, testSourceMethods, testTargetMethods):
            if outputMethod == targetMethod:
                self.numPerfectPredictions += 1
            if outputMethod is None:
                self.numFailedPredictions += 1
            else:
                idealDist = sourceMethod.getEditDistanceTo(targetMethod)
                actualDist = outputMethod.getEditDistanceTo(targetMethod)
                editDistDecreases.append(idealDist - actualDist)
        self.avgEditDistDecrease = sum(editDistDecreases) / len(editDistDecreases)

        # obtain perfect prediction and failed prediction ratios
        self.perfectPredictionRatio = self.numPerfectPredictions / len(self.outputMethods)
        self.failedPredictionRatio = self.numFailedPredictions / len(self.outputMethods)

        # get the training statistics
        self.trainingStats = model.getTrainingStats()

# Cell
def annotateBarPlot(plot) -> None:
    """
    Annotates the bars of the given `plot` with their values.
    """

    maxSize = max([abs(p.get_height()) for p in plot.patches])

    for p in plot.patches:
        plot.annotate(
            "{:.1f}".format(p.get_height()),
            (p.get_x() + p.get_width() / 2, p.get_height() + maxSize * 0.05 * (1 if p.get_height() >= 0 else -1)),
            ha = "center",
            va = "center",
            size = 10
        )

# Cell
def plotBar(
    data: pd.DataFrame,
    title: str,
    xLabel: str,
    yLabel: str,
    annotate: bool = True
) -> pd.DataFrame:
    """
    Plots an annotated bar graph with the given parameters and returns the `data` DataFrame.
    """

    minVal = data.values.min()
    maxVal = data.values.max()
    maxSize = max(abs(minVal), abs(maxVal))

    # create plot
    plot = data.plot.bar(
        ylim = (0 if minVal >= 0 else minVal - maxSize * 0.15, 0 if maxVal <= 0 else maxVal + maxSize * 0.15),
        title = title,
        xlabel = xLabel,
        ylabel = yLabel,
        rot = 0
    )

    # annotate if necessary
    if annotate:
        annotateBarPlot(plot)

    return data

# Cell
def plotTrainingAccuracies(
    evaluations: List[HephaestusModelEvaluation],
    lineLabels: List[str],
    title: str
) -> None:
    """
    Plots the training accuracies of the given `evaluations` as a line graph per training step. The line labels are the
    given `lineLabels`. The title is the given `title`.
    """

    # create dataframe
    data = {}
    for evaluation, label in zip(evaluations, lineLabels):
        data[label] = evaluation.trainingStats["trainAccuracy"].to_list()
    frame = pd.DataFrame(data, index = evaluations[0].trainingStats["step"].to_list())

    # create plot
    plot = frame.plot.line(
        title = title,
        xlabel = "Training Step",
        ylabel = "Training Accuracy %",
        grid = True
    )

# Cell
def plotPerfectPredictionAccuracies(
    evaluations: List[HephaestusModelEvaluation],
    xLabels: List[str],
    title: str
) -> pd.DataFrame:
    """
    Plots the perfect prediction accuracies of the given `evaluations` as a bar graph. The x-axis labels are the
    given `xlabels`. The title of the graph is the given `title`. Returns the dataframe used to create the graph.
    """
    return plotBar(
        data = pd.DataFrame({
            "acc%": [100 * e.perfectPredictionRatio for e in evaluations]
        }, index = xLabels),
        title = title,
        xLabel = "Model",
        yLabel = "Perfect Prediction Accuracy %"
    )

# Cell
def plotAllPerfectPredictionAccuracies(
    evaluations: Dict[str, HephaestusModelEvaluation],
    xLabels: List[str]
) -> None:
    """
    Plots the perfect predictions accuracies of all given `evaluations` as a nested bar graph. The x-axis labels are the
    given `xLabels`. Returns the dataframe used to create the graph.
    """

    data = {}
    for key in evaluations:
        data[key] = [100 * e.perfectPredictionRatio for e in evaluations[key]]
    frame = pd.DataFrame(data, index = xLabels)

    return plotBar(
        data = pd.DataFrame(data, index = xLabels),
        title = "Perfect Prediction Accuracies of All Models",
        xLabel = "Model",
        yLabel = "Perfect Prediction Accuracy %",
        annotate = False
    )

# Cell
def plotAvgEditDistDecreases(
    evaluations: List[HephaestusModelEvaluation],
    xLabels: List[str],
    title: str
) -> pd.DataFrame:
    """
    Plots the average edit distance decreases of the given `evaluations` as a bar graph. The x-axis labels are the
    given `xlabels`. The title of the graph is the given `title`. Returns the dataframe used to create the graph.
    """
    return plotBar(
        data = pd.DataFrame({
            "decrease": [e.avgEditDistDecrease for e in evaluations]
        }, index = xLabels),
        title = title,
        xLabel = "Model",
        yLabel = "Average Edit Distance Decrease"
    )

# Cell
def plotAllAvgEditDistDecreases(
    evaluations: Dict[str, HephaestusModelEvaluation],
    xLabels: List[str]
) -> None:
    """
    Plots the average edit distance decreases of all given `evaluations` as a nested bar graph. The x-axis labels are the
    given `xLabels`. Returns the dataframe used to create the graph.
    """

    data = {}
    for key in evaluations:
        data[key] = [e.avgEditDistDecrease for e in evaluations[key]]
    frame = pd.DataFrame(data, index = xLabels)

    return plotBar(
        data = pd.DataFrame(data, index = xLabels),
        title = "Average Edit Distance Decreases of All Models",
        xLabel = "Model",
        yLabel = "Average Edit Distance Decrease",
        annotate = False
    )

# Cell
def plotFailedPredictionRates(
    evaluations: List[HephaestusModelEvaluation],
    xLabels: List[str],
    title: str
) -> pd.DataFrame:
    """
    Plots the failed prediction rates of the given `evaluations` as a bar graph. The x-axis labels are the
    given `xlabels`. The title of the graph is the given `title`. Returns the dataframe used to create the graph.
    """
    return plotBar(
        data = pd.DataFrame({
            "rate%": [100 * e.failedPredictionRatio for e in evaluations]
        }, index = xLabels),
        title = title,
        xLabel = "Model",
        yLabel = "Failed Prediction Rate %"
    )