# legacy_test/test_steady_cloak.py
"""Steady Cloak — legacy test (open-loop, no DRL).

Applies constant rear-cylinder rotation [0, -5.1, 5.1]*U0 to suppress
vortex shedding. Matches the OID_analysis/scripts/collect_steady_cloak.py
collection procedure exactly.

No DTW comparison — steady cloak was never benchmarked with DTW in
SR_analysis (similarity: ---). This test produces controlled.npz for
field analysis (CCD/OID).

Usage: conda run -n pycuda_3_10 python test_steady_cloak.py --device 0
"""

import argparse
import json
import os
import sys
import time

import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
_SRC = os.path.join(_REPO, "src")
_DRL = os.path.join(_SRC, "drl_pinball")
for p in [_REPO, _SRC, _DRL]:
    if p not in sys.path:
        sys.path.insert(0, p)

from LegacyCelerisLab import FlowField  # noqa: E402
from LegacyCelerisLab import utils as legacy_utils  # noqa: E402

from legacy_test.core.legacy_env_builder import (  # noqa: E402
    L0, U0, DATA_TYPE, FIFO_LEN,
    _center_y, _stabilize,
    SENSOR_RADIUS, PINBALL_RADIUS,
)

CONFIG_DIR = os.path.join(_REPO, "configs", "legacy_configs")
SAMPLE_INTERVAL = 800
ACTION_BIAS = (0.0, -5.1, 5.1)
NUM_STEPS = 200
OUT_DIR = os.path.join(os.path.dirname(__file__), "output", "steady_cloak")


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


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--device", type=int, default=0)
    ap.add_argument("--out", type=str, default=OUT_DIR)
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    log("=== Steady Cloak: Legacy Test (open-loop) ===")

    # Build pinball + sensors env (no disturbance cylinder)
    cuda_cfg = legacy_utils.load_cuda_config(os.path.join(CONFIG_DIR, "config_cuda.json"))
    field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json"))
    field_cfg = field_cfg._replace(viscosity=0.004)

    ff = FlowField(field_cfg, cuda_cfg, device_id=args.device)
    cy = _center_y(ff)

    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor((40.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS)  # id=0,1,2
    ff.add_cylinder((30.0 * L0, cy, 0.0), PINBALL_RADIUS)                 # front, id=3
    ff.add_cylinder((31.3 * L0, cy + 0.75 * L0, 0.0), PINBALL_RADIUS)     # top, id=4
    ff.add_cylinder((31.3 * L0, cy - 0.75 * L0, 0.0), PINBALL_RADIUS)     # bottom, id=5

    n_obj = ff.obs.size // 2
    assert n_obj == 6

    _stabilize(ff, 6)

    # Record target: clean channel (sensors only, no pinball)
    ff_clean = FlowField(field_cfg, cuda_cfg, device_id=args.device)
    cy_c = _center_y(ff_clean)
    for y_off in [2.0, 0.0, -2.0]:
        ff_clean.add_sensor((40.0 * L0, cy_c + y_off * L0, 0.0), SENSOR_RADIUS)
    _stabilize(ff_clean, 3)
    target = np.empty((0, 6), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff_clean.run(SAMPLE_INTERVAL, np.zeros(3, dtype=DATA_TYPE))
        target = np.vstack((target, ff_clean.obs.copy()[0:6]))
    np.savez_compressed(os.path.join(args.out, "target.npz"), target_states=target)
    del ff_clean

    # Apply constant rear-cylinder rotation (matching OID analysis collector)
    # front=0, bottom=-5.1*U0, top=5.1*U0
    action_arr = np.zeros(n_obj, dtype=DATA_TYPE)
    action_arr[3] = 0.0                        # front
    action_arr[4] = float(ACTION_BIAS[1] * U0)  # top = -5.1*U0
    action_arr[5] = float(ACTION_BIAS[2] * U0)  # bottom = 5.1*U0
    log(f"  Rotation: front=0, top={action_arr[4]:.6f}, bottom={action_arr[5]:.6f}")

    # Let steady cloak stabilize (matching OID: 100 SI steps)
    for _ in range(100):
        ff.run(SAMPLE_INTERVAL, action_arr)

    # Record signals
    sig_s = np.zeros((NUM_STEPS, 6), dtype=np.float32)
    sig_f = np.zeros((NUM_STEPS, 6), dtype=np.float32)
    for s in range(NUM_STEPS):
        ff.run(SAMPLE_INTERVAL, action_arr)
        obs = ff.obs.copy()[0:12]
        sig_s[s] = obs[0:6]
        sig_f[s] = obs[6:12]

    save_actions = np.zeros((NUM_STEPS, 3), dtype=np.float32)
    np.savez_compressed(os.path.join(args.out, "controlled.npz"),
                        sensors=sig_s, forces=sig_f, actions=save_actions,
                        rewards=np.zeros(NUM_STEPS, dtype=np.float32))

    # Check force balance (steady cloak should suppress lift oscillations)
    front_fy_mean = float(np.mean(sig_f[:, 1]))
    top_fy_mean = float(np.mean(sig_f[:, 3]))
    bot_fy_mean = float(np.mean(sig_f[:, 5]))
    lift_rms = float(np.sqrt(np.mean(sig_f[:, 1]**2 + sig_f[:, 3]**2 + sig_f[:, 5]**2)))

    log(f"  Forces: front_fy={front_fy_mean:+.6f}, top_fy={top_fy_mean:+.6f}, bottom_fy={bot_fy_mean:+.6f}")
    log(f"  Lift RMS: {lift_rms:.6f}")

    # Check: lift oscillations should be near zero (successful cloaking)
    # Uncontrolled: ~0.05; Controlled: ~0.005 (10x reduction)
    passed = lift_rms < 0.01
    log(f"  {'PASS' if passed else 'FAIL'} (lift RMS={lift_rms:.6f} < 0.01)")

    if passed:
        result = {"passed": True, "lift_rms": float(lift_rms)}
    else:
        result = {"passed": False, "lift_rms": float(lift_rms)}

    with open(os.path.join(args.out, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    del ff
    return 0 if passed else 1


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