import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import argparse
import os

def get_rc_params():
    colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
    rcParams = plt.rcParams
    # Use LaTeX for rendering
    # Setup fonts
    rcParams["text.usetex"] = True
    rcParams["font.family"] = "serif"
    rcParams["font.size"] = 12
    rcParams["axes.labelsize"] = 12
    rcParams["axes.titlesize"] = 12
    rcParams["legend.fontsize"] = 10
    rcParams["xtick.labelsize"] = 10
    rcParams["ytick.labelsize"] = 10
    # Figure size and resolution
    rcParams["figure.figsize"] = (6, 4)
    rcParams["figure.dpi"] = 300
    # Use colors from the palette
    rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors])
    # Grid
    rcParams["axes.grid"] = True
    rcParams["grid.alpha"] = 0.5
    rcParams["grid.linestyle"] = "--"
    # Point ticks to the inside of the axes
    rcParams["xtick.direction"] = "in"
    rcParams["ytick.direction"] = "in"
    rcParams["xtick.top"] = True
    rcParams["ytick.right"] = True
    return rcParams

parser = argparse.ArgumentParser(description="Plot timing results from CSV file.")
parser.add_argument("csv_file", help="Path to the CSV file containing timing results.")
args = parser.parse_args()

plt.rcParams.update(get_rc_params())
df = pd.read_csv(args.csv_file)
for column in df.columns[1:]:
    df[column] = df[column] / df["N"] # Normalize by N

fig, ax = plt.subplots(figsize=(6,3))
x = np.arange(len(df))
bar_width = 0.2
for i, column in enumerate(df.columns[1:]):
    ax.bar(x + i * bar_width, df[column], width=bar_width, label=column)
ax.set_xticks(x + bar_width / 2)
ax.set_xticklabels(df["N"])
ax.set_xlabel("N (number of values)")
ax.set_ylabel("Time (seconds) per N")
ax.legend(ncol=3)

fig.tight_layout()

output_directory = os.path.abspath(os.path.join(__file__, "../../../../projects/project1/include"))
output_file = os.path.basename(args.csv_file).replace(".csv", ".pdf")
fig.savefig(os.path.join(output_directory, output_file))