#!/usr/bin/env pvpython
from paraview.simple import *
import os, csv

ROOT = "/home/yyq/OpenFOAM/yyq-v2206/run/meshRobustness_ESM2"
MATCH = f"{ROOT}/tables/matched_times.csv"

cases = {
    "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",
}

matched = {}
with open(MATCH, newline="") as f:
    for r in csv.DictReader(f):
        matched[(r["case_key"], r["target_s"])] = float(r["actual_s"])

def apply_preset_safe(lut, names):
    for n in names:
        try:
            lut.ApplyPreset(n, True)
            return
        except:
            pass

def export_panel(case_key, time_value, zplane, field_mode, out_file):
    case_dir = cases[case_key]
    foam = os.path.join(case_dir, os.path.basename(case_dir) + ".foam")

    ResetSession()
    reader = OpenDataFile(foam)
    try:
        reader.MeshRegions = ['internalMesh']
    except:
        pass
    try:
        reader.CellArrays = ['T','p','Qdot','kappa','U']
    except:
        pass

    view = CreateView("RenderView")
    view.ViewSize = [2400, 400]
    view.Background = [1.0, 1.0, 1.0]
    view.OrientationAxesVisibility = 0
    view.InteractionMode = '2D'

    sl = Slice(Input=reader)
    sl.SliceType = 'Plane'
    sl.SliceType.Origin = [0.06, 0.01, zplane]
    sl.SliceType.Normal = [0, 0, 1]

    c2p = CellDatatoPointData(Input=sl)

    disp = Show(c2p, view)
    disp.Representation = 'Surface'

    if field_mode == "T":
        ColorBy(disp, ('POINTS','T'))
        lut = GetColorTransferFunction('T')
        apply_preset_safe(lut, ['Viridis (matplotlib)', 'Viridis'])
        lut.RescaleTransferFunction(293.0, 2800.0)
        pwf = GetOpacityTransferFunction('T')
        pwf.RescaleTransferFunction(293.0, 2800.0)
    elif field_mode == "Ux":
        ColorBy(disp, ('POINTS','U','X'))
        lut = GetColorTransferFunction('U')
        apply_preset_safe(lut, ['Cool to Warm', 'Blue to Red Rainbow'])
        lut.RescaleTransferFunction(-40.0, 120.0)
        pwf = GetOpacityTransferFunction('U')
        pwf.RescaleTransferFunction(-40.0, 120.0)
    else:
        raise RuntimeError("field_mode must be T or Ux")

    disp.SetScalarBarVisibility(view, False)

    ctr = Contour(Input=c2p)
    ctr.ContourBy = ['POINTS','T']
    ctr.Isosurfaces = [1200.0]
    cdisp = Show(ctr, view)
    cdisp.Representation = 'Surface'
    cdisp.ColorArrayName = [None, '']
    cdisp.DiffuseColor = [0.0, 0.0, 0.0]
    cdisp.AmbientColor = [0.0, 0.0, 0.0]
    cdisp.LineWidth = 2.0
    cdisp.SetScalarBarVisibility(view, False)

    UpdatePipeline(time=time_value, proxy=reader)
    UpdatePipeline(time=time_value, proxy=sl)
    UpdatePipeline(time=time_value, proxy=c2p)
    UpdatePipeline(time=time_value, proxy=ctr)

    cam = GetActiveCamera()
    cam.ParallelProjection = 1
    cam.SetFocalPoint(0.06, 0.01, 0.0)
    cam.SetPosition(0.06, 0.01, 1.0)
    cam.SetViewUp(0, 1, 0)
    view.CameraParallelScale = 0.01

    os.makedirs(os.path.dirname(out_file), exist_ok=True)
    SaveScreenshot(out_file, view, ImageResolution=[2400,400])

# Fig. S12: center-plane T at 0.40 / 0.50 / 0.60 ms
targets = ["0.00040000","0.00050000","0.00060000"]
for case_key in ["coarse","baseline","fine"]:
    for tgt in targets:
        export_panel(
            case_key=case_key,
            time_value=matched[(case_key, tgt)],
            zplane=0.010,
            field_mode="T",
            out_file=f"{ROOT}/raw_slices/T_center/{case_key}_t{tgt}.tif"
        )

# Fig. S13: center-plane Ux at 0.40 / 0.50 / 0.60 ms
for case_key in ["coarse","baseline","fine"]:
    for tgt in targets:
        export_panel(
            case_key=case_key,
            time_value=matched[(case_key, tgt)],
            zplane=0.010,
            field_mode="Ux",
            out_file=f"{ROOT}/raw_slices/Ux_center/{case_key}_t{tgt}.tif"
        )

# Fig. S15: T at t=0.60 ms on z = 0.003 / 0.010 / 0.017
zlist = [0.003, 0.010, 0.017]
zlab  = ["0.003","0.010","0.017"]
for case_key in ["coarse","baseline","fine"]:
    for z, zl in zip(zlist, zlab):
        export_panel(
            case_key=case_key,
            time_value=matched[(case_key, "0.00060000")],
            zplane=z,
            field_mode="T",
            out_file=f"{ROOT}/raw_slices/T_span/{case_key}_z{zl}_t0.00060000.tif"
        )
