#!/usr/bin/env python3
from PIL import Image
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable

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"] = 8

CASE_ROWS = [
    ("coarse",   "Coarse\n256x24x24"),
    ("baseline", "Baseline\n320x30x30"),
    ("fine",     "Fine\n384x36x36"),
]

TIMES = ["0.40 ms", "0.50 ms", "0.60 ms"]
TKEYS = ["0.00040000", "0.00050000", "0.00060000"]

ZCOLS = [("0.003", "z = 0.003 m"), ("0.010", "z = 0.010 m"), ("0.017", "z = 0.017 m")]

letters = [f"({chr(ord('a')+i)})" for i in range(9)]

# local flame-head window
XMAX_MM = 30.0
YMAX_MM = 20.0
FULL_X_MM = 240.0

W_IN = 174/25.4  # 174 mm, as in guideline for large-size journals

def load_rgb(path):
    return np.asarray(Image.open(path).convert("RGB"))

def nonwhite_bbox(arr, thr=248):
    mask = np.any(arr < thr, axis=2)
    ys, xs = np.where(mask)
    if len(xs) == 0:
        return 0, 0, arr.shape[1], arr.shape[0]
    x0 = max(0, int(xs.min()) - 4)
    x1 = min(arr.shape[1], int(xs.max()) + 5)
    y0 = max(0, int(ys.min()) - 4)
    y1 = min(arr.shape[0], int(ys.max()) + 5)
    return x0, y0, x1, y1

def crop_local_panel(path, xmax_mm=XMAX_MM):
    arr = load_rgb(path)
    x0, y0, x1, y1 = nonwhite_bbox(arr)
    domain = arr[y0:y1, x0:x1, :]
    h, w, _ = domain.shape

    frac = xmax_mm / FULL_X_MM
    xcut = max(20, min(w, int(round(w * frac))))
    local = domain[:, :xcut, :]

    # trim again after the left-window crop
    xx0, yy0, xx1, yy1 = nonwhite_bbox(local)
    local = local[yy0:yy1, xx0:xx1, :]
    return local

def assemble_3x3(in_pattern, out_tif, out_png, col_labels, cbar_label, cmap_name, vmin, vmax):
    fig, axs = plt.subplots(3, 3, figsize=(W_IN, 5.6))
    k = 0

    for i, (ckey, row_label) in enumerate(CASE_ROWS):
        for j, col_label in enumerate(col_labels):
            ax = axs[i, j]
            img = crop_local_panel(in_pattern(ckey, j))

            ax.imshow(img, extent=[0, XMAX_MM, 0, YMAX_MM], origin='lower', aspect='auto')
            ax.set_xlim(0, XMAX_MM)
            ax.set_ylim(0, YMAX_MM)

            ax.set_xticks([0, 10, 20, 30])
            ax.set_yticks([0, 10, 20])

            ax.tick_params(axis='both', labelsize=7, length=2, pad=1)

            if i < 2:
                ax.set_xticklabels([])
            else:
                ax.set_xlabel("x (mm)", fontsize=8)

            if j == 0:
                ax.set_ylabel("y (mm)", fontsize=8)
                ax.text(-0.34, 0.5, row_label, rotation=90, transform=ax.transAxes,
                        ha="center", va="center", fontsize=8)
            else:
                ax.set_yticklabels([])

            if i == 0:
                ax.set_title(col_label, fontsize=9, pad=4)

            ax.text(0.02, 0.96, letters[k], transform=ax.transAxes,
                    ha="left", va="top", fontsize=8, fontweight="bold", color="black")

            for sp in ax.spines.values():
                sp.set_linewidth(0.6)

            k += 1

    cax = fig.add_axes([0.92, 0.14, 0.015, 0.72])
    sm = ScalarMappable(norm=Normalize(vmin=vmin, vmax=vmax), cmap=cmap_name)
    cb = fig.colorbar(sm, cax=cax)
    cb.set_label(cbar_label, fontsize=8)
    cb.ax.tick_params(labelsize=7)

    plt.subplots_adjust(left=0.10, right=0.90, bottom=0.10, top=0.92, wspace=0.08, hspace=0.14)
    fig.savefig(out_tif, dpi=600, pil_kwargs={"compression": "tiff_lzw"})
    fig.savefig(out_png, dpi=300)
    plt.close(fig)

# --- Fig. S12: center-plane temperature slices ---
assemble_3x3(
    in_pattern=lambda c, j: f"{ROOT}/raw_slices/T_center/{c}_t{TKEYS[j]}.tif",
    out_tif=f"{ROOT}/plots/Fig_S12.tif",
    out_png=f"{ROOT}/plots/Fig_S12.png",
    col_labels=TIMES,
    cbar_label="T (K)",
    cmap_name="viridis",
    vmin=293.0,
    vmax=2800.0
)

# --- Fig. S13: center-plane Ux slices ---
assemble_3x3(
    in_pattern=lambda c, j: f"{ROOT}/raw_slices/Ux_center/{c}_t{TKEYS[j]}.tif",
    out_tif=f"{ROOT}/plots/Fig_S13.tif",
    out_png=f"{ROOT}/plots/Fig_S13.png",
    col_labels=TIMES,
    cbar_label=r"$U_x$ (m s$^{-1}$)",
    cmap_name="coolwarm",
    vmin=-40.0,
    vmax=120.0
)

# --- Fig. S15: spanwise temperature slices at 0.60 ms ---
assemble_3x3(
    in_pattern=lambda c, j: f"{ROOT}/raw_slices/T_span/{c}_z{ZCOLS[j][0]}_t0.00060000.tif",
    out_tif=f"{ROOT}/plots/Fig_S15.tif",
    out_png=f"{ROOT}/plots/Fig_S15.png",
    col_labels=[z[1] for z in ZCOLS],
    cbar_label="T (K)",
    cmap_name="viridis",
    vmin=293.0,
    vmax=2800.0
)
