#!/usr/bin/env python3
import os, re, glob, csv, math, subprocess
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = "/home/yyq/OpenFOAM/yyq-v2206/run/meshRobustness_ESM2"
RAWO = f"{ROOT}/raw_overlay_v2"

plt.rcParams["font.family"] = "sans-serif"
plt.rcParams["font.sans-serif"] = ["Arial","Helvetica","Liberation Sans","DejaVu Sans"]
plt.rcParams["font.size"] = 8

cases = {
    "coarse": {
        "label": "256×24×24",
        "name": "Coarse",
        "dir": f"{RAWO}/coarse",
        "style": dict(color="black", linestyle="--", linewidth=1.8),
    },
    "baseline": {
        "label": "320×30×30",
        "name": "Baseline",
        "dir": f"{RAWO}/baseline",
        "style": dict(color="#1f77b4", linestyle="-", linewidth=1.8),
    },
    "fine": {
        "label": "384×36×36",
        "name": "Fine",
        "dir": f"{RAWO}/fine",
        "style": dict(color="#d62728", linestyle="-.", linewidth=1.8),
    },
}

times_focus = [0.00052000, 0.00058000, 0.00062000]

def full_last_time(case_dir):
    txt = subprocess.check_output(
        ["bash","-lc", f"source /usr/lib/openfoam/openfoam2206/etc/bashrc && foamListTimes -case {case_dir} -noZero"],
        text=True
    )
    vals = []
    for s in txt.split():
        try:
            vals.append(float(s))
        except:
            pass
    return max(vals) if vals else float("nan")

def read_csv_numeric(path):
    if (not os.path.exists(path)) or os.path.getsize(path) == 0:
        return None, None
    with open(path, newline="") as f:
        lines = [ln for ln in f.read().splitlines() if ln.strip()]
    if not lines:
        return None, None
    rr = list(csv.reader(lines))
    if not rr or not rr[0]:
        return None, None
    header = rr[0]
    cols = {h: [] for h in header}
    for row in rr[1:]:
        if len(row) < len(header):
            row = row + [""]*(len(header)-len(row))
        for h, v in zip(header, row):
            try:
                cols[h].append(float(v))
            except:
                cols[h].append(float("nan"))
    return header, {k: np.array(v, dtype=float) for k,v in cols.items()}

def pick_key(fieldnames, keys):
    fmap = {f: f.strip().lower() for f in fieldnames}
    for key in keys:
        kl = key.lower()
        for f, fl in fmap.items():
            if fl == kl or kl in fl:
                return f
    raise RuntimeError(f"Cannot find {keys} in {fieldnames}")

def ftime(path):
    m = re.search(r't([0-9]+\.[0-9]+)\.csv$', os.path.basename(path))
    if not m:
        raise RuntimeError(path)
    return float(m.group(1))

def get_contour_points(path):
    header, data = read_csv_numeric(path)
    if header is None:
        return None
    xk = pick_key(header, ["Points:0", "Points_0", "x"])
    yk = pick_key(header, ["Points:1", "Points_1", "y"])
    x = data[xk]
    y = data[yk]
    if len(x) == 0 or np.all(np.isnan(x)):
        return None
    return x*1000.0, y*1000.0

def nearest_existing_time(case_key, ttarget):
    files = sorted(glob.glob(f"{cases[case_key]['dir']}/contour_center/T1200_z010_t*.csv"), key=ftime)
    vals = [ftime(f) for f in files]
    return min(vals, key=lambda v: abs(v-ttarget))

def centerline_profile(case_key, ttarget):
    t = nearest_existing_time(case_key, ttarget)
    path = f"{cases[case_key]['dir']}/centerline/centerline_t{t:.8f}.csv"
    header, data = read_csv_numeric(path)
    xk = pick_key(header, ["Points:0", "Points_0", "x"])
    pk = pick_key(header, ["p"])
    x = data[xk]*1000.0
    dp = (data[pk] - 101325.0)/1000.0
    return x, dp

# ---- metrics for S14 and Table S5 ----
rows = []
for ckey, meta in cases.items():
    croot = meta["dir"]
    contour_files = sorted(glob.glob(f"{croot}/contour_center/T1200_z010_t*.csv"), key=ftime)
    center_map = {ftime(p): p for p in glob.glob(f"{croot}/centerline/centerline_t*.csv")}

    for cf in contour_files:
        t = ftime(cf)
        if t > 0.00062000 + 1e-12:
            continue

        pts = get_contour_points(cf)
        if pts is None:
            continue
        x_mm, y_mm = pts
        x = x_mm/1000.0
        y = y_mm/1000.0

        y0 = 0.010
        center_eps = 6.0e-4
        branch_eps = 8.0e-4

        center_mask = np.abs(y - y0) <= center_eps
        if np.any(center_mask):
            x_center = np.nanmax(x[center_mask])
        else:
            idx = np.argsort(np.abs(y - y0))[:10]
            x_center = np.nanmax(x[idx])

        x_tip = np.nanmax(x)
        upper_mask = y >= y0 + branch_eps
        lower_mask = y <= y0 - branch_eps
        if np.any(upper_mask):
            x_tip = max(x_tip, np.nanmax(x[upper_mask]))
        if np.any(lower_mask):
            x_tip = max(x_tip, np.nanmax(x[lower_mask]))

        Db = x_tip - x_center

        lf = center_map.get(t, None)
        if lf is None:
            continue

        header_l, data_l = read_csv_numeric(lf)
        xcl = pick_key(header_l, ["Points:0", "Points_0", "x"])
        pcol = pick_key(header_l, ["p"])
        xc = data_l[xcl]
        pc = data_l[pcol]

        ahead_mask = (xc >= x_center) & (xc <= x_center + 0.010)
        p_ahead_max = np.nanmax(pc[ahead_mask]) if np.any(ahead_mask) else np.nan
        dp_ahead = (p_ahead_max - 101325.0)/1000.0 if np.isfinite(p_ahead_max) else np.nan

        rows.append({
            "case_key": ckey,
            "time_s": t,
            "time_ms": t*1e3,
            "Db_mm": Db*1e3,
            "dp_ahead_kPa": dp_ahead,
        })

rows = sorted(rows, key=lambda r: (r["case_key"], r["time_s"]))

def rows_of_case(ck):
    return [r for r in rows if r["case_key"] == ck]

def nearest_val(rr, ttarget, key):
    return min(rr, key=lambda r: abs(r["time_s"]-ttarget))[key]

def onset(rr, key, thr):
    for r in rr:
        if np.isfinite(r[key]) and r[key] >= thr:
            return r["time_s"]
    return float("nan")

# ---- Fig S12: center-plane T1200 contour overlays ----
fig, axs = plt.subplots(1, 3, figsize=(174/25.4, 2.65))
letters = ["(a)", "(b)", "(c)"]
for j, ttarget in enumerate(times_focus):
    ax = axs[j]
    for ckey in ["coarse", "baseline", "fine"]:
        t = nearest_existing_time(ckey, ttarget)
        path = f"{cases[ckey]['dir']}/contour_center/T1200_z010_t{t:.8f}.csv"
        pts = get_contour_points(path)
        if pts is None:
            continue
        x_mm, y_mm = pts
        ax.scatter(x_mm, y_mm, s=2.0, label=cases[ckey]["label"], color=cases[ckey]["style"]["color"])
    ax.set_xlim(8, 28)
    ax.set_ylim(4, 16)
    ax.set_xlabel("x (mm)")
    if j == 0:
        ax.set_ylabel("y (mm)")
    else:
        ax.set_yticklabels([])
    ax.set_title(f"{ttarget*1e3:.2f} ms", fontsize=9)
    ax.text(0.02, 0.96, letters[j], transform=ax.transAxes, ha="left", va="top", fontsize=8, fontweight="bold")
    for sp in ax.spines.values():
        sp.set_linewidth(0.6)
axs[0].legend(frameon=False, fontsize=7, loc="lower right")
plt.subplots_adjust(left=0.08, right=0.99, bottom=0.18, top=0.88, wspace=0.12)
fig.savefig(f"{ROOT}/plots/Fig_S12.tif", dpi=600, pil_kwargs={"compression":"tiff_lzw"})
fig.savefig(f"{ROOT}/plots/Fig_S12.png", dpi=300)
plt.close(fig)

# ---- Fig S13: centerline overpressure profiles ----
fig, axs = plt.subplots(1, 3, figsize=(174/25.4, 2.65))
for j, ttarget in enumerate(times_focus):
    ax = axs[j]
    for ckey in ["coarse", "baseline", "fine"]:
        x_mm, dp = centerline_profile(ckey, ttarget)
        mask = (x_mm >= 0.0) & (x_mm <= 35.0)
        ax.plot(x_mm[mask], dp[mask], label=cases[ckey]["label"], **cases[ckey]["style"])
    ax.set_xlim(0, 35)
    ax.set_xlabel("x (mm)")
    if j == 0:
        ax.set_ylabel(r"$\Delta p$ (kPa)")
    else:
        ax.set_yticklabels([])
    ax.set_title(f"{ttarget*1e3:.2f} ms", fontsize=9)
    ax.text(0.02, 0.96, letters[j], transform=ax.transAxes, ha="left", va="top", fontsize=8, fontweight="bold")
    for sp in ax.spines.values():
        sp.set_linewidth(0.6)
axs[0].legend(frameon=False, fontsize=7, loc="best")
plt.subplots_adjust(left=0.08, right=0.99, bottom=0.18, top=0.88, wspace=0.12)
fig.savefig(f"{ROOT}/plots/Fig_S13.tif", dpi=600, pil_kwargs={"compression":"tiff_lzw"})
fig.savefig(f"{ROOT}/plots/Fig_S13.png", dpi=300)
plt.close(fig)

# ---- Fig S14: metrics to 0.62 ms ----
fig, axs = plt.subplots(1, 3, figsize=(174/25.4, 2.8))
styles_line = {
    "coarse":   dict(color="black",   linestyle="--", linewidth=1.5),
    "baseline": dict(color="#1f77b4", linestyle="-",  linewidth=1.7),
    "fine":     dict(color="#d62728", linestyle="-.", linewidth=1.5),
}
for ckey in ["coarse","baseline","fine"]:
    rr = rows_of_case(ckey)
    tms = [r["time_ms"] for r in rr]
    dbm = [r["Db_mm"] for r in rr]
    dpk = [r["dp_ahead_kPa"] for r in rr]
    axs[0].plot(tms, dbm, label=cases[ckey]["label"], **styles_line[ckey])
    axs[1].plot(tms, dpk, label=cases[ckey]["label"], **styles_line[ckey])
    axs[2].plot(dbm, dpk, label=cases[ckey]["label"], **styles_line[ckey])
axs[0].set_xlabel("t (ms)")
axs[0].set_ylabel(r"$D_b$ (mm)")
axs[0].set_xlim(0.0, 0.62)
axs[1].set_xlabel("t (ms)")
axs[1].set_ylabel(r"$\Delta p_{ahead}$ (kPa)")
axs[1].set_xlim(0.0, 0.62)
axs[2].set_xlabel(r"$D_b$ (mm)")
axs[2].set_ylabel(r"$\Delta p_{ahead}$ (kPa)")
for i, ax in enumerate(axs):
    ax.text(0.02, 0.96, f"({chr(ord('a')+i)})", transform=ax.transAxes, ha="left", va="top", fontsize=8, fontweight="bold")
axs[0].legend(frameon=False, fontsize=7, loc="best")
plt.subplots_adjust(left=0.08, right=0.99, bottom=0.20, top=0.94, wspace=0.35)
fig.savefig(f"{ROOT}/plots/Fig_S14.eps")
fig.savefig(f"{ROOT}/plots/Fig_S14.tif", dpi=600, pil_kwargs={"compression":"tiff_lzw"})
fig.savefig(f"{ROOT}/plots/Fig_S14.png", dpi=300)
plt.close(fig)

# ---- Fig S15: spanwise contour overlays at 0.62 ms ----
fig, axs = plt.subplots(1, 3, figsize=(174/25.4, 2.65))
zlist = ["0.003", "0.010", "0.017"]
for j, z in enumerate(zlist):
    ax = axs[j]
    for ckey in ["coarse", "baseline", "fine"]:
        path = f"{cases[ckey]['dir']}/contour_span/T1200_z{z}_t0.00062000.csv"
        pts = get_contour_points(path)
        if pts is None:
            continue
        x_mm, y_mm = pts
        ax.scatter(x_mm, y_mm, s=2.0, label=cases[ckey]["label"], color=cases[ckey]["style"]["color"])
    ax.set_xlim(8, 28)
    ax.set_ylim(4, 16)
    ax.set_xlabel("x (mm)")
    if j == 0:
        ax.set_ylabel("y (mm)")
    else:
        ax.set_yticklabels([])
    ax.set_title(f"z = {z} m", fontsize=9)
    ax.text(0.02, 0.96, letters[j], transform=ax.transAxes, ha="left", va="top", fontsize=8, fontweight="bold")
    for sp in ax.spines.values():
        sp.set_linewidth(0.6)
axs[0].legend(frameon=False, fontsize=7, loc="lower right")
plt.subplots_adjust(left=0.08, right=0.99, bottom=0.18, top=0.88, wspace=0.12)
fig.savefig(f"{ROOT}/plots/Fig_S15.tif", dpi=600, pil_kwargs={"compression":"tiff_lzw"})
fig.savefig(f"{ROOT}/plots/Fig_S15.png", dpi=300)
plt.close(fig)

# ---- Table S5 sampled at 0.52 / 0.58 / 0.62 ms ----
case_dirs = {
    "coarse": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening_m256x24x24",
    "baseline": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening",
    "fine": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening_m384x36x36",
}
table_rows = []
for ckey in ["coarse","baseline","fine"]:
    rr = rows_of_case(ckey)
    lst = full_last_time(case_dirs[ckey])
    table_rows.append({
        "Mesh": cases[ckey]["label"],
        "Last stable output (s)": "" if math.isnan(lst) else f"{lst:.8f}",
        "Db @ 0.52 ms (mm)": f"{nearest_val(rr,0.00052,'Db_mm'):.4f}",
        "Db @ 0.58 ms (mm)": f"{nearest_val(rr,0.00058,'Db_mm'):.4f}",
        "Db @ 0.62 ms (mm)": f"{nearest_val(rr,0.00062,'Db_mm'):.4f}",
        "dpAhead @ 0.52 ms (kPa)": f"{nearest_val(rr,0.00052,'dp_ahead_kPa'):.4f}",
        "dpAhead @ 0.58 ms (kPa)": f"{nearest_val(rr,0.00058,'dp_ahead_kPa'):.4f}",
        "dpAhead @ 0.62 ms (kPa)": f"{nearest_val(rr,0.00062,'dp_ahead_kPa'):.4f}",
        "t(dpAhead >= 3 kPa) (s)": "" if math.isnan(onset(rr,'dp_ahead_kPa',3.0)) else f"{onset(rr,'dp_ahead_kPa',3.0):.8f}",
    })

with open(f"{ROOT}/tables/Table_S5.csv", "w", newline="") as f:
    header = list(table_rows[0].keys())
    w = csv.DictWriter(f, fieldnames=header)
    w.writeheader()
    w.writerows(table_rows)
