#!/usr/bin/env python3
"""Reproduce pipeline inference: legacy models on 1280x512 config.

Loads pre-trained PPO models from models/old/ and models/250525/,
runs deterministic inference on config_lbm_pinball.json (parabolic, bounce-back).

Key differences from train pipeline:
  - 1280x512 grid (not 2000x600)
  - Legacy norm (force_norm_fact + sens_deviation + sens_norm_fact)
  - Old action scaling with biases (match legacy env)
  - No VecNormalize (old models use manual norm)
  - DummyEnv for PPO.load()

Usage:
    conda run -n pycuda_3_10 python -u infer_reproduce.py --device-id 1
    conda run -n pycuda_3_10 python -u infer_reproduce.py --device-id 1 --scene karman
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import pycuda.driver as cuda; cuda.init()

_REPO = str(Path(__file__).resolve().parents[3])
_SRC = Path(_REPO) / "src"
for p in [_REPO, str(_SRC)]:
    if p not in sys.path:
        sys.path.insert(0, p)

import torch
from torch.nn import Module as TorchModule
from stable_baselines3 import PPO

from CelerisLab import Simulation
from CelerisLab.common.render import compute_vorticity, render_vorticity_field
from CelerisLab.lbm.initializers import add_vortex
from drl_pinball.reproduce.configs.model_inventory import ModelInventory
from drl_pinball.eval.scene_manifest import REPRODUCE_SCENES

L0 = 20.0
U0 = 0.01
NX = 1280
NY = 512
CENTER_Y = float(NY - 1) / 2.0
RADIUS = L0 / 2.0
FIFO_LEN = 150
WARMUP = int(4.0 * NX / U0)

# Karman/Vortex/Steady geometry (1280x512)
DIST_X = 10.0 * L0   # 200
PB_FRONT_X = 30.0 * L0   # 600
PB_REAR_X = 31.3 * L0    # 626
SENSOR_X = 40.0 * L0     # 800

# Illusion geometry (shifted left)
ILL_PB_FRONT_X = 19.0 * L0   # 380
ILL_PB_REAR_X = 20.3 * L0    # 406
ILL_SENSOR_X = 30.0 * L0     # 600
ILL_TARGET_X = 20.0 * L0     # 400

VORTEX_RADIUS = 2.0 * L0
SR_DATA = _SRC / "SR_analysis" / "data"
_OUT_BASE = Path(__file__).resolve().parent / "output" / "reproduce"


class Sin(TorchModule):
    def __init__(self): super().__init__()
    def forward(self, x): return torch.sin(x)


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# ---------------------------------------------------------------------------
# Legacy helpers
# ---------------------------------------------------------------------------
class ActionSmoother:
    def __init__(self, weight=0.1):
        self.weight = weight; self._state = None
    def __call__(self, target):
        t = np.asarray(target, dtype=np.float32)
        if self._state is None:
            self._state = t.copy()
        else:
            self._state = (1.0 - self.weight) * self._state + self.weight * t
        return self._state.copy()
    def reset(self, value=None):
        self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None


def get_cc(sim, sid):
    nx, ny = sim.lbm_cfg.nx, sim.lbm_cfg.ny
    cells_arr, _ = sim.bodies.get(sid).get_sensor_list(nx, ny)
    return float(len(cells_arr))


def action_to_omega_legacy(action_norm, scale=8.0, bias=(0.0, -4.0, 4.0)):
    b = np.array(bias, dtype=np.float32)
    sv = (np.asarray(action_norm, dtype=np.float32) * scale + b) * U0
    return -sv / RADIUS


def load_legacy_norm(ref_dir: str) -> Dict[str, Any]:
    with open(os.path.join(ref_dir, "norm.json")) as f:
        d = json.load(f)
    return {
        "force_norm_fact": np.float32(d["force_norm_fact"]),
        "sens_deviation": np.array(d["sens_deviation"], dtype=np.float32),
        "sens_norm_fact": np.array(d["sens_norm_fact"], dtype=np.float32),
    }


def normalize_obs(obs_slice, norm):
    forces = obs_slice[6:12] / norm["force_norm_fact"]
    sens = (obs_slice[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]
    return np.clip(np.hstack([forces, sens]), -1.0, 1.0).astype(np.float32)


def save_vorticity(sim, out_path, cylinders, nx=NX, ny=NY):
    macro = sim.get_macroscopic()
    vort = compute_vorticity(macro["ux"], macro["uy"])
    render_vorticity_field(vort, nx=nx, ny=ny, out_path=str(out_path),
                           cylinders=cylinders, vmin=-0.03, vmax=0.03)


# ---------------------------------------------------------------------------
# Obs readers (legacy ordering)
# ---------------------------------------------------------------------------
def read_obs_karman(sim, dist_id, sensor_ids, pinball_ids, cc):
    """dist_fx, dist_fy, s0_ux, s0_uy, s1_ux, s1_uy, s2_ux, s2_uy,
       front_fx, front_fy, top_fx, top_fy, bot_fx, bot_fy"""
    obs = list(sim.read_force(dist_id, normalize=True))
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend([float(s[0]) * cc, float(s[1]) * cc])
    for pid in pinball_ids:
        obs.extend(sim.read_force(pid, normalize=True))
    return np.array(obs, dtype=np.float32)


def read_obs_6obj(sim, sensor_ids, pinball_ids, cc):
    obs = []
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend([float(s[0]) * cc, float(s[1]) * cc])
    for pid in pinball_ids:
        obs.extend(sim.read_force(pid, normalize=True))
    return np.array(obs, dtype=np.float32)


# ---------------------------------------------------------------------------
# Karman Cloak
# ---------------------------------------------------------------------------
def run_karman_reproduce(scene: Dict[str, Any], device_id: int, out_dir: Path) -> None:
    log(f"=== Reproduce: Karman Cloak ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = scene["si"]
    num_steps = scene["num_steps"]
    model_name = scene["model_name"]
    scale = scene["action_scale"]
    bias = scene["action_bias"]

    # Phase 1: Disturbance + sensors, record target
    log("  Phase 1: Recording target...")
    sim = Simulation(lbm_config_path=scene["config_path"], device_id=device_id)
    dist_id = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0)
    sensor_ids = [
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    log(f"  Sensor cell count: {cc}")

    target_states = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_obs_karman(sim, dist_id, sensor_ids, [], cc)
        target_states[i] = obs[2:8]
    np.savez_compressed(out_dir / "target.npz", target_states=target_states)

    # Save target vorticity
    save_vorticity(sim, out_dir / "vorticity_target.png",
                   [((DIST_X, CENTER_Y), 1.0 * L0),
                    ((SENSOR_X, CENTER_Y + 40.0), 5.0),
                    ((SENSOR_X, CENTER_Y), 5.0),
                    ((SENSOR_X, CENTER_Y - 40.0), 5.0)])

    # Phase 2: Add pinball
    log("  Phase 2: Adding pinball...")
    n0 = sim.bodies.count
    sim.add_body("circle", center=(PB_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.sync_bodies()
    fid, tid, bid = list(range(n0, n0 + 3))
    sim.run(WARMUP, zero_obs=True)

    # Load legacy norm
    ref_dir = str(SR_DATA / "karman" / "karman_re100")
    norm = load_legacy_norm(ref_dir)

    # Bias FIFO
    bias_norm = np.array([0.0, -1.0, 1.0])
    bias_omega = action_to_omega_legacy(bias_norm, scale=scale, bias=bias)
    ema = ActionSmoother(weight=0.1); ema.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)
    sim.snapshot()

    # DRL inference
    log("  Phase 3: DRL inference...")
    sim.restore()
    ema.reset(bias_omega.copy())

    model = ModelInventory().load(model_name, device="cpu")
    log(f"  Model: {model_name} on CPU")

    obs_init = read_obs_karman(sim, dist_id, sensor_ids, [fid, tid, bid], cc)
    obs_norm = normalize_obs(obs_init[2:14], norm)

    sig_s = np.zeros((num_steps, 6), dtype=np.float32)
    sig_f = np.zeros((num_steps, 6), dtype=np.float32)
    sig_a = np.zeros((num_steps, 3), dtype=np.float32)
    fifo = deque(maxlen=FIFO_LEN)
    [fifo.append(target_states[i, :].copy()) for i in range(FIFO_LEN)]

    for step in range(num_steps):
        action, _ = model.predict(obs_norm, deterministic=True)
        action = np.asarray(action, dtype=np.float32).flatten()
        target_omega = action_to_omega_legacy(action, scale=scale, bias=bias)
        smoothed = ema(target_omega)

        sim.set_body(fid, omega=smoothed[0]); sim.set_body(tid, omega=smoothed[1]); sim.set_body(bid, omega=smoothed[2])
        sim.run(SI, zero_obs=True)

        obs = read_obs_karman(sim, dist_id, sensor_ids, [fid, tid, bid], cc)
        sl = obs[2:14]
        fifo.append(sl[0:6].copy())
        sig_s[step] = sl[0:6]
        sig_f[step] = sl[6:12]
        sig_a[step] = action
        obs_norm = normalize_obs(sl, norm)

    save_vorticity(sim, out_dir / "vorticity_controlled.png", [
        ((DIST_X, CENTER_Y), 1.0 * L0),
        ((PB_FRONT_X, CENTER_Y), RADIUS),
        ((PB_REAR_X, CENTER_Y + 15.0), RADIUS),
        ((PB_REAR_X, CENTER_Y - 15.0), RADIUS),
    ])
    sim.close()

    np.savez_compressed(out_dir / "signals.npz",
                        sensors=sig_s, forces=sig_f, actions=sig_a)
    metrics = {"dtw_sim_v5": -1.0}  # DTW computed by generate_report later
    with open(out_dir / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Steady Cloak (open-loop)
# ---------------------------------------------------------------------------
def run_steady_reproduce(scene: Dict[str, Any], device_id: int, out_dir: Path) -> None:
    log(f"=== Reproduce: Steady Cloak ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = scene["si"]
    num_steps = scene["num_steps"]
    surf_vel = scene.get("open_loop_surf_vel", (0.0, -5.1, 5.1))
    bias_surf = np.array(surf_vel, dtype=np.float32) * U0

    sim = Simulation(lbm_config_path=scene["config_path"], device_id=device_id)
    sensor_ids = [
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(PB_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])

    save_vorticity(sim, out_dir / "vorticity_zero.png", [
        ((PB_FRONT_X, CENTER_Y), RADIUS),
        ((PB_REAR_X, CENTER_Y + 15.0), RADIUS),
        ((PB_REAR_X, CENTER_Y - 15.0), RADIUS),
    ])

    bias_omega = -bias_surf / RADIUS
    ema = ActionSmoother(weight=0.1); ema.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema(bias_omega)
        sim.set_body(3, omega=s[0]); sim.set_body(4, omega=s[1]); sim.set_body(5, omega=s[2])
        sim.run(SI, zero_obs=True)
    sim.snapshot()
    sim.restore()
    ema.reset(bias_omega.copy())

    sig_s = np.zeros((num_steps, 6), dtype=np.float32)
    sig_f = np.zeros((num_steps, 6), dtype=np.float32)
    for step in range(num_steps):
        smoothed = ema(bias_omega)
        sim.set_body(3, omega=smoothed[0]); sim.set_body(4, omega=smoothed[1]); sim.set_body(5, omega=smoothed[2])
        sim.run(SI, zero_obs=True)
        obs = read_obs_6obj(sim, sensor_ids, [3, 4, 5], cc)
        sig_s[step] = obs[0:6]
        sig_f[step] = obs[6:12]

    save_vorticity(sim, out_dir / "vorticity_controlled.png", [
        ((PB_FRONT_X, CENTER_Y), RADIUS),
        ((PB_REAR_X, CENTER_Y + 15.0), RADIUS),
        ((PB_REAR_X, CENTER_Y - 15.0), RADIUS),
    ])
    sim.close()

    np.savez_compressed(out_dir / "signals.npz", sensors=sig_s, forces=sig_f,
                        actions=np.zeros((num_steps, 3), dtype=np.float32))
    with open(out_dir / "metrics.json", "w") as f:
        json.dump({"dtw_sim_v5": -1.0}, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Illusion
# ---------------------------------------------------------------------------
def run_illusion_reproduce(scene: Dict[str, Any], device_id: int, out_dir: Path) -> None:
    scene_id = scene["scene_id"]
    log(f"=== Reproduce: Illusion {scene_id} ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = scene["si"]
    num_steps = scene["num_steps"]
    model_name = scene["model_name"]
    scale = scene["action_scale"]
    bias = scene["action_bias"]
    target_diam = scene.get("target_diam", 1.0)

    # Load legacy data
    ref_dir = str(SR_DATA / "illusion" / scene_id)
    norm = load_legacy_norm(ref_dir)
    with open(os.path.join(ref_dir, "target_harmonics.json")) as f:
        target_harmonics = json.load(f)
    legacy_target = np.load(os.path.join(ref_dir, "target.npz"))["target_states"]

    def gen_target_at(t):
        D = len(target_harmonics)
        result = np.zeros(D, dtype=np.float32)
        for d, h in enumerate(target_harmonics):
            val = np.float32(h["dc"])
            for amp, freq, phase in zip(h["amps"], h["freqs"], h["phases"]):
                val += amp * np.cos(2.0 * np.pi * freq * t + phase)
            result[d] = val
        return result

    sim = Simulation(lbm_config_path=scene["config_path"], device_id=device_id)
    sensor_ids = [
        sim.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(ILL_PB_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(ILL_PB_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(ILL_PB_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    fid, tid, bid = 3, 4, 5

    # Bias FIFO
    bias_surf = np.array([0.0, -1.0, 1.0], dtype=np.float32) * U0
    bias_omega = -bias_surf / RADIUS
    ema_b = ActionSmoother(weight=0.1); ema_b.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema_b(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)
    sim.snapshot()

    # DRL inference
    sim.restore()
    ema = ActionSmoother(weight=0.1); ema.reset(bias_omega.copy())
    model = ModelInventory().load(model_name, device="cpu")
    log(f"  Model: {model_name} on CPU")

    obs_init = read_obs_6obj(sim, sensor_ids, [fid, tid, bid], cc)
    obs_norm_base = normalize_obs(obs_init, norm)
    t0 = gen_target_at(0)
    target_cd = np.float32(t0[0] / norm["force_norm_fact"])
    target_cl = np.float32(t0[1] / norm["force_norm_fact"])
    obs_norm = np.clip(np.hstack([obs_norm_base, [target_cd, target_cl]]), -1.0, 1.0).astype(np.float32)

    sig_s = np.zeros((num_steps, 6), dtype=np.float32)
    sig_f = np.zeros((num_steps, 6), dtype=np.float32)
    sig_a = np.zeros((num_steps, 3), dtype=np.float32)

    for step in range(num_steps):
        action, _ = model.predict(obs_norm, deterministic=True)
        action = np.asarray(action, dtype=np.float32).flatten()
        target_omega = action_to_omega_legacy(action, scale=scale, bias=bias)
        smoothed = ema(target_omega)

        sim.set_body(fid, omega=smoothed[0]); sim.set_body(tid, omega=smoothed[1]); sim.set_body(bid, omega=smoothed[2])
        sim.run(SI, zero_obs=True)

        obs = read_obs_6obj(sim, sensor_ids, [fid, tid, bid], cc)
        sig_s[step] = obs[0:6]
        sig_f[step] = obs[6:12]
        sig_a[step] = action

        forces_n = obs[6:12] / norm["force_norm_fact"]
        sens_n = (obs[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]
        t_h = gen_target_at(step)
        target_cd = np.float32(t_h[0] / norm["force_norm_fact"])
        target_cl = np.float32(t_h[1] / norm["force_norm_fact"])
        obs_norm = np.clip(np.hstack([forces_n, sens_n, [target_cd, target_cl]]), -1.0, 1.0).astype(np.float32)

    save_vorticity(sim, out_dir / "vorticity_controlled.png", [
        ((ILL_PB_FRONT_X, CENTER_Y), RADIUS),
        ((ILL_PB_REAR_X, CENTER_Y + 15.0), RADIUS),
        ((ILL_PB_REAR_X, CENTER_Y - 15.0), RADIUS),
    ])
    sim.close()

    np.savez_compressed(out_dir / "signals.npz",
                        sensors=sig_s, forces=sig_f, actions=sig_a)
    with open(out_dir / "metrics.json", "w") as f:
        json.dump({"dtw_sim_v5": -1.0}, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Vortex
# ---------------------------------------------------------------------------
def run_vortex_reproduce(scene: Dict[str, Any], device_id: int, out_dir: Path) -> None:
    scene_id = scene["scene_id"]
    vortex_type = scene["vortex_type"]
    model_name = scene["model_name"]
    SI = scene["si"]
    num_steps = scene["num_steps"]
    scale = scene["action_scale"]
    bias = scene["action_bias"]
    vortex_strength = scene["vortex_strength_factor"] * U0

    log(f"=== Reproduce: Vortex {vortex_type} ===")
    out_dir.mkdir(parents=True, exist_ok=True)

    ref_dir = str(SR_DATA / "vortex" / scene_id)
    norm = load_legacy_norm(ref_dir)

    # Stage 1: Sensors only, record target with vortex at x=10
    log("  Stage 1: Recording target with vortex...")
    sim = Simulation(lbm_config_path=scene["config_path"], device_id=device_id)
    sensor_ids = [
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    sim.snapshot()

    add_vortex(sim.field, center=(10.0 * L0, CENTER_Y), radius=VORTEX_RADIUS,
               strength=vortex_strength, vortex_type=vortex_type)
    target_states = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_obs_6obj(sim, sensor_ids, [], cc)
        target_states[i] = obs
    np.savez_compressed(out_dir / "target.npz", target_states=target_states)

    # Stage 2: Add pinball + vortex at x=15
    log("  Stage 2: Adding pinball + vortex at x=15...")
    sim.restore()
    n0 = sim.bodies.count
    sim.add_body("circle", center=(PB_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PB_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.sync_bodies()
    fid, tid, bid = list(range(n0, n0 + 3))

    bias_surf = np.array([0.0, -4.0, 4.0], dtype=np.float32) * U0
    bias_omega = -bias_surf / RADIUS
    ema_init = ActionSmoother(weight=0.1); ema_init.reset(np.zeros(3, dtype=np.float32))
    for _ in range(100):
        s = ema_init(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)

    add_vortex(sim.field, center=(15.0 * L0, CENTER_Y), radius=VORTEX_RADIUS,
               strength=vortex_strength, vortex_type=vortex_type)
    sim.snapshot()

    # Bias FIFO
    sim.restore()
    ema_b = ActionSmoother(weight=0.1); ema_b.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema_b(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)
    sim.snapshot()

    # DRL inference
    log("  Stage 3: DRL inference...")
    sim.restore()
    ema = ActionSmoother(weight=0.1); ema.reset(bias_omega.copy())
    model = ModelInventory().load(model_name, device="cpu")
    log(f"  Model: {model_name} on CPU")

    obs_init = read_obs_6obj(sim, sensor_ids, [fid, tid, bid], cc)
    obs_norm = normalize_obs(obs_init, norm)

    sig_s = np.zeros((num_steps, 6), dtype=np.float32)
    sig_f = np.zeros((num_steps, 6), dtype=np.float32)
    sig_a = np.zeros((num_steps, 3), dtype=np.float32)

    for step in range(num_steps):
        action, _ = model.predict(obs_norm, deterministic=True)
        action = np.asarray(action, dtype=np.float32).flatten()
        target_omega = action_to_omega_legacy(action, scale=scale, bias=bias)
        smoothed = ema(target_omega)

        sim.set_body(fid, omega=smoothed[0]); sim.set_body(tid, omega=smoothed[1]); sim.set_body(bid, omega=smoothed[2])
        sim.run(SI, zero_obs=True)

        obs = read_obs_6obj(sim, sensor_ids, [fid, tid, bid], cc)
        sig_s[step] = obs[0:6]
        sig_f[step] = obs[6:12]
        sig_a[step] = action
        forces_n = obs[6:12] / norm["force_norm_fact"]
        sens_n = (obs[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]
        obs_norm = np.clip(np.hstack([forces_n, sens_n]), -1.0, 1.0).astype(np.float32)

    save_vorticity(sim, out_dir / "vorticity_controlled.png", [
        ((PB_FRONT_X, CENTER_Y), RADIUS),
        ((PB_REAR_X, CENTER_Y + 15.0), RADIUS),
        ((PB_REAR_X, CENTER_Y - 15.0), RADIUS),
    ])
    sim.close()

    np.savez_compressed(out_dir / "signals.npz",
                        sensors=sig_s, forces=sig_f, actions=sig_a)
    with open(out_dir / "metrics.json", "w") as f:
        json.dump({"dtw_sim_v5": -1.0}, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Scene dispatcher
# ---------------------------------------------------------------------------
SCENE_RUNNERS = {
    "karman_cloak": run_karman_reproduce,
    "steady_cloak": run_steady_reproduce,
    "illusion_075L": run_illusion_reproduce,
    "illusion_1L": run_illusion_reproduce,
    "illusion_15L": run_illusion_reproduce,
    "vortex_lamb": run_vortex_reproduce,
    "vortex_taylor": run_vortex_reproduce,
}


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    parser = argparse.ArgumentParser(description="Evaluate reproduce models")
    parser.add_argument("--device-id", type=int, default=1)
    parser.add_argument("--scene", type=str, default=None,
                        help="Run single scene (e.g. karman_cloak, illusion_1L)")
    args = parser.parse_args()

    log(f"GPU: {args.device_id}")

    scenes = REPRODUCE_SCENES
    if args.scene:
        key = args.scene
        scenes = [s for s in REPRODUCE_SCENES if s["scene_id"] == key]
        if not scenes:
            log(f"ERROR: No scene matching '{args.scene}'")
            return 1

    for scene in scenes:
        scene_id = scene["scene_id"]
        out_dir = _OUT_BASE / scene_id
        runner = SCENE_RUNNERS.get(scene_id)
        if runner is None:
            log(f"WARNING: No runner for scene {scene_id}, skipping.")
            continue
        runner(scene, args.device_id, out_dir)

    log("All reproduce scenes complete.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
