#!/usr/bin/env python
# coding: utf-8

import os
import re
import matplotlib.pyplot as plt
from datetime import timedelta

# ---------------- SETTINGS ----------------
base_dir = os.getcwd()
slurm_dir = os.path.join(base_dir, "slurm_jobs")
pattern_elapsed = re.compile(r"✅ Elapsed:\s*(\d+):(\d+):(\d+)")

# ---------------- PARSING ----------------
data = {}  # {(W, Epp, Ewp): [(L, seconds), ...]}

for fname in os.listdir(slurm_dir):
    if not fname.endswith(".out"):
        continue
    fpath = os.path.join(slurm_dir, fname)

    # extract parameters from filename, e.g. L20_W30_Epp-2.0_Ewp-3.0.out
    mL = re.search(r"L(\d+)", fname)
    mW = re.search(r"W(-?\d+(?:\.\d+)?)", fname)
    mEpp = re.search(r"Epp(-?\d+(?:\.\d+)?)", fname)
    mEwp = re.search(r"Ewp(-?\d+(?:\.\d+)?)", fname)

    if not (mL and mW and mEpp and mEwp):
        print(f"⚠️ Skipping {fname}: missing parameters.")
        continue

    L = int(mL.group(1))
    W = float(mW.group(1))
    Epp = float(mEpp.group(1))
    Ewp = float(mEwp.group(1))

    with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read()

    mtime = pattern_elapsed.search(content)
    if not mtime:
        print(f"⏳ No elapsed time found in {fname}")
        continue

    h, m, s = map(int, mtime.groups())
    seconds = h * 3600 + m * 60 + s

    key = (W, Epp, Ewp)
    data.setdefault(key, []).append((L, seconds))

    print(f"✅ {fname}: Elapsed {timedelta(seconds=seconds)}")

# ---------------- PLOT ----------------
if not data:
    print("❌ No elapsed times found.")
    exit()

plt.figure(figsize=(8, 6))

for (W, Epp, Ewp), values in sorted(data.items()):
    values.sort(key=lambda x: x[0])
    Ls, times = zip(*values)
    label = f"W={W}, Epp={Epp}, Ewp={Ewp}"
    plt.plot(Ls, times, "o-", label=label)

plt.xlabel("L (chain length)")
plt.ylabel("Elapsed time (s)")
plt.title("Simulation time vs L")
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.savefig("elapsed_vs_L.png", dpi=200)
#plt.show()

print("\n✅ Plot saved as elapsed_vs_L.png")