from pathlib import Path
import json

import jpype
import numpy as np
from mph.model import Model as MphModel

root = Path(r"C:\Users\jiwei\Documents\GitHub\sim-proj\dev-docs\comsol-hbm-hbn-repro")
model_dir = root / "model"
output_dir = root / "output"
model_dir.mkdir(parents=True, exist_ok=True)
output_dir.mkdir(parents=True, exist_ok=True)

case_slug = "arxiv_2510_11461_hbm_hbn_repro"
case_selector = root / "input" / "current-interposer.txt"
interposer = case_selector.read_text(encoding="utf-8").strip().lower() if case_selector.exists() else "hbn"
if interposer not in {"hbn", "si"}:
    raise ValueError(f"unsupported interposer selector {interposer!r}; expected hbn or si")

thickness_selector = root / "input" / "current-interposer-thickness-um.txt"
power_selector = root / "input" / "current-gpu-power-w.txt"
interposer_t_um = float(thickness_selector.read_text(encoding="utf-8").strip()) if thickness_selector.exists() else 300.0
gpu_power_w = float(power_selector.read_text(encoding="utf-8").strip()) if power_selector.exists() else 100.0

case_id = f"{interposer}_p{gpu_power_w:g}w_t{interposer_t_um:g}um"
checkpoint = model_dir / f"{case_slug}_02_{case_id}.mph"

try:
    model.component().remove("comp1")
except Exception:
    pass
for collection_name in ["study", "sol"]:
    collection = getattr(model, collection_name)()
    try:
        existing = [str(tag) for tag in list(collection.tags())]
    except Exception:
        existing = []
    for tag in existing:
        try:
            collection.remove(tag)
        except Exception:
            pass

model.label(f"{case_slug}_02_baseline_{interposer}.mph")
try:
    model.title("GPU-HBM h-BN interposer reproduction baseline")
except Exception:
    pass
model.modelPath(str(model_dir))

for stale in ["paper_id", "work_root", "assumption_level"]:
    try:
        model.param().remove(stale)
    except Exception:
        pass

# Paper-grounded parameters plus representative dimensions needed for reproduction.
params = {
    "T_amb": ("293.15[K]", "Ambient temperature for convective boundaries"),
    "h_top": ("250[W/(m^2*K)]", "Forced convection coefficient; paper range is 150-350 W/(m^2*K)"),
    "h_bottom": ("10[W/(m^2*K)]", "Natural convection coefficient at substrate bottom"),
    "gpu_power": (f"{gpu_power_w:g}[W]", "Baseline TDP / heat input anchor from the paper sweep"),
    "pkg_lx": ("20[mm]", "Assumed package footprint x"),
    "pkg_ly": ("20[mm]", "Assumed package footprint y"),
    "substrate_t": ("0.5[mm]", "Assumed silicon substrate thickness"),
    "interposer_t": (f"{interposer_t_um:g}[um]", "h-BN interposer thickness near reported saturation"),
    "gpu_lx": ("10[mm]", "Assumed GPU die footprint x"),
    "gpu_ly": ("10[mm]", "Assumed GPU die footprint y"),
    "gpu_t": ("150[um]", "Assumed GPU active layer thickness"),
    "lid_t": ("0.2[mm]", "Assumed heat spreader thickness"),
}
for tag, (value, desc) in params.items():
    model.param().set(tag, value)
    model.param().descr(tag, desc)

comp = model.component().create("comp1", True)
geom = comp.geom().create("geom1", 3)
geom.lengthUnit("mm")

def block(tag, label, size, pos):
    feat = geom.create(tag, "Block")
    feat.label(label)
    feat.set("size", jpype.JArray(jpype.JString)(size))
    feat.set("pos", jpype.JArray(jpype.JString)(pos))
    feat.set("selresult", "on")
    feat.set("selresultshow", "dom")
    return feat

block("blk_sub", "Silicon substrate", ["pkg_lx", "pkg_ly", "substrate_t"], ["-pkg_lx/2", "-pkg_ly/2", "0"])
block("blk_int", "h-BN interposer", ["pkg_lx", "pkg_ly", "interposer_t"], ["-pkg_lx/2", "-pkg_ly/2", "substrate_t"])
block(
    "blk_gpu",
    "GPU heat source die",
    ["gpu_lx", "gpu_ly", "gpu_t"],
    ["-gpu_lx/2", "-gpu_ly/2", "substrate_t+interposer_t"],
)
block(
    "blk_lid",
    "Copper heat spreader",
    ["pkg_lx", "pkg_ly", "lid_t"],
    ["-pkg_lx/2", "-pkg_ly/2", "substrate_t+interposer_t+gpu_t"],
)

geom.run()

def make_boundary_box(tag, label, zmin, zmax):
    sel = comp.selection().create(tag, "Box")
    sel.label(label)
    sel.geom("geom1", 2)
    sel.set("entitydim", "2")
    sel.set("xmin", "-pkg_lx/2-1[mm]")
    sel.set("xmax", "pkg_lx/2+1[mm]")
    sel.set("ymin", "-pkg_ly/2-1[mm]")
    sel.set("ymax", "pkg_ly/2+1[mm]")
    sel.set("zmin", zmin)
    sel.set("zmax", zmax)
    sel.set("condition", "inside")
    return sel

make_boundary_box(
    "sel_top",
    "Top forced convection face",
    "substrate_t+interposer_t+gpu_t+lid_t-1[um]",
    "substrate_t+interposer_t+gpu_t+lid_t+1[um]",
)
make_boundary_box("sel_bottom", "Bottom natural convection face", "-1[um]", "1[um]")

def material(tag, label, selection, k, rho, cp):
    mat = comp.material().create(tag, "Common")
    mat.label(label)
    mat.propertyGroup("def").set("thermalconductivity", jpype.JArray(jpype.JString)(k))
    mat.propertyGroup("def").set("density", jpype.JArray(jpype.JString)([rho]))
    mat.propertyGroup("def").set("heatcapacity", jpype.JArray(jpype.JString)([cp]))
    mat.selection().named(selection)
    return mat

material("mat_sub", "Silicon substrate", "geom1_blk_sub_dom", ["140[W/(m*K)]"], "2329[kg/m^3]", "700[J/(kg*K)]")
material("mat_gpu", "Silicon GPU die", "geom1_blk_gpu_dom", ["140[W/(m*K)]"], "2329[kg/m^3]", "700[J/(kg*K)]")
material("mat_lid", "Copper heat spreader", "geom1_blk_lid_dom", ["400[W/(m*K)]"], "8960[kg/m^3]", "385[J/(kg*K)]")

if interposer == "hbn":
    material(
        "mat_int",
        "Anisotropic h-BN interposer",
        "geom1_blk_int_dom",
        [
            "751[W/(m*K)]", "0", "0",
            "0", "751[W/(m*K)]", "0",
            "0", "0", "10[W/(m*K)]",
        ],
        "2100[kg/m^3]",
        "800[J/(kg*K)]",
    )
else:
    material("mat_int", "Silicon interposer", "geom1_blk_int_dom", ["140[W/(m*K)]"], "2329[kg/m^3]", "700[J/(kg*K)]")

ht = comp.physics().create("ht", "HeatTransfer", "geom1")

hs = ht.create("hs_gpu", "HeatSource", 3)
hs.label("GPU volumetric heat generation")
hs.selection().named("geom1_blk_gpu_dom")
hs.set("Q0", "gpu_power/(gpu_lx*gpu_ly*gpu_t)")

top = ht.create("hf_top", "HeatFluxBoundary", 2)
top.label("Top forced convection")
top.selection().named("sel_top")
top.set("HeatFluxType", "ConvectiveHeatFlux")
top.set("h", "h_top")
top.set("Text", "T_amb")

bottom = ht.create("hf_bottom", "HeatFluxBoundary", 2)
bottom.label("Bottom natural convection")
bottom.selection().named("sel_bottom")
bottom.set("HeatFluxType", "ConvectiveHeatFlux")
bottom.set("h", "h_bottom")
bottom.set("Text", "T_amb")

mesh = comp.mesh().create("mesh1")
mesh.autoMeshSize(5)
mesh.run()

std = model.study().create("std1")
std.create("stat", "Stationary")
model.sol().create("sol1")
model.sol("sol1").study("std1")
model.sol("sol1").createAutoSequence("std1")
model.sol("sol1").runAll()

wrapped = MphModel(model)
temps_c = np.asarray(wrapped.evaluate("T", unit="degC"), dtype=float).ravel()
finite = temps_c[np.isfinite(temps_c)]

model.save(str(checkpoint))

result = {
    "step": "baseline_solve",
    "case_id": case_id,
    "interposer": interposer,
    "gpu_power_w": gpu_power_w,
    "interposer_t_um": interposer_t_um,
    "checkpoint": str(checkpoint),
    "mesh_size": "normal",
    "temperature_degC": {
        "min": float(np.min(finite)),
        "max": float(np.max(finite)),
        "mean": float(np.mean(finite)),
        "count": int(finite.size),
    },
    "top_boundary_entities": [int(x) for x in list(comp.selection("sel_top").entities())],
    "bottom_boundary_entities": [int(x) for x in list(comp.selection("sel_bottom").entities())],
    "status": "ok",
}

result_path = output_dir / f"{case_id}_result.json"
result_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
result["result_path"] = str(result_path)

_result = result