from matplotlib.patches import Patch
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Ignore the warning generated by Matplotlib for setting a y-lim of 0.01 on a
# log scale
import warnings
warnings.simplefilter("ignore", UserWarning)

# Formatted dataset names
FORMATTED_AA = r'\texttt{WH-aa}'
FORMATTED_J = r'\texttt{WH-j}'
FORMATTED_B_50 = r'\texttt{b-50}'
FORMATTED_B_100 = r'\texttt{b-100}'
FORMATTED_GKA = r'\texttt{GKA}'

# Formatted generator names
FORMATTED_ER = r'Erd\"os-R\'enyi'
FORMATTED_CL = r'Chung-Lu'
FORMATTED_BA = r'Barab\'asi-Albert'
FORMATTED_TO = r'Tunable-OCT'


def _get_dataset(name):
    if 'aa' in name:
        return FORMATTED_AA
    elif 'j' in name:
        return FORMATTED_J
    elif 'b-50' in name:
        return FORMATTED_B_50
    elif 'b-100' in name:
        return FORMATTED_B_100
    elif 'gka' in name:
        return FORMATTED_GKA


def _get_generator(name):
    if '-er' in name:
        return FORMATTED_ER
    elif '-cl' in name:
        return FORMATTED_CL
    elif '-ba' in name:
        return FORMATTED_BA
    elif '-to' in name:
        return FORMATTED_TO


def _get_dataset_order():
    return [FORMATTED_AA,
            FORMATTED_J,
            FORMATTED_GKA,
            FORMATTED_B_50,
            FORMATTED_B_100]


def _get_generator_order():
    return [FORMATTED_ER,
            FORMATTED_TO,
            FORMATTED_CL,
            FORMATTED_BA]


def construct_df(exact_filename, solver1_name, solver2_name, yaxis_title):
    rows = []
    solver1 = {}
    solver2 = {}
    opt = {}
    datasets = set()

    with open(exact_filename, 'r') as infile:
        # Discard header
        infile.readline()
        for line in infile.readlines():
            line = line.split(',')
            datasets.add(line[0])
            solver = line[1]
            if solver == solver1_name:
                solver1[line[0]] = float(line[2])
                opt[line[0]] = float(line[3])
            elif solver == solver2_name:
                solver2[line[0]] = float(line[2])

    for dataset in datasets:
        if dataset in solver1 and dataset in solver2 and dataset in opt:
            row = [dataset,
                   max(solver1[dataset], 0.01),
                   solver2[dataset],
                   opt[dataset],
                   _get_dataset(dataset),
                   _get_generator(dataset)]
            rows.append(row)
    df = pd.DataFrame(rows, columns=['name',
                                     solver1_name,
                                     solver2_name,
                                     "$OPT'$",
                                     'Dataset',
                                     'Generator'])

    df[yaxis_title] = df[solver1_name] / df[solver2_name]
    return df


def _compute_plot(solver1, solver2, title, exact_data_filename,
                  output_filename, palette):
    df = construct_df(exact_data_filename, solver1, solver2, 'Run Time')

    # Prepare a new plot
    sns.set(font_scale=1.7)
    plt.clf()

    # Compute the scatterplot
    ax = sns.relplot(
        data=df,
        x="$OPT'$",
        y='Run Time',
        col='Dataset',
        col_wrap=3,
        col_order=_get_dataset_order(),
        hue='Generator',
        hue_order=_get_generator_order(),
        s=60,
        palette=palette,
        alpha=0.8,
        linewidth=0,
        legend=False,  # We generate the legend manually
    )

    # Add legend
    legend_colors = zip(_get_generator_order(), palette)
    ax.add_legend(
        title=r'\textbf{Generator}',
        legend_data={
            name: Patch(color=color)
            for name, color in legend_colors
        },
        label_order=_get_generator_order(),
        bbox_to_anchor=(
            0.78,   # Percent of total width
            0.72,   # Percent of total height
            0,
            0
        ),
    )

    # Modify titles
    # Top title
    plt.subplots_adjust(top=0.9)
    ax.fig.suptitle(title)

    # Facet titles
    ax.set_ylabels(r'\textbf{Run Time Ratio}')
    ax.set_xlabels(r'\textbf{Minimum OCT}')

    # Adjust scales
    ax.set(yscale="log", ylim=(0.01, 1000), xlim=(-3, 100))

    # Give the lonely plot some friends (move to bottom row next to last plot).
    lone_plot = ax.fig.axes[2]
    pos = lone_plot.get_position()
    lone_plot.set_position([
        pos.x0,
        ax.fig.axes[4].get_position().y0,
        pos.width,
        pos.height
    ])

    # Add facet titles and hlines
    titles = _get_dataset_order()
    for axis in ax.fig.axes:
        axis.set_title(titles.pop(0))
        axis.axhline(y=1, color='black', dashes=[3, 3])

    # Save to file
    ax.savefig(output_filename)


if __name__ == "__main__":
    # Use latex text
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')

    # Define color palette
    colors = ["dusty purple", "windows blue", "amber", "faded green"]
    palette = sns.xkcd_palette(colors)

    # Generate plots
    _compute_plot(solver1='vc',
                  solver2='ilp_1t',
                  title=r'Relative Run Times of \textsf{VC} to \textsf{ILP}',
                  exact_data_filename='results/exact_results.csv',
                  output_filename='figures/figure2.pdf',
                  palette=palette)
    print('Generated VC vs ILP_1t (figure2)')
    _compute_plot(solver1='ilp_1t',
                  solver2='ilp',
                  title=r'Relative Run Times of 1-Thread \textsf{ILP} to 4-Thread \textsf{ILP}',
                  exact_data_filename='results/exact_results.csv',
                  output_filename='figures/figure3.pdf',
                  palette=palette)
    print('Generated ILP_1t vs ILP (figure3)')
