# Base Dependencies

# -----------------

from pathlib import Path

from os.path import join as pjoin

# Local Dependencies

# ------------------

from evaluation.io import collect_step_times

# 3rd-Party Dependencies

# ----------------------

import matplotlib.pyplot as plt

import pandas as pd

import seaborn as sns

# Constants

# ---------

COLOR_PALETTE = "Set2"

# Auxiliar Functions

# ------------------

def rename_strategies(x):

    if x == "BatchLC" or x == "BatchBALD":

        return "BatchLC / BatchBALD"

    else:

        return x

# Main Functions

# --------------

def step_time_barplot():

    """Plots a strip plot showing the step time results for the DDI Extraction corpus"""

    title = "Average Step Time per Method and Query Strategy"

    output_file = Path(pjoin("results", "step_time_barplot.png"))

    # collect results

    n2c2_results = collect_step_times(Path(pjoin("results", "n2c2", "all")))

    ddi_results = collect_step_times(Path(pjoin("results", "ddi")))

    n2c2_results["Corpus"] = ["n2c2"] * len(n2c2_results)

    ddi_results["Corpus"] = ["DDI"] * len(ddi_results)

    results = pd.concat([n2c2_results, ddi_results], ignore_index=True)

    # convert step time to minutes

    results["iter_time (average)"] = results["iter_time (average)"].apply(

        lambda x: x / 60

    )

    results["strategy"] = results["strategy"].apply(rename_strategies)

    # plot

    sns.set()

    sns.set_style("whitegrid")  # grid style

    colors = sns.color_palette(COLOR_PALETTE)

    fig, axes = plt.subplots(2)

    g = sns.FacetGrid(results, row="Corpus", aspect=2, legend_out=True, sharex=False)

    g.map(sns.barplot, "iter_time (average)", "method", "strategy", palette=colors)

    g.add_legend(title="Query Strategy")

    g.set_ylabels("Method")

    g.set_xlabels("Step Time (minutes)")

    g.set_yticklabels(["RF", "BiLSTM", "CBERT", "CBERT-pairs"])

    sns.move_legend(

        g,

        "lower center",

        bbox_to_anchor=(0.45, 1),

        ncol=3,

        title=None,

        frameon=True,

    )

    # ax.bar_label(ax.containers[0], fmt='%.f%%')

    if title:

        # plt.title(title)

        pass

    if output_file:

        plt.savefig(output_file, bbox_inches="tight")

    else:

        plt.show()

    plt.clf()