DynamisLab/src/OID_analysis/scripts/collect_controlled.py

# OID_analysis/scripts/collect_controlled.py
"""
Collect DRL-controlled rollout (q_ctl) for Karman cloak and illusion scenes.
Generates field time series for Delta-q_ctl computation.

Usage:
    # Karman:
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \
        --scene karman_re100 --device 1 --steps 500

    # Illusion (3 diameters):
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \
        --scene illusion_1.0L --device 3 --steps 500
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \
        --scene illusion_0.75L --device 3 --steps 500
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \
        --scene illusion_1.5L --device 3 --steps 500
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque

import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from LegacyCelerisLab import FlowField  # noqa: E402
from OID_analysis.utils.cfd_interface import (  # noqa: E402
    load_legacy_configs, get_velocity_field, save_vorticity_png, vorticity_from_ddf,
    load_ppo_model, scale_action, build_observation, compute_similarity,
    calc_lag, calc_dtw_sim, analyze_harmonics, gen_target_states_at,
)
from OID_analysis.configs import (  # noqa: E402
    get_scene, get_scene_list, model_path_for_scene, LEGACY_CFG_DIR,
)

DATA_TYPE = np.float32
L0 = 20.0
CENTER_Y = (512 - 1) / 2.0
FIFO_LEN = 150
CONV_LEN_DEFAULT = 30
CONV_LEN_ILLUSION = 36


def collect_single(scene_name: str, device_id: int, n_steps: int) -> dict:
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    ac_scale = cfg["action_scale"]
    ac_bias = cfg["action_bias"]
    n_obj = cfg["n_objects_env"]
    s_dim = cfg["s_dim"]
    source = cfg.get("source", "")

    out_dir = data_dir_for_scene(scene_name)
    os.makedirs(out_dir, exist_ok=True)

    # Check DDF+FIFO checkpoint exists
    ddf_ckpt_path = os.path.join(out_dir, "ddf_checkpoint.npy")
    fifo_ckpt_path = os.path.join(out_dir, "fifo_checkpoint.npy")
    has_blk = os.path.isfile(ddf_ckpt_path) and os.path.isfile(fifo_ckpt_path)

    # Load legacy configs
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(u0))

    if not has_blk:
        print(f"  No DDF checkpoint found, building env from scratch ...")
        ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

        if cfg["has_disturbance"]:
            # Karman layout: dist_cyl first
            ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
            for y_off in [2.0, 0.0, -2.0]:
                ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)
            n_phase1 = 4
            ff.run(int(4 * 1280 / u0), np.zeros(n_phase1, dtype=DATA_TYPE))
        else:
            # Illusion layout: sensors first
            for y_off in [2.0, 0.0, -2.0]:
                ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)

        # Add pinball
        ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0)
        ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
        ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)

        ff.run(int(4 * 1280 / u0), np.zeros(n_obj, dtype=DATA_TYPE))
        ff.get_ddf()
        ff.save_ddf()

        # Norm
        obs_slice_start = cfg["obs_slice"][0]
        obs_slice_end = cfg["obs_slice"][1]
        fifo = deque(maxlen=FIFO_LEN)
        for _ in range(FIFO_LEN):
            ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
            fifo.append(ff.obs.copy()[obs_slice_start:obs_slice_end])
        temp = np.array(fifo, dtype=DATA_TYPE)
        force_norm_fact = 6.0 * float(np.max(np.abs(temp[:, 6:12])))
        sens_deviation = np.mean(temp[:, 0:6], axis=0).astype(DATA_TYPE)
        sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
        for i in range(6):
            sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp[:, i] - sens_deviation[i])))

        norm = {"force_norm_fact": force_norm_fact,
                "sens_deviation": sens_deviation.tolist(),
                "sens_norm_fact": sens_norm_fact.tolist()}

        # Preset-action FIFO (matches legacy env)
        ff.apply_ddf()
        bias_arr = np.zeros(n_obj, dtype=DATA_TYPE)
        if cfg["has_disturbance"]:
            bias_arr[4] = -4.0 * u0
            bias_arr[5] = 4.0 * u0
        else:
            bias_arr[4] = -1.0 * u0
            bias_arr[5] = 1.0 * u0
        fifo.clear()
        for _ in range(FIFO_LEN):
            ff.run(si, bias_arr)
            fifo.append(ff.obs.copy()[obs_slice_start:obs_slice_end])
        save_states_arr = np.array(fifo, dtype=DATA_TYPE)

        # Save checkpoint
        ff.get_ddf()
        np.save(ddf_ckpt_path, ff.ddf)
        np.save(fifo_ckpt_path, save_states_arr)

        with open(os.path.join(out_dir, "norm.json"), "w") as f:
            json.dump(norm, f, indent=2)
        print(f"  Checkpoint saved to {out_dir}")
    else:
        print(f"  Loading DDF+FIFO checkpoint from {out_dir}")
        # Load norm
        with open(os.path.join(out_dir, "norm.json")) as f:
            norm = json.load(f)
        # Rebuild env to get a fresh FlowField
        ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
        if cfg["has_disturbance"]:
            ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
            for y_off in [2.0, 0.0, -2.0]:
                ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)
        else:
            for y_off in [2.0, 0.0, -2.0]:
                ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)
        ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0)
        ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
        ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)

        # Restore DDF
        ff.ddf = np.load(ddf_ckpt_path)
        ff.apply_ddf()
        print(f"  DDF checkpoint restored")

    # Save config
    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump({k: str(v) if not isinstance(v, (int, float, list, bool)) else v
                   for k, v in cfg.items()}, f, indent=2)

    # ---- Target signals (needed for s_dim=14 illusion) ----
    target_states = None
    target_harmonics = None
    if cfg_sid == "illusion":
        target_path = os.path.join(out_dir, "target.npz")
        harm_path = os.path.join(out_dir, "target_harmonics.json")
        if os.path.isfile(target_path) and os.path.isfile(harm_path):
            target_data = np.load(target_path)
            target_states = target_data["target_states"]
            with open(harm_path) as f:
                target_harmonics = json.load(f)
            print(f"  Target loaded: {target_states.shape}")
        else:
            print(f"  WARNING: no target found at {target_path}")

    # ---- PPO inference ----
    obs_slice_start = cfg["obs_slice"][0]
    obs_slice_end = cfg["obs_slice"][1]

    # Load checkpoint FIFO state
    load_state = np.load(fifo_ckpt_path)
    fifo = deque(maxlen=FIFO_LEN)
    for s in load_state:
        fifo.append(s)

    model_path = model_path_for_scene(scene_name)
    if model_path is None:
        raise ValueError(f"No model path for {scene_name}")

    print(f"  Loading model: {model_path}  (s_dim={s_dim})")
    model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=s_dim, a_dim=3)
    model.set_random_seed(19)

    obs = np.zeros(s_dim, dtype=np.float32)
    sens_c, forc_c, act_c, ux_list, uy_list = [], [], [], [], []

    for step in range(n_steps):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        act_c.append(action.copy())

        # Build omega array
        temp = np.zeros(n_obj, dtype=DATA_TYPE)
        omega = (action * ac_scale + np.array(ac_bias, dtype=np.float32)) * u0
        temp[n_obj - 3:] = omega

        ff.context.push()
        ff.run(si, temp)
        ff.context.pop()

        obs_slice = ff.obs.copy()[obs_slice_start:obs_slice_end]
        fifo.append(obs_slice)
        sens_c.append(obs_slice[0:6])
        forc_c.append(obs_slice[6:12])

        # Build next observation
        forces_norm = obs_slice[6:12] / norm["force_norm_fact"]
        sens_norm = (obs_slice[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"]

        if s_dim == 14 and target_harmonics is not None:
            target_vals = gen_target_states_at(step, target_harmonics)
            t_cd_n = float(target_vals[0]) / norm["force_norm_fact"]
            t_cl_n = float(target_vals[1]) / norm["force_norm_fact"]
            obs = np.clip(np.hstack([forces_norm, sens_norm, t_cd_n, t_cl_n]), -1.0, 1.0).astype(np.float32)
        else:
            obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)

        # Save field every step (for Delta-q_ctl POD)
        ux, uy = get_velocity_field(ff, u0=u0)
        ux_list.append(ux)
        uy_list.append(uy)

    # Save
    sens_arr = np.array(sens_c, dtype=np.float32)
    forc_arr = np.array(forc_c, dtype=np.float32)
    act_arr = np.array(act_c, dtype=np.float32)

    # Compute similarity
    conv_len = cfg.get("conv_len", CONV_LEN_DEFAULT)
    if target_states is not None:
        if cfg_sid == "karman":
            sim = compute_similarity(target_states, sens_arr, conv_len)
        elif cfg_sid == "illusion":
            # For illusion, target_states[:, 2:8] has the sensor references
            target_sensors = target_states[:, 2:8] if target_states.shape[1] >= 8 else target_states
            sim = compute_similarity(target_sensors, sens_arr, conv_len)
        else:
            sim = 0.0
        print(f"  similarity = {sim:.4f}")

    np.savez(os.path.join(out_dir, "controlled.npz"),
             sensors=sens_arr, forces=forc_arr, actions=act_arr)
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.stack(ux_list), uy=np.stack(uy_list))

    omega_viz = vorticity_from_ddf(ff, u0=u0)
    save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"),
                       omega_viz, title=f"{scene_name} controlled")

    result = {"scene": scene_name, "n_steps": n_steps,
              "similarity": float(sim) if target_states is not None else 0.0}
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    del ff, model
    print(f"  Saved {n_steps} snapshots to {out_dir}")
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, required=True,
                    help="Scene name: karman_re100, illusion_0.75L, illusion_1.0L, illusion_1.5L")
    ap.add_argument("--device", type=int, default=3)
    ap.add_argument("--steps", type=int, default=500)
    args = ap.parse_args()

    all_scenes = get_scene_list()
    if args.scene not in all_scenes:
        print(f"Unknown scene: {args.scene}. Available PPO scenes: "
              f"{[s for s in all_scenes if get_scene(s).get('source') == 'PPO_inference']}")
        return 1

    t0 = time.time()
    r = collect_single(args.scene, args.device, args.steps)
    print(f"Done in {time.time() - t0:.1f}s: sim={r.get('similarity', 0):.4f}")


if __name__ == "__main__":
    main()