#!/usr/bin/env python3
"""Test a trained AFEM model on alternative scatterer geometries.

Supports: square, multi-circle, and the original circle.

Usage:
    python src/test_media.py                          # uses src/test_config.yaml
    python src/test_media.py --k-test 30.0 --geometry circle
    python src/test_media.py --config my_test.yaml    # custom config

All test parameters live in the YAML config. CLI args serve as overrides.
"""

import argparse
import copy
import os
import sys
import time
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from torch_geometric.data import Batch

_project_root = Path(__file__).resolve().parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from src.network import create_model
from src.utils import load_checkpoint, load_config, setup_helmholtz_config
from src.helmholtz_alt import (
    HelmholtzProblemSquare,
    HelmholtzProblemMultiCircle,
    create_helmholtz_problem_square,
    create_helmholtz_problem_multi_circle,
)


# ═══════════════════════════════════════════════════════════════════════
#  Geometry factory mapping
# ═══════════════════════════════════════════════════════════════════════

_GEOMETRY_FACTORIES = {
    "square": create_helmholtz_problem_square,
    "multi_circle": create_helmholtz_problem_multi_circle,
    "circle": None,  # default HelmholtzProblem
}


# ═══════════════════════════════════════════════════════════════════════
#  Epsilon_r property patching
# ═══════════════════════════════════════════════════════════════════════

def _patch_epsilon_r(env):
    inner_fp = env.fem_problem.fem_problem
    if hasattr(inner_fp, "eps_r_at_midpoints"):
        def _eps_r(self):
            return inner_fp.eps_r_at_midpoints(self.mesh)
        type(env)._epsilon_r_elements = property(_eps_r)


# ═══════════════════════════════════════════════════════════════════════
#  Fine FEM reference (computed once, interpolated later)
# ═══════════════════════════════════════════════════════════════════════

def _compute_fine_fem_reference(env, n_refine: int = 2):
    """Compute fine-FEM reference on initial mesh + n_refine uniform refinement."""
    from skfem import Basis, ElementTriP1

    fp = env.fem_problem.fem_problem
    ref_mesh = copy.deepcopy(env.mesh)
    for _ in range(n_refine):
        ref_mesh = ref_mesh.refined(np.arange(ref_mesh.t.shape[1]))
    ref_basis = Basis(ref_mesh, ElementTriP1())
    ref_sol = fp.calculate_solution(ref_basis, cache=False)

    # Interpolate to coarse mesh vertices
    pts = env.mesh.p.T
    finder = ref_mesh.element_finder()
    cells = finder(*pts.T)
    cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1)

    i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells]
    p = ref_mesh.p
    x, y = pts[:, 0], pts[:, 1]
    x0, y0 = p[0, i0], p[1, i0]
    x1, y1 = p[0, i1], p[1, i1]
    x2, y2 = p[0, i2], p[1, i2]
    denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    denom = np.where(np.abs(denom) < 1e-15, 1.0, denom)
    w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom
    w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom
    w2 = 1.0 - w0 - w1
    u_ref_on_coarse = w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2]
    return u_ref_on_coarse, ref_mesh, ref_sol


def _interpolate_ref_to_mesh(target_pts, ref_mesh, ref_sol):
    """Interpolate cached reference solution to arbitrary mesh vertices."""
    finder = ref_mesh.element_finder()
    cells = finder(*target_pts.T)
    cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1)

    i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells]
    p = ref_mesh.p
    x, y = target_pts[:, 0], target_pts[:, 1]
    x0, y0 = p[0, i0], p[1, i0]
    x1, y1 = p[0, i1], p[1, i1]
    x2, y2 = p[0, i2], p[1, i2]
    denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    denom = np.where(np.abs(denom) < 1e-15, 1.0, denom)
    w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom
    w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom
    w2 = 1.0 - w0 - w1
    return w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2]


def _compute_ref_grid(env, n_refine: int = 3, resolution: int = 200):
    """Compute fine reference on a regular grid for smooth heatmaps."""
    from skfem import Basis, ElementTriP1

    fp = env.fem_problem.fem_problem
    ref_mesh = copy.deepcopy(env.mesh)
    for _ in range(n_refine):
        ref_mesh = ref_mesh.refined(np.arange(ref_mesh.t.shape[1]))
    ref_basis = Basis(ref_mesh, ElementTriP1())
    ref_sol = fp.calculate_solution(ref_basis, cache=False)

    boundary = fp._domain._boundary
    x_vec = np.linspace(boundary[0], boundary[2], resolution)
    y_vec = np.linspace(boundary[1], boundary[3], resolution)
    X, Y = np.meshgrid(x_vec, y_vec)
    grid_pts = np.column_stack([X.ravel(), Y.ravel()])

    U_grid = np.zeros(len(grid_pts), dtype=np.complex128)
    batch_size = 4096
    for start in range(0, len(grid_pts), batch_size):
        end = min(start + batch_size, len(grid_pts))
        batch = grid_pts[start:end]
        finder = ref_mesh.element_finder()
        cells = finder(*batch.T)
        cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1)

        i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells]
        p = ref_mesh.p
        x, y = batch[:, 0], batch[:, 1]
        x0, y0 = p[0, i0], p[1, i0]
        x1, y1 = p[0, i1], p[1, i1]
        x2, y2 = p[0, i2], p[1, i2]
        denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
        denom = np.where(np.abs(denom) < 1e-15, 1.0, denom)
        w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom
        w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom
        w2 = 1.0 - w0 - w1
        U_grid[start:end] = w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2]

    return {"X": X, "Y": Y, "E_scat": U_grid.reshape(resolution, resolution)}


def _compute_step_error(scalar, u_ref) -> float:
    if u_ref is None:
        return float("nan")
    diff = np.abs(scalar - u_ref)
    denom = np.linalg.norm(np.abs(u_ref))
    if denom < 1e-12:
        denom = 1.0
    return float(np.linalg.norm(diff) / denom)


# ═══════════════════════════════════════════════════════════════════════
#  Visualization
# ═══════════════════════════════════════════════════════════════════════

def _render_field(ax, triang, values, title, vmin, vmax, show_mesh=True):
    tcf = ax.tripcolor(triang, values, shading="gouraud", cmap="jet",
                       vmin=vmin, vmax=vmax)
    if show_mesh and triang is not None:
        n = triang.triangles.shape[0]
        ax.triplot(triang, lw=(0.5 if n < 500 else 0.3), color="black",
                   alpha=(0.7 if n < 2000 else 0.5))
    ax.set_aspect("equal")
    ax.set_title(title, fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
    return tcf


def _draw_scatterer(ax, geometry: str, env):
    fp = env.fem_problem.fem_problem
    if geometry == "square":
        sq = getattr(fp, "_sq_cx", 0.5), getattr(fp, "_sq_cy", 0.5)
        hs = getattr(fp, "_sq_half", 0.2)
        ang = getattr(fp, "_sq_angle", 0.0)
        corners = np.array([
            [-hs, -hs], [hs, -hs], [hs, hs], [-hs, hs], [-hs, -hs]
        ])
        if ang != 0:
            c, s = np.cos(ang), np.sin(ang)
            corners = corners @ np.array([[c, -s], [s, c]]).T
        corners[:, 0] += sq[0]
        corners[:, 1] += sq[1]
        ax.plot(corners[:, 0], corners[:, 1], color="cyan", linewidth=1.5,
                linestyle="--")
    elif geometry == "multi_circle":
        circles = getattr(fp, "_circles", [])
        for c in circles:
            theta = np.linspace(0, 2 * np.pi, 128)
            ax.plot(c["cx"] + c["radius"] * np.cos(theta),
                    c["cy"] + c["radius"] * np.sin(theta),
                    color="cyan", linewidth=1.5, linestyle="--")
    elif geometry == "circle":
        cx = getattr(fp, "_cx", 0.5)
        cy = getattr(fp, "_cy", 0.5)
        r = getattr(fp, "_radius", 0.2)
        theta = np.linspace(0, 2 * np.pi, 128)
        ax.plot(cx + r * np.cos(theta), cy + r * np.sin(theta),
                color="cyan", linewidth=1.5, linestyle="--")


def _save_pngs(steps, stem, checkpoint_path, k, geometry, env, ref_grid):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.tri as tri

    per_step_dir = f"{stem}_steps"
    os.makedirs(os.path.dirname(stem) or ".", exist_ok=True)
    os.makedirs(per_step_dir, exist_ok=True)

    # ── Overview grid ──
    n = len(steps)
    ncols = min(n, 4)
    nrows = (n + ncols - 1) // ncols
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3.5 * nrows))
    axes_flat = np.array([axes]) if nrows * ncols == 1 else np.array(axes).flatten()

    for i, step_data in enumerate(steps):
        mesh, scalar, err_val, n_elem = step_data[:4]
        pts = mesh.p.T
        tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
        s = np.abs(scalar) if np.iscomplexobj(scalar) else scalar
        vmin, vmax = s.min(), s.max()
        if vmax - vmin < 1e-12:
            vmin, vmax = vmin - 0.5, vmax + 0.5
        tcf = _render_field(axes_flat[i], tg, s,
                            f"Step {i}: {n_elem} elem, err={err_val:.4f}",
                            vmin, vmax)
        fig.colorbar(tcf, ax=axes_flat[i], fraction=0.046, pad=0.04)
        _draw_scatterer(axes_flat[i], geometry, env)

    for j in range(n, len(axes_flat)):
        axes_flat[j].set_visible(False)

    fig.subplots_adjust(left=0.04, right=0.90, top=0.90, bottom=0.06,
                        wspace=0.15, hspace=0.30)
    geo_label = {"square": "Square", "multi_circle": "Multi-Circle",
                 "circle": "Circle"}.get(geometry, geometry)
    fig.suptitle(
        f"Helmholtz |E_scat| [{geo_label}] — {os.path.basename(checkpoint_path)}\n"
        f"k={k:.1f}  eps_r info in scatterer overlay",
        fontsize=12,
    )
    fig.savefig(f"{stem}.png", dpi=200, bbox_inches="tight")
    plt.close(fig)
    print(f"[Viz] Overview → {stem}.png")

    # ── Per-step panels (FEM + Reference + Error) ──
    for i, step_data in enumerate(steps):
        mesh, scalar, err_val, n_elem = step_data[:4]
        u_ref_at_verts = step_data[4] if len(step_data) > 4 else None

        pts = mesh.p.T
        tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
        coarse_val = np.abs(scalar) if np.iscomplexobj(scalar) else scalar

        fig2, axes2 = plt.subplots(1, 3, figsize=(18, 6))
        axes2 = list(np.atleast_1d(axes2))

        # Panel 1: FEM
        cvmin, cvmax = coarse_val.min(), coarse_val.max()
        if cvmax - cvmin < 1e-12:
            cvmin, cvmax = cvmin - 0.5, cvmax + 0.5
        tcf1 = _render_field(axes2[0], tg, coarse_val,
                             f"Step {i}: FEM |E_scat| ({n_elem} elem)",
                             cvmin, cvmax)
        _draw_scatterer(axes2[0], geometry, env)
        fig2.colorbar(tcf1, ax=axes2[0], fraction=0.046, pad=0.04)

        # Panel 2: Fine FEM reference on grid
        if ref_grid is not None:
            g = ref_grid
            gm = np.abs(g["E_scat"])
            mvmin, mvmax = gm.min(), gm.max()
            if mvmax - mvmin < 1e-12:
                mvmin, mvmax = mvmin - 0.5, mvmax + 0.5
            im2 = axes2[1].pcolormesh(g["X"], g["Y"], gm,
                                      shading="gouraud", cmap="jet",
                                      vmin=mvmin, vmax=mvmax)
            axes2[1].set_title("Fine FEM Ref |E_scat|", fontsize=9)
            axes2[1].set_aspect("equal")
            axes2[1].set_xticks([])
            axes2[1].set_yticks([])
            _draw_scatterer(axes2[1], geometry, env)
            fig2.colorbar(im2, ax=axes2[1], fraction=0.046, pad=0.04)

        # Panel 3: Pointwise error
        if u_ref_at_verts is not None:
            u_fem_abs = np.abs(scalar)
            u_ref_abs = np.abs(u_ref_at_verts)
            error_abs = np.abs(u_fem_abs - u_ref_abs)
            evmin, evmax = 0.0, error_abs.max() or 1.0
            if evmax - evmin < 1e-12:
                evmax = evmin + 1.0
            tcf3 = _render_field(axes2[2], tg, error_abs,
                                 f"||FEM|−|Ref||  L2={err_val:.4f}",
                                 evmin, evmax)
            _draw_scatterer(axes2[2], geometry, env)
            fig2.colorbar(tcf3, ax=axes2[2], fraction=0.046, pad=0.04)

        fig2.tight_layout()
        fig2.savefig(f"{per_step_dir}/step{i:02d}.png", dpi=150,
                     bbox_inches="tight")
        plt.close(fig2)

    print(f"[Viz] Per-step PNGs → {per_step_dir}/  ({n} files)")


# ═══════════════════════════════════════════════════════════════════════
#  Scatterer config injection
# ═══════════════════════════════════════════════════════════════════════

def _inject_scatterer_config(base_config: dict, geometry: str, sc_cfg: dict, k_test: float):
    """Inject scatterer params from test config into the base config's helmholtz section.

    Returns (config, factory) where factory is the geometry-specific create function.
    """
    hc = (base_config.setdefault("environment", {})
          .setdefault("mesh_refinement", {})
          .setdefault("fem", {})
          .setdefault("helmholtz", {}))

    sc = hc.setdefault("scatterer", {})
    sc["mode"] = "fixed"
    sc["eps_r"] = float(sc_cfg.get("eps_r", 3.0))

    if geometry == "square":
        sc["square"] = {
            "cx": float(sc_cfg.get("cx", 0.5)),
            "cy": float(sc_cfg.get("cy", 0.5)),
            "half_side": float(sc_cfg.get("half_side", 0.15)),
            "angle": float(sc_cfg.get("angle", 0.0)),
        }
    elif geometry == "multi_circle":
        circles_raw = sc_cfg.get("circles", [])
        circles = []
        for c in circles_raw:
            circles.append({
                "cx": float(c["cx"]), "cy": float(c["cy"]),
                "radius": float(c["radius"]),
                "eps_r": float(c.get("eps_r", sc_cfg.get("eps_r", 3.0))),
            })
        sc["circles"] = circles
    elif geometry == "circle":
        sc["cx"] = float(sc_cfg.get("cx", 0.5))
        sc["cy"] = float(sc_cfg.get("cy", 0.5))
        sc["radius"] = float(sc_cfg.get("radius", 0.2))

    hc["wave_number_mode"] = "fixed"
    hc["wave_number"] = float(k_test)

    factory = _GEOMETRY_FACTORIES.get(geometry)
    return base_config, factory


# ═══════════════════════════════════════════════════════════════════════
#  Main test function
# ═══════════════════════════════════════════════════════════════════════

def test_alt_media(
    base_config: dict,
    test_cfg: dict,
    cli_overrides: Optional[dict] = None,
):
    """Run AFEM inference with config-driven parameters.

    Args:
        base_config: loaded from config.yaml (model/network/algo)
        test_cfg: loaded from test_config.yaml (test-specific params)
        cli_overrides: optional CLI arg overrides dict
    """
    ov = cli_overrides or {}

    # ── Resolve parameters: test_cfg < CLI override ──
    tc = test_cfg.get("test", {})
    ref_cfg = test_cfg.get("reference", {})
    sc_cfg = test_cfg.get("scatterer", {})

    geometry = ov.get("geometry") or tc.get("geometry", "circle")
    checkpoint_path = ov.get("checkpoint") or tc.get("checkpoint", "checkpoints/model_final.pt")
    output_path = ov.get("output") or tc.get("output", "result/test_media.png")
    seed = ov.get("seed") or tc.get("seed", 99)
    k_test = ov.get("k_test") or test_cfg.get("k_test", 8.0)
    n_refine_vertex = ov.get("n_refine_vertex") or ref_cfg.get("n_refine_vertex", 2)
    n_refine_grid = ov.get("n_refine_grid") or ref_cfg.get("n_refine_grid", 3)
    grid_resolution = ov.get("grid_resolution") or ref_cfg.get("grid_resolution", 200)

    # Allow CLI override of scatterer params
    for key in ("cx", "cy", "radius", "eps_r", "half_side", "angle"):
        if ov.get(key) is not None:
            sc_cfg[key] = ov[key]
    if ov.get("circles") is not None:
        sc_cfg["circles"] = ov["circles"]

    algo = base_config.get("algorithm", {})

    # ── 1. Inject scatterer config ──
    config, factory = _inject_scatterer_config(
        copy.deepcopy(base_config), geometry, sc_cfg, k_test)

    # ── 2. Create env with alt factory ──
    import environment.fem_problem as fem_problem_module

    _orig_create = None
    if factory is not None:
        _orig_create = fem_problem_module.create_helmholtz_problem
        fem_problem_module.create_helmholtz_problem = factory

    from environment.mesh_refinement import MeshRefinement
    env = MeshRefinement(
        environment_config=config.get("environment", {}).get("mesh_refinement", {}),
        seed=seed,
    )

    # ── 3. Load model ──
    model = create_model(env, config.get("network", {}), algo.get("ppo", {}))
    load_checkpoint(model, checkpoint_path)
    model.eval()
    dev = next(model.parameters()).device
    print(f"[Device] {dev}")
    model = model.to(dev)

    # ── 4. Reset env ──
    print(f"[Test] Geometry: {geometry}  k={k_test:.3f}")
    obs = env.reset()

    # ── 5. Patch epsilon_r_elements (after reset) ──
    _patch_epsilon_r(env)

    # Restore original factory
    if _orig_create is not None:
        fem_problem_module.create_helmholtz_problem = _orig_create

    # ── 6. Print scatterer info ──
    fp = env.fem_problem.fem_problem
    if geometry == "square":
        print(f"[Test]   Square: center=({getattr(fp, '_sq_cx', 0.5):.3f}, "
              f"{getattr(fp, '_sq_cy', 0.5):.3f})  half_side={getattr(fp, '_sq_half', 0.2):.3f}")
    elif geometry == "multi_circle":
        circles_attr = getattr(fp, "_circles", [])
        for i, c in enumerate(circles_attr):
            print(f"[Test]   Circle {i}: center=({c['cx']:.3f}, {c['cy']:.3f})  "
                  f"r={c['radius']:.3f}  eps_r={c['eps_r']:.1f}")
    elif geometry == "circle":
        print(f"[Test]   Circle: center=({getattr(fp, '_cx', 0.5):.3f}, "
              f"{getattr(fp, '_cy', 0.5):.3f})  r={getattr(fp, '_radius', 0.2):.3f}")

    # ── 7. Compute fine-FEM reference ONCE on initial mesh ──
    n_init = env.mesh.t.shape[1]
    print(f"[Test] Initial mesh: {n_init} elements")
    print(f"[Test] Computing fine-FEM reference (n_refine_vertex={n_refine_vertex}, "
          f"n_refine_grid={n_refine_grid}, grid={grid_resolution})...")

    t0 = time.time()
    u_ref_initial, ref_mesh, ref_sol = _compute_fine_fem_reference(env, n_refine=n_refine_vertex)
    ref_grid = _compute_ref_grid(env, n_refine=n_refine_grid, resolution=grid_resolution)
    print(f"[Test] Reference ready ({time.time() - t0:.1f}s, grid {ref_grid['X'].shape})")

    # ── 8. Run inference ──
    stem = output_path.rsplit(".", 1)[0] if "." in output_path else output_path
    init_mesh = env.mesh
    init_sol = env.scalar_solution
    init_err = _compute_step_error(init_sol, u_ref_initial)
    steps = [(init_mesh, init_sol, init_err, env.num_agents, u_ref_initial)]

    n_elem_init = env.num_agents
    print(f"  Step  0: reward=---  err={init_err:.4f}  elements={n_elem_init}")

    done = False
    step_idx = 0
    total_reward = 0.0
    while not done:
        obs_g = obs.to(dev)
        with torch.no_grad():
            actions, _, _ = model(Batch.from_data_list([obs_g]), deterministic=True)
        obs, reward, done, info = env.step(actions.cpu().numpy())
        step_r = float(np.sum(reward))
        total_reward += step_r
        step_idx += 1

        # Interpolate cached reference to current mesh vertices (no re-solve)
        u_ref_current = _interpolate_ref_to_mesh(env.mesh.p.T, ref_mesh, ref_sol)
        step_err = _compute_step_error(env.scalar_solution, u_ref_current)
        steps.append((env.mesh, env.scalar_solution, step_err, env.num_agents,
                      u_ref_current))

        print(f"  Step {step_idx:2d}: reward={step_r:+.4f} err={step_err:.4f}  "
              f"elements={info.get('num_elements', '?')}  "
              f"sel={info.get('selected_count', 0)}  "
              f"done={done}")

    print(f"\n[Test] total_reward={total_reward:.4f}  final_err={steps[-1][2]:.4f}  "
          f"final_elements={steps[-1][3]}")

    # ── 9. Visualize ──
    _save_pngs(steps, stem, checkpoint_path, k_test, geometry, env, ref_grid)
    print(f"[Viz] Done → {output_path}")


# ═══════════════════════════════════════════════════════════════════════
#  CLI
# ═══════════════════════════════════════════════════════════════════════

def _load_yaml(path: str) -> dict:
    """Load a YAML file, resolving relative paths against project root."""
    import yaml
    if not os.path.isabs(path):
        path = os.path.join(_project_root, path)
    with open(path, "r") as f:
        return yaml.safe_load(f)


def main():
    parser = argparse.ArgumentParser(
        description="Test AFEM trained model on alternative scatterer geometries")

    # Config
    parser.add_argument("--config", default="src/test_config.yaml",
                        help="Test config YAML (default: src/test_config.yaml)")

    # Test scenario overrides
    parser.add_argument("--geometry", choices=["square", "multi_circle", "circle"],
                        help="Scatterer geometry (overrides config)")
    parser.add_argument("--checkpoint", help="Model checkpoint path (overrides config)")
    parser.add_argument("--output", help="Output image path (overrides config)")
    parser.add_argument("--seed", type=int, help="Random seed (overrides config)")
    parser.add_argument("--k-test", type=float, help="Wave number (overrides config)")

    # Scatterer overrides
    parser.add_argument("--cx", type=float, help="Scatterer center x")
    parser.add_argument("--cy", type=float, help="Scatterer center y")
    parser.add_argument("--radius", type=float, help="Scatterer radius (circle)")
    parser.add_argument("--eps-r", type=float, help="Dielectric constant eps_r")
    parser.add_argument("--half-side", type=float, help="Half side length (square)")
    parser.add_argument("--angle", type=float, help="Rotation angle in radians (square)")
    parser.add_argument("--circles", nargs="*", default=None,
                        help="Circle specs: 'cx,cy,radius[,eps_r]' (multi_circle)")

    # Reference computation overrides
    parser.add_argument("--n-refine-vertex", type=int,
                        help="Uniform refinement levels for vertex error reference")
    parser.add_argument("--n-refine-grid", type=int,
                        help="Uniform refinement levels for grid heatmap reference")
    parser.add_argument("--grid-resolution", type=int,
                        help="Grid resolution N for heatmap (N x N)")

    args = parser.parse_args()

    # ── Load test config ──
    test_cfg = _load_yaml(args.config)

    # ── Load base config ──
    base_config_path = test_cfg.get("base_config", "src/config.yaml")
    base_config = _load_yaml(base_config_path)

    # ── Build CLI overrides dict (only non-None values) ──
    cli_overrides = {}
    for key in ("geometry", "checkpoint", "output", "seed", "k_test",
                "cx", "cy", "radius", "eps_r", "half_side", "angle",
                "n_refine_vertex", "n_refine_grid", "grid_resolution"):
        val = getattr(args, key.replace("-", "_"), None)
        if val is not None:
            cli_overrides[key] = val

    # Parse --circles if provided
    if args.circles is not None:
        circles = []
        for spec in args.circles:
            parts = [float(x.strip()) for x in spec.split(",")]
            circles.append({
                "cx": parts[0], "cy": parts[1], "radius": parts[2],
                "eps_r": parts[3] if len(parts) > 3 else 3.0,
            })
        cli_overrides["circles"] = circles

    # ── Set seeds ──
    seed = cli_overrides.get("seed", test_cfg.get("test", {}).get("seed", 99))
    torch.manual_seed(seed)
    np.random.seed(seed)

    test_alt_media(
        base_config=base_config,
        test_cfg=test_cfg,
        cli_overrides=cli_overrides,
    )


if __name__ == "__main__":
    main()
