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

ROOT = "/home/yyq/OpenFOAM/yyq-v2206/run/meshRobustness_ESM2"
plt.rcParams["font.family"] = "sans-serif"
plt.rcParams["font.sans-serif"] = ["Arial","Helvetica","Liberation Sans","DejaVu Sans"]
plt.rcParams["font.size"] = 9

cases = {
    "coarse": {
        "case_dir": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening_m256x24x24",
        "csv_root": f"{ROOT}/raw_csv/coarse",
        "mesh": "256×24×24",
        "label": "Coarse",
    },
    "baseline": {
        "case_dir": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening",
        "csv_root": f"{ROOT}/raw_csv/baseline",
        "mesh": "320×30×30",
        "label": "Baseline",
    },
    "fine": {
        "case_dir": "/home/yyq/OpenFOAM/yyq-v2206/run/S10_phi1p0_screening_m384x36x36",
        "csv_root": f"{ROOT}/raw_csv/fine",
        "mesh": "384×36×36",
        "label": "Fine",
    },
}

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 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 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

    reader = csv.reader(lines)
    rows = list(reader)
    if not rows:
        return None, None

    header = rows[0]
    if not header or all((not h.strip()) for h in header):
        return None, None

    cols = {h: [] for h in header}
    for row in rows[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 columns {keys} in {fieldnames}")

rows = []

for ckey, meta in cases.items():
    croot = meta["csv_root"]
    contour_files = sorted(glob.glob(f"{croot}/contour/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.00060000 + 1e-12:
            continue

        fieldnames_c, data_c = read_csv_numeric(cf)
        if fieldnames_c is None:
            continue

        xcol = pick_key(fieldnames_c, ["Points:0", "Points_0", "x"])
        ycol = pick_key(fieldnames_c, ["Points:1", "Points_1", "y"])

        x = data_c[xcol]
        y = data_c[ycol]

        if len(x) == 0 or np.all(np.isnan(x)):
            continue

        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

        fieldnames_l, data_l = read_csv_numeric(lf)
        if fieldnames_l is None:
            continue

        xcl = pick_key(fieldnames_l, ["Points:0", "Points_0", "x"])
        pcol = pick_key(fieldnames_l, ["p"])
        try:
            ucol = pick_key(fieldnames_l, ["U:0", "U_0", "ux"])
        except:
            ucol = pick_key(fieldnames_l, ["u"])

        xc = data_l[xcl]
        pc = data_l[pcol]
        uc = data_l[ucol]

        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

        wake_mask = (xc >= max(0.0, x_center - 0.020)) & (xc <= x_center)
        u_back = np.nanmin(uc[wake_mask]) if np.any(wake_mask) else np.nan

        rows.append({
            "case_key": ckey,
            "case_label": meta["label"],
            "mesh": meta["mesh"],
            "time_s": t,
            "time_ms": t * 1e3,
            "x_tip_m": x_tip,
            "x_center_m": x_center,
            "Db_m": Db,
            "Db_mm": Db * 1e3,
            "dp_ahead_kPa": dp_ahead,
            "U_back_ms": u_back,
        })

rows = sorted(rows, key=lambda r: (r["case_key"], r["time_s"]))
if not rows:
    raise RuntimeError("No valid contour rows were parsed. Check contour CSV contents.")

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

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

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

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

tab_rows = []
for ckey, meta in cases.items():
    rr = rows_of_case(ckey)
    lst = full_last_time(meta["case_dir"])
    tab_rows.append({
        "Mesh": meta["mesh"],
        "LastStableOutput_s": "" if math.isnan(lst) else f"{lst:.8f}",
        "Db_0p40ms_mm": f"{near_val(rr,0.0004,'Db_mm'):.4f}",
        "Db_0p50ms_mm": f"{near_val(rr,0.0005,'Db_mm'):.4f}",
        "Db_0p60ms_mm": f"{near_val(rr,0.0006,'Db_mm'):.4f}",
        "dpAhead_0p40ms_kPa": f"{near_val(rr,0.0004,'dp_ahead_kPa'):.4f}",
        "dpAhead_0p50ms_kPa": f"{near_val(rr,0.0005,'dp_ahead_kPa'):.4f}",
        "dpAhead_0p60ms_kPa": f"{near_val(rr,0.0006,'dp_ahead_kPa'):.4f}",
        "t_Db_ge_5mm_s": "" if math.isnan(onset(rr,'Db_mm',5.0)) else f"{onset(rr,'Db_mm',5.0):.8f}",
        "t_dpAhead_ge_3kPa_s": "" if math.isnan(onset(rr,'dp_ahead_kPa',3.0)) else f"{onset(rr,'dp_ahead_kPa',3.0):.8f}",
    })

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

fig, axs = plt.subplots(1, 3, figsize=(174/25.4, 2.8))
styles = {
    "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]["mesh"], **styles[ckey])
    axs[1].plot(tms, dpk, label=cases[ckey]["mesh"], **styles[ckey])
    axs[2].plot(dbm, dpk, label=cases[ckey]["mesh"], **styles[ckey])

axs[0].set_xlabel("t (ms)")
axs[0].set_ylabel(r"$D_b$ (mm)")
axs[0].set_xlim(0.0, 0.60)

axs[1].set_xlabel("t (ms)")
axs[1].set_ylabel(r"$\Delta p_{ahead}$ (kPa)")
axs[1].set_xlim(0.0, 0.60)

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=9, fontweight="bold")

axs[0].legend(frameon=False, fontsize=8, 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)
