"""openEMS replication of the Yang & Rahmat-Samii validation cases.

Apples-to-apples with the RayRF projects built by build_and_run.py:
identical cell size (uniform cubic), identical domain extents (board +
10 mm air + 4 mm PML inside the mesh), identical geometry coordinates
(shifted so the board center is the origin), identical 9-21 GHz
Gaussian band, 50 ohm in-plane lumped port across the same 0.5 mm feed
gap, EndCriteria 1e-5 with a 1M step bound (effectively uncapped).

Cases:
  pmc_inf   - dipole 1.5 mm above a PMC boundary face (zmin = PMC)
  pec_sheet - dipole 1.5 mm above a finite 25 x 25 mm PEC sheet
  ebg       - the 7x7 mushroom array (3.0 mm patches, 0.5 mm gaps,
              1.0 mm eps_r 2.20 substrate, one 1-cell via per patch)

Output: runs_openems/<case>/port_s11.csv with the same columns as the
RayRF backend CSV (freq_hz, s11_db, zin_real, zin_imag) plus
summary.json with cells / steps / engine throughput / wall time.

Run with the dedicated venv:
  RayRF-Misc/RayRF-vs-openEMS/.venv-openems/Scripts/python.exe
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time

import numpy as np

# h5py BEFORE CSXCAD/openEMS (HDF5 DLL ordering on this box).
import h5py  # noqa: F401

from CSXCAD import ContinuousStructure  # noqa: E402
from openEMS import openEMS  # noqa: E402

CELL_DEFAULT = 0.25
BOARD = 25.0
AIR = 10.0
PML_MM = 4.0
F_MIN, F_MAX = 9e9, 21e9
R0 = 50.0

L_DIP, STRIP, GAP, H_GND = 10.0, 0.5, 0.5, 1.5
PATCH_W, PERIOD, SUB_H, SUB_EPS = 3.0, 3.5, 1.0, 2.20


class RedirectFD:
    def __init__(self, fd: int, path: str):
        self.fd, self.path = fd, path

    def __enter__(self):
        sys.stdout.flush()
        self._saved = os.dup(self.fd)
        self._f = open(self.path, "wb")
        os.dup2(self._f.fileno(), self.fd)
        return self

    def __exit__(self, *exc):
        sys.stdout.flush()
        os.dup2(self._saved, self.fd)
        os.close(self._saved)
        self._f.close()


def lines(lo: float, hi: float, cell: float) -> list:
    n = int(round((hi - lo) / cell))
    return list(np.linspace(lo, hi, n + 1))


def build(case: str, cell: float, max_ts: int, end_crit: float):
    fdtd = openEMS(NrTS=max_ts, EndCriteria=end_crit)
    fdtd.SetGaussExcite(0.5 * (F_MIN + F_MAX), 0.5 * (F_MAX - F_MIN))

    pml_cells = int(round(PML_MM / cell))
    bc = [f"PML_{pml_cells}"] * 6
    if case == "pmc_inf":
        bc[4] = "PMC"  # zmin
    fdtd.SetBoundaryCond(bc)

    csx = ContinuousStructure()
    fdtd.SetCSX(csx)
    mesh = csx.GetGrid()
    mesh.SetDeltaUnit(1e-3)

    half = BOARD / 2.0 + AIR + PML_MM  # 26.5 in plane
    if case == "pmc_inf":
        # dipole plane at z=0, PMC wall 1.5 mm below, air+PML above
        z_lo, z_dip = -H_GND, 0.0
        z_hi = AIR + PML_MM
    elif case == "freespace":
        # dipole alone, air+PML both sides
        z_lo, z_dip = -(AIR + PML_MM), 0.0
        z_hi = AIR + PML_MM
    else:
        # ground at z=0, dipole at 1.5, air+PML both sides
        z_lo, z_dip = -(AIR + PML_MM), H_GND
        z_hi = H_GND + AIR + PML_MM

    mesh.AddLine("x", lines(-half, half, cell))
    mesh.AddLine("y", lines(-half, half, cell))
    mesh.AddLine("z", lines(z_lo, z_hi, cell))
    n_cells = (len(lines(-half, half, cell)) - 1) ** 2 * \
        (len(lines(z_lo, z_hi, cell)) - 1)

    pec = csx.AddMetal("pec")

    # Dipole arms (zero-thickness strips), gap 0.5 mm centered at x=0.
    arm_in, arm_out = GAP / 2.0, L_DIP / 2.0
    yh = STRIP / 2.0
    pec.AddBox([-arm_out, -yh, z_dip], [-arm_in, +yh, z_dip], priority=10)
    pec.AddBox([+arm_in, -yh, z_dip], [+arm_out, +yh, z_dip], priority=10)

    if case == "pec_sheet":
        pec.AddBox([-BOARD / 2, -BOARD / 2, 0.0],
                   [+BOARD / 2, +BOARD / 2, 0.0], priority=10)

    if case == "ebg":
        pec.AddBox([-BOARD / 2, -BOARD / 2, 0.0],
                   [+BOARD / 2, +BOARD / 2, 0.0], priority=10)
        sub = csx.AddMaterial("substrate", epsilon=SUB_EPS)
        sub.AddBox([-BOARD / 2, -BOARD / 2, 0.0],
                   [+BOARD / 2, +BOARD / 2, SUB_H], priority=0)
        # Same half-cell array offset as the RayRF project, so each via
        # center sits on a cell center and the via is one exact cell.
        off = cell / 2.0
        centers = [off + PERIOD * (i - 3) for i in range(7)]
        vh = cell / 2.0
        for cy in centers:
            for cx in centers:
                pec.AddBox([cx - PATCH_W / 2, cy - PATCH_W / 2, SUB_H],
                           [cx + PATCH_W / 2, cy + PATCH_W / 2, SUB_H],
                           priority=10)
                pec.AddBox([cx - vh, cy - vh, 0.0],
                           [cx + vh, cy + vh, SUB_H], priority=10)

    port = fdtd.AddLumpedPort(
        1, R0,
        [-arm_in, -yh, z_dip], [+arm_in, +yh, z_dip],
        "x", excite=1.0, priority=5,
    )
    return fdtd, port, n_cells


RE_SPEED = re.compile(r"Speed:\s*([\d.]+)\s*MC/s")
RE_TS = re.compile(r"Time step.*?:\s*([0-9.eE+-]+)\s*s")
RE_ITER = re.compile(r"Time for (\d+) iterations")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("case", choices=["pmc_inf", "pec_sheet", "ebg",
                                     "freespace"])
    ap.add_argument("--cell", type=float, default=CELL_DEFAULT)
    ap.add_argument("--max-ts", type=int, default=1_000_000)
    ap.add_argument("--end-criteria", type=float, default=1e-5)
    ap.add_argument("--tag", default="")
    args = ap.parse_args()

    here = os.path.dirname(os.path.abspath(__file__))
    tag = args.case + (args.tag or "")
    sim_path = os.path.abspath(os.path.join(
        here, "..", "runs_openems", tag))
    os.makedirs(sim_path, exist_ok=True)

    fdtd, port, n_cells = build(args.case, args.cell, args.max_ts,
                                args.end_criteria)
    print(f"[{tag}] cells={n_cells/1e6:.2f} M  cell={args.cell} mm",
          flush=True)

    log_path = os.path.join(sim_path, "openems_stdout.log")
    t0 = time.perf_counter()
    with RedirectFD(1, log_path):
        fdtd.Run(sim_path, cleanup=False)
    wall = time.perf_counter() - t0

    log_text = open(log_path, "r", errors="replace").read()
    speeds = [float(m.group(1)) for m in RE_SPEED.finditer(log_text)]
    iters = [int(m.group(1)) for m in RE_ITER.finditer(log_text)]

    freqs = np.linspace(F_MIN, F_MAX, 501)
    port.CalcPort(sim_path, freqs)
    s11 = port.uf_ref / port.uf_inc
    zin = port.uf_tot / port.if_tot

    csv_path = os.path.join(sim_path, "port_s11.csv")
    with open(csv_path, "w") as fh:
        fh.write("freq_hz,s11_db,zin_real,zin_imag\n")
        for f, s, z in zip(freqs, s11, zin):
            fh.write(f"{f:.6e},{20*np.log10(max(abs(s),1e-12)):.6f},"
                     f"{z.real:.6f},{z.imag:.6f}\n")

    i = int(np.argmin(20 * np.log10(np.abs(s11) + 1e-12)))
    summary = {
        "case": tag, "cell_mm": args.cell, "n_cells": n_cells,
        "wall_s": round(wall, 1),
        "mc_per_s_engine": speeds[-1] if speeds else None,
        "iterations": iters[-1] if iters else None,
        "dip_ghz": round(freqs[i] / 1e9, 3),
        "dip_db": round(float(20 * np.log10(abs(s11[i]))), 2),
        "zin_at_dip": [round(zin[i].real, 1), round(zin[i].imag, 1)],
    }
    with open(os.path.join(sim_path, "summary.json"), "w") as fh:
        json.dump(summary, fh, indent=2)
    print(json.dumps(summary), flush=True)


if __name__ == "__main__":
    main()
