#!/usr/bin/env python3
# reproduce/phase3_reproduce.py
"""Phase 3: DRL inference with legacy-compatible config + reference comparison.

Uses config_lbm_pinball_legacy_compat.json (regularized inlet, NBB equivalent)
instead of the default config_lbm_pinball.json. After inference, compares
output against SR_analysis reference data.

Scenes: karman_re100, steady_cloak, illusion_1L, vortex_lamb

Usage:
    conda run -n pycuda_3_10 python phase3_reproduce.py --device 0
    conda run -n pycuda_3_10 python phase3_reproduce.py --device 0 --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

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

_REPO = str(Path(__file__).resolve().parents[3])
_SRC = Path(_REPO) / "src"
_DRL = _SRC / "drl_pinball"
for p in [_REPO, str(_SRC), str(_DRL)]:
    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  # noqa: E402
from CelerisLab.common.render import compute_vorticity, render_vorticity_field  # noqa: E402
from CelerisLab.lbm.initializers import add_vortex  # noqa: E402
from drl_pinball.reproduce.configs.model_inventory import ModelInventory  # noqa: E402
from reproduce.core.drl_comparator import compare_reproduce_output  # noqa: E402

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
LEGACY_COMPAT_CFG = "configs/config_lbm_pinball_legacy_compat.json"
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)

# Standard geometry
DIST_X = 10.0 * L0
PB_FRONT_X = 30.0 * L0
PB_REAR_X = 31.3 * L0
SENSOR_X = 40.0 * L0

# Illusion geometry
ILL_PB_FRONT_X = 19.0 * L0
ILL_PB_REAR_X = 20.3 * L0
ILL_SENSOR_X = 30.0 * L0
ILL_TARGET_X = 20.0 * L0

SR_DATA = _SRC / "SR_analysis" / "data"
_THIS_DIR = Path(__file__).resolve().parent
OUT_BASE = _THIS_DIR / "output" / "phase3"


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)


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(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 read_obs_karman(sim, dist_id, sensor_ids, pinball_ids, cc):
    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)


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)


# ---------------------------------------------------------------------------
# Karman Cloak Re100
# ---------------------------------------------------------------------------
def run_karman(device_id: int, out_dir: Path) -> None:
    log("=== Phase 3: Karman Cloak Re100 ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = 800
    num_steps = 200
    scale, bias = 8.0, (0.0, -4.0, 4.0)
    ref_dir = str(SR_DATA / "karman" / "karman_re100")
    norm = load_legacy_norm(ref_dir)

    # Phase 1: Disturbance + sensors, record target
    sim = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, 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 CC: {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)

    # Phase 2: Add 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)

    # Bias FIFO
    bias_norm = np.array([0.0, -1.0, 1.0])
    bias_omega = action_to_omega(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
    sim.restore()
    ema.reset(bias_omega.copy())

    model = ModelInventory().load("d1a3o12_re100", device="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)

    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(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]
        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)

    # Compare against reference
    log("  Comparing against SR_analysis reference...")
    result = compare_reproduce_output(ref_dir, str(out_dir), label="karman_re100", conv_len=30)
    with open(out_dir / "result.json", "w") as f:
        json.dump(result, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Steady Cloak
# ---------------------------------------------------------------------------
def run_steady(device_id: int, out_dir: Path) -> None:
    log("=== Phase 3: Steady Cloak ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = 800; num_steps = 200
    surf_vel = (0.0, -5.1, 5.1)
    bias_surf = np.array(surf_vel, dtype=np.float32) * U0

    sim = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, 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])

    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))

    ref_dir = str(SR_DATA / "steady" / "steady")
    result = compare_reproduce_output(ref_dir, str(out_dir), label="steady", conv_len=30)
    with open(out_dir / "result.json", "w") as f:
        json.dump(result, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Illusion 1L
# ---------------------------------------------------------------------------
def run_illusion(device_id: int, out_dir: Path, diam_L: float = 1.0, si: int = 600) -> None:
    # diam_L: 0.75 → "illusion_0.75L", 1.0 → "illusion_1L", 1.5 → "illusion_1.5L"
    if diam_L == int(diam_L):
        label_suffix = str(int(diam_L))
    else:
        label_suffix = str(diam_L).rstrip('0')
    label = f"illusion_{label_suffix}L"
    # Map diameter to model name
    model_map = {
        0.75: "d1a3o14_250525_imit_075L_2U_400S",
        1.0:  "d1a3o14_250525_imit_1L_2U_600S",
        1.5:  "d1a3o14_250525_imit_15L_2U",
    }
    model_name = model_map[diam_L]
    log(f"=== Phase 3: Illusion {label} ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    scale, bias = 8.0, (0.0, -2.0, 2.0)
    ref_dir = str(SR_DATA / "illusion" / label)
    norm = load_legacy_norm(ref_dir)

    # Load target harmonics
    with open(os.path.join(ref_dir, "target_harmonics.json")) as f:
        target_harmonics = json.load(f)

    def gen_target_at(t):
        D = len(target_harmonics)
        vals = np.zeros(D, dtype=np.float32)
        for d, h in enumerate(target_harmonics):
            val = float(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)
            vals[d] = val
        return vals

    # Phase 1: Record target on new CFD
    sim_t = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, device_id=device_id)
    sim_t.add_body("circle", center=(ILL_TARGET_X, CENTER_Y, 0.0), radius=diam_L * L0)
    s_ids_t = [
        sim_t.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim_t.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim_t.add_body("sensor", center=(ILL_SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim_t.initialize(); sim_t.run(WARMUP, zero_obs=True)
    cc_t = get_cc(sim_t, s_ids_t[0])
    target = np.zeros((FIFO_LEN, 8), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim_t.run(si, zero_obs=True)
        f = list(sim_t.read_force(0, normalize=True))
        s = [float(sim_t.read_sensor(sid, normalize=True)[d]) * cc_t for sid in s_ids_t for d in range(2)]
        target[i] = np.hstack([f, s])
    sim_t.close()
    np.savez_compressed(out_dir / "target.npz", target_states=target)

    # Phase 2: Pinball + sensors
    sim = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, 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])
    log(f"  Sensor CC: {cc}")

    # Bias FIFO (init bias = [0, -1, 1] * U0)
    init_bias_surf = np.array([0.0, -1.0, 1.0]) * U0
    init_bias_omega = -init_bias_surf / RADIUS
    ema = ActionSmoother(weight=0.1); ema.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema(init_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(action_to_omega(np.array([0.0, -1.0, 1.0]), scale=scale, bias=bias))

    # DRL inference
    model = ModelInventory().load(model_name, device="cpu")
    obs_init = read_obs_6obj(sim, sensor_ids, [3, 4, 5], cc)
    obs_12 = normalize_obs(obs_init, norm)
    target_cd = (gen_target_at(0)[0] - norm["sens_deviation"][0]) / norm["sens_norm_fact"][0]
    target_cl = (gen_target_at(0)[1] - norm["sens_deviation"][1]) / norm["sens_norm_fact"][1]
    obs_norm = np.clip(np.hstack([obs_12, [target_cd, target_cl]]), -1.0, 1.0).astype(np.float32)

    num_steps = 200
    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(action, scale=scale, bias=bias)
        smoothed = ema(target_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]
        sig_a[step] = action

        obs_12 = normalize_obs(obs, norm)
        tgt = gen_target_at(step)
        target_cd = (tgt[0] - norm["sens_deviation"][0]) / norm["sens_norm_fact"][0]
        target_cl = (tgt[1] - norm["sens_deviation"][1]) / norm["sens_norm_fact"][1]
        obs_norm = np.clip(np.hstack([obs_12, [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)
    result = compare_reproduce_output(ref_dir, str(out_dir), label=label, conv_len=36)
    with open(out_dir / "result.json", "w") as f:
        json.dump(result, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Vortex Lamb
# ---------------------------------------------------------------------------
def run_vortex(device_id: int, out_dir: Path, vortex_type: str = "lamb") -> None:
    log(f"=== Phase 3: Vortex {vortex_type} ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    SI = 800; num_steps = 150; scale, bias = 4.0, (0.0, -4.0, 4.0)
    ref_dir = str(SR_DATA / "vortex" / f"vortex_{vortex_type}")
    norm = load_legacy_norm(ref_dir)
    strength = 0.5 * U0 if vortex_type == "lamb" else 0.03 * U0

    # Phase 1: Sensors only + vortex, record target
    sim_t = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, device_id=device_id)
    s_ids_t = [
        sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim_t.initialize(); sim_t.run(WARMUP, zero_obs=True)
    sim_t.field.download_ddf()
    add_vortex(sim_t.field, (10.0 * L0, CENTER_Y), 2.0 * L0, strength, vortex_type)
    cc_t = get_cc(sim_t, s_ids_t[0])

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim_t.run(SI, zero_obs=True)
        for j, sid in enumerate(s_ids_t):
            s = sim_t.read_sensor(sid, normalize=True)
            target[i, j * 2] = float(s[0]) * cc_t
            target[i, j * 2 + 1] = float(s[1]) * cc_t
    sim_t.close()
    np.savez_compressed(out_dir / "target.npz", target_states=target)

    # Phase 2: Pinball + sensors + vortex
    sim = Simulation(lbm_config_path=LEGACY_COMPAT_CFG, 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)
    sim.field.download_ddf()
    add_vortex(sim.field, (15.0 * L0, CENTER_Y), 2.0 * L0, strength, vortex_type)
    cc = get_cc(sim, sensor_ids[0])
    log(f"  Sensor CC: {cc}")

    # Bias FIFO
    bias_omega = action_to_omega(np.array([0.0, -1.0, 1.0]), 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(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())

    # DRL inference
    model = ModelInventory().load(f"vortex_{vortex_type}", device="cpu")
    obs_init = read_obs_6obj(sim, sensor_ids, [3, 4, 5], 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(action, scale=scale, bias=bias)
        smoothed = ema(target_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]
        sig_a[step] = action
        obs_norm = normalize_obs(obs, norm)

    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)
    result = compare_reproduce_output(ref_dir, str(out_dir), label=f"vortex_{vortex_type}", conv_len=30)
    with open(out_dir / "result.json", "w") as f:
        json.dump(result, f, indent=2)
    log("  Done.")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Phase 3: DRL reproduce with legacy-compat config")
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--scene", type=str, default="all",
                    help="Scene: karman, steady, illusion_1L, vortex_lamb, vortex_taylor, all")
    args = ap.parse_args()

    all_scenes = {
        "karman": lambda: run_karman(args.device, OUT_BASE / "karman"),
        "steady": lambda: run_steady(args.device, OUT_BASE / "steady"),
        "illusion_1L": lambda: run_illusion(args.device, OUT_BASE / "illusion_1L", 1.0, 600),
        "vortex_lamb": lambda: run_vortex(args.device, OUT_BASE / "vortex_lamb", "lamb"),
        "vortex_taylor": lambda: run_vortex(args.device, OUT_BASE / "vortex_taylor", "taylor"),
    }

    if args.scene == "all":
        for name, func in all_scenes.items():
            func()
    else:
        for s in args.scene.split(","):
            s = s.strip()
            if s in all_scenes:
                all_scenes[s]()

    log("\nPhase 3 complete.")
