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

parser = argparse.ArgumentParser(description='Plot scaling data.')
parser.add_argument('input_file', type=str, help='Path to the input CSV file')
parser.add_argument("--times", "-t", action="store_true", help="Plot times instead of iterations")
parser.add_argument("--no-loglog", "-nL", action="store_true", help="Do not use log-log scale")

args = parser.parse_args()
data = pd.read_csv(args.input_file)

x = data["N"]
fig, ax = plt.subplots(figsize=(8, 5))
if args.times:
    ax.plot(x, data["Jacobi_time"], label='Jacobi Time', color='C0', linestyle='-', marker='o')
    ax.plot(x, data["Arma_time"], label='Armadillo Time', color='C1', linestyle='-', marker='o')
    ax.set_ylabel('Time (s)')
    ax.set_title('Scaling of Computation Time')
    ax.legend()
    ax.grid()
else:
    ax.set_ylabel('Number of Iterations')
    ax.set_title('Scaling of Jacobi Iterations')
    ax.plot(x, data["Jacobi_iterations"], label='Jacobi Iterations', color='C0', linestyle='-', marker='o')
    ax.grid()

if not args.no_loglog:
    ax.set_xscale('log')
    ax.set_yscale('log')
ax.set_xlabel('Matrix Size N')
dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../projects/project2/include"))
filename = args.input_file.split("/")[-1].replace(".txt", "")
if args.times:
    fig.savefig(f"{dir}/{filename}_times.pdf")
else:
    fig.savefig(f"{dir}/{filename}_iterations.pdf")