import sys
import os
import numpy as np
import matplotlib.pyplot as plt
import glob
import re

# Check for reference file argument
if len(sys.argv) != 2:
    print("Usage: python plot_convergence.py <ref_output_path>")
    sys.exit(1)

ref_output = sys.argv[1]
ref_dt_match = re.search(r"zs_vec_dt_([0-9.]+)", os.path.basename(ref_output))

if not ref_dt_match:
    print("Error: Could not extract dt from reference output filename.")
    sys.exit(1)

ref_dt_str = ref_dt_match.group(1)
ref_vec = np.loadtxt(ref_output)
ymin, ymax = np.min(ref_vec), np.max(ref_vec)

# ---- Plot 1: Convergence Curve with Iteration Count ----
data = np.loadtxt("results/convergence_errors.txt", skiprows=1)
dts, errors, iters = data[:, 0], data[:, 1], data[:, 2]

fig, ax1 = plt.subplots()

color_l2 = "tab:blue"
ax1.set_xlabel("Time step size " + r"$(\Delta t)$")
ax1.set_ylabel("Relative " + r"$L_2$" + "-error in " + r"$h$", color=color_l2)
ax1.loglog(dts, errors, 'o-', label=r"Relative $L_2$-error", color=color_l2)
ax1.loglog(dts, errors[0]*(dts/dts[0])**1, '--', label=r"$\mathcal{O}(\Delta t)$", alpha=0.5, color="gray")
ax1.loglog(dts, errors[0]*(dts/dts[0])**2, '--', label=r"$\mathcal{O}(\Delta t^2)$", alpha=0.5, color="gray")
ax1.tick_params(axis='y', labelcolor=color_l2)
ax1.grid(True, which="both", linestyle="--")

# Second y-axis for iteration count
ax2 = ax1.twinx()
color_iter = "tab:red"
ax2.set_ylabel("Average number of iterations per time-step", color=color_iter)
ax2.semilogx(dts, iters, 's--', color=color_iter, label="Iterations")
ax2.tick_params(axis='y', labelcolor=color_iter)

# Combine legends
lines_1, labels_1 = ax1.get_legend_handles_labels()
lines_2, labels_2 = ax2.get_legend_handles_labels()
ax1.legend(lines_1 + lines_2, labels_1 + labels_2, loc='upper left')

os.makedirs("results", exist_ok=True)
#plt.title("Convergence with Iteration Count")
plt.tight_layout()
plt.savefig("results/convergence_plot.png")

# ---- Plot 2: Surface vectors ----
plt.figure()

# Plot reference first
plt.plot(ref_vec, label=f"dt = {ref_dt_str} (ref)", linestyle="solid", linewidth=2, color='black')

# Load test surfaces based on dt values from convergence file
test_files = [f"results/zs_vec_dt_{dt:.5f}".rstrip("0").rstrip(".") + ".txt" for dt in dts]

line_styles = ['-', '--', '-.', ':']
markers = ['o', 's', 'v', '^', 'D', 'x', '*']
style_cycle = [(ls, m) for ls in line_styles for m in markers]

for i, file in enumerate(test_files):
    if os.path.abspath(file) == os.path.abspath(ref_output):
        continue

    if not os.path.exists(file):
        print(f"Warning: Skipping missing file {file}")
        continue

    zs_vec = np.loadtxt(file)
    ls, m = style_cycle[i % len(style_cycle)]
    label = os.path.basename(file).replace("zs_vec_", "").replace(".txt", "")
    plt.plot(zs_vec, label=label, linestyle=ls, marker=m, markersize=4)

plt.ylim(ymin, ymax)
plt.xlabel("Index")
plt.ylabel("z value")
plt.title("Surface Vectors vs. Reference")
plt.legend()
plt.grid(True)
plt.savefig("results/surfaces_plot.png")
plt.show()