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

parser = argparse.ArgumentParser(description='Plot buckling beam data.')
parser.add_argument('input_file', type=str, help='Path to the input CSV file')

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

x = np.linspace(0, 1, len(data))
fig, axs = plt.subplots(1, 2, figsize=(12, 5), sharex=True)
for col in data.columns:
    c = f"C{int(col[-1])}"
    ls = "-" if "jacobi" in col else "--"
    label = col.replace("_", " ").title()
    axs[0].plot(x, data[col], label=label, color=c, linestyle=ls)
    if "analytical" in col:
        jacobi_col = col.replace("analytical", "jacobi")
        if jacobi_col in data.columns:
            error = np.abs(data[col] - data[jacobi_col])
            axs[1].plot(x, error, label=f'Error {label.split(" ")[-1]}', color=c, linestyle=':')
axs[1].set_xlabel('x')
axs[1].set_ylabel('Error')
axs[1].set_title('Buckling Beam Errors')
axs[1].legend()
axs[1].grid()
axs[0].set_xlabel('x')
axs[0].set_ylabel('Displacement')
axs[0].set_title('Buckling Beam Displacements')
axs[0].legend()
axs[0].grid()

dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../projects/project2/include"))
filename = args.input_file.split("/")[-1].replace(".txt", "")
fig.savefig(f"{dir}/{filename}_plot.pdf")