#!/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"] = 9

CASE_ROWS = [
    ("coarse",   "Coarse\n256×24×24"),
    ("baseline", "Baseline\n320×30×30"),
    ("fine",     "Fine\n384×36×36"),
]
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)]

W_IN = 174/25.4  # 174 mm

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

def assemble_3x3(in_pattern, out_tif, out_png, col_labels, row_defs, cmap, vmin, vmax, cbar_label):
    fig, axs = plt.subplots(3, 3, figsize=(W_IN, 4.9))
    k = 0
    for i, (ckey, rlab) in enumerate(row_defs):
        for j, col in enumerate(col_labels):
            ax = axs[i, j]
            img = load_img(in_pattern(ckey, j))
            ax.imshow(img)
            ax.set_axis_off()
            ax.text(0.015, 0.96, letters[k], transform=ax.transAxes,
                    ha="left", va="top", fontsize=9, fontweight="bold", color="black")
            if i == 0:
                ax.set_title(col, fontsize=9, pad=5)
            if j == 0:
                ax.text(-0.11, 0.5, rlab, rotation=90, transform=ax.transAxes,
                        ha="center", va="center", fontsize=9)
            k += 1

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

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

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,
    row_defs=CASE_ROWS,
    cmap="viridis",
    vmin=293.0,
    vmax=2800.0,
    cbar_label="T (K)"
)

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,
    row_defs=CASE_ROWS,
    cmap="coolwarm",
    vmin=-40.0,
    vmax=120.0,
    cbar_label=r"$U_x$ (m s$^{-1}$)"
)

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],
    row_defs=CASE_ROWS,
    cmap="viridis",
    vmin=293.0,
    vmax=2800.0,
    cbar_label="T (K)"
)
