# legacy_test/test_illusion_1L.py
"""Illusion 1L — legacy test.

Builds the illusion environment with LegacyCelerisLab, loads the
d1a3o14_250525_imit_1L_2U_600S PPO model (S_DIM=14), runs
deterministic inference, and compares against SR_analysis reference.

KEY: Uses REFERENCE norm and REFERENCE target_harmonics from
SR_analysis, NOT builder-computed values. The PPO model was trained
with these exact norm values.

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

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__), "..", "..", ".."))
_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
    FIFO_LEN, CONV_LEN, U0, L0, DATA_TYPE,
    _nu_from_re, _center_y, _stabilize,
    SENSOR_RADIUS, PINBALL_RADIUS,
)
from legacy_test.core.model_loader import load_model  # noqa: E402
from legacy_test.core.comparator import compare_scene  # noqa: E402
from legacy_test.core.io_helpers import save_signals  # noqa: E402
from legacy_test.core.dtw_metrics import gen_target_states_at  # noqa: E402

SAMPLE_INTERVAL = 600
ACTION_SCALE = 8.0
ACTION_BIAS = (0.0, -2.0, 2.0)
NUM_STEPS = 200
REF_DIR = os.path.join(_SRC, "SR_analysis", "data", "illusion", "illusion_1L")
OUT_DIR = os.path.join(os.path.dirname(__file__), "output", "illusion_1L")

CONFIG_DIR = os.path.join(_REPO, "configs", "legacy_configs")


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("=== Illusion 1L: Legacy Test ===")

    # Load REFERENCE norm, harmonics, and save_states (model was trained with these)
    with open(os.path.join(REF_DIR, "norm.json")) as f:
        ref_norm = json.load(f)
    with open(os.path.join(REF_DIR, "target_harmonics.json")) as f:
        target_harmonics = json.load(f)

    f_nf = float(ref_norm["force_norm_fact"])
    s_dev = np.array(ref_norm["sens_deviation"], dtype=np.float32)
    s_nf = np.array(ref_norm["sens_norm_fact"], dtype=np.float32)

    # Build the EXACT same env as legacy_env_imit.py __init__
    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=float(_nu_from_re(100.0)))

    # Phase 1: Target cylinder + sensors (record target, extract harmonics)
    ff_target = FlowField(field_cfg, cuda_cfg, device_id=args.device)
    cy = _center_y(ff_target)
    ff_target.add_cylinder((20.0 * L0, cy, 0.0), 1.0 * L0)
    for y_off in [2.0, 0.0, -2.0]:
        ff_target.add_sensor((30.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS)
    _stabilize(ff_target, 4)

    target_states = np.empty((0, 8), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff_target.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
        target_states = np.vstack((target_states, ff_target.obs.copy()[0:8]))

    # Save our target for reference
    np.savez_compressed(os.path.join(args.out, "target.npz"), target_states=target_states)
    del ff_target

    # Phase 2: Pinball + sensors (exactly matching legacy_env_imit __init__)
    ff = FlowField(field_cfg, cuda_cfg, device_id=args.device)
    cy2 = _center_y(ff)

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

    n_obj = ff.obs.size // 2
    assert n_obj == 6, f"Expected 6 objects, got {n_obj}"

    _stabilize(ff, 6)
    ff.get_ddf()
    ff.save_ddf()  # pre-bias checkpoint

    # Norm collection (from zero-action FIFO — matches ref but confirms consistency)
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, np.zeros(6, dtype=DATA_TYPE))
        fifo.append(ff.obs.copy()[0:12])

    # Bias FIFO: init_bias = [0, -1*U0, 1*U0] (matching legacy line 143)
    ff.apply_ddf()  # restore pre-bias
    init_bias = (0.0, -1.0, 1.0)
    bias_arr = np.zeros(6, dtype=DATA_TYPE)
    bias_arr[3] = float(init_bias[0] * U0)
    bias_arr[4] = float(init_bias[1] * U0)
    bias_arr[5] = float(init_bias[2] * U0)

    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, bias_arr)
        fifo.append(ff.obs.copy()[0:12])

    save_states = np.array(list(fifo), dtype=DATA_TYPE)

    # CRITICAL: save DDF AFTER bias FIFO (matching legacy line 147-148)
    ff.get_ddf()
    ff.save_ddf()

    log(f"  ref force_norm_fact = {f_nf:.6f}")
    log(f"  ref sens_deviation = {s_dev}")

    model = load_model("d1a3o14_250525_imit_1L_2U_600S")
    log("Model loaded on CPU")

    # DRL inference: reset goes to POST-bias state (legacy save_ddf on line 148)
    ff.restore_ddf()
    ff.apply_ddf()
    fifo = deque(maxlen=FIFO_LEN)
    for row in save_states:
        fifo.append(row.copy())

    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)

    # Build initial observation using REFERENCE norm
    raw = ff.obs.copy()[0:12]
    forces_norm = raw[6:12] / f_nf
    sens_norm = (raw[0:6] - s_dev) / s_nf
    obs_12 = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)
    tgt = gen_target_states_at(0, target_harmonics)
    obs = np.clip(np.hstack([obs_12, [tgt[0] / f_nf, tgt[1] / f_nf]]), -1.0, 1.0).astype(np.float32)

    for step in range(NUM_STEPS):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        sig_a[step] = action.copy()

        action_arr = np.zeros(6, dtype=DATA_TYPE)
        action_arr[3:] = (action * ACTION_SCALE + np.array(ACTION_BIAS, dtype=np.float32)) * U0

        ff.context.push()
        try:
            ff.run(SAMPLE_INTERVAL, action_arr)
        finally:
            ff.context.pop()

        raw = ff.obs.copy()[0:12]
        fifo.append(raw)
        sig_s[step] = raw[0:6]
        sig_f[step] = raw[6:12]

        forces_norm = raw[6:12] / f_nf
        sens_norm = (raw[0:6] - s_dev) / s_nf
        obs_12 = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)
        tgt = gen_target_states_at(step + 1, target_harmonics)
        obs = np.clip(np.hstack([obs_12, [tgt[0] / f_nf, tgt[1] / f_nf]]), -1.0, 1.0).astype(np.float32)

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

    log("Comparing against reference...")
    result = compare_scene(REF_DIR, sig_s, sig_f, sig_a, conv_len=36, label="illusion_1L")
    with open(os.path.join(args.out, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    log(f"PASS" if result["passed"] else "FAIL")
    del ff
    return 0 if result["passed"] else 1


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