# legacy_test/test_illusion_remaining.py
"""Illusion 0.75L and 1.5L — legacy tests.

Usage: conda run -n pycuda_3_10 python test_illusion_remaining.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 legacy_test.core.legacy_env_builder import (  # noqa: E402
    build_illusion, FIFO_LEN, CONV_LEN, U0, DATA_TYPE,
)
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, save_target, save_norm  # noqa: E402
from legacy_test.core.dtw_metrics import gen_target_states_at  # noqa: E402

ACTION_SCALE = 8.0
ACTION_BIAS = (0.0, -2.0, 2.0)
NUM_STEPS = 200
OUT_BASE = os.path.join(os.path.dirname(__file__), "output")

# Configs: (label, diameter_L, sample_interval, model_name, ref_subdir)
ILLUSION_CASES = [
    ("illusion_075L", 0.75, 400, "d1a3o14_250525_imit_075L_2U_400S", "illusion_0.75L"),
    ("illusion_15L",  1.50, 800, "d1a3o14_250525_imit_15L_2U",       "illusion_1.5L"),
]


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


def run_one(device_id: int, label: str, diam_L: float, si: int,
            model_name: str, ref_subdir: str) -> dict:
    log(f"=== {label}: Legacy Test ===")
    ref_dir = os.path.join(_SRC, "SR_analysis", "data", "illusion", ref_subdir)
    out_dir = os.path.join(OUT_BASE, label)
    os.makedirs(out_dir, exist_ok=True)

    data = build_illusion(device_id=device_id, target_diameter_L=diam_L,
                          sample_interval=si, action_bias=ACTION_BIAS)
    ff = data["flow_field"]
    target_harmonics = data["target_harmonics"]
    norm = data["norm"]
    n_obj = norm.get("n_obj_total", 6)
    f_nf = float(norm["force_norm_fact"])
    s_dev = np.array(norm["sens_deviation"], dtype=np.float32)
    s_nf = np.array(norm["sens_norm_fact"], dtype=np.float32)

    save_target(out_dir, data["target_states"]); save_norm(out_dir, norm)

    model = load_model(model_name)
    log(f"  Model: {model_name}")

    ff.restore_ddf(); ff.apply_ddf()
    init_bias = (0.0, -1.0, 1.0)
    bias_arr = np.zeros(n_obj, 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)

    for _ in range(FIFO_LEN):
        ff.run(si, bias_arr)

    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)

    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)
    tcd = tgt[0] / f_nf if f_nf > 1e-12 else 0.0
    tcl = tgt[1] / f_nf if f_nf > 1e-12 else 0.0
    obs = np.clip(np.hstack([obs_12, [tcd, tcl]]), -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(n_obj, dtype=DATA_TYPE)
        action_arr[3:] = (action * ACTION_SCALE + np.array(ACTION_BIAS, dtype=np.float32)) * U0

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

        raw = ff.obs.copy()[0:12]
        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)
        tcd = tgt[0] / f_nf if f_nf > 1e-12 else 0.0
        tcl = tgt[1] / f_nf if f_nf > 1e-12 else 0.0
        obs = np.clip(np.hstack([obs_12, [tcd, tcl]]), -1.0, 1.0).astype(np.float32)

    save_signals(out_dir, sig_s, sig_f, sig_a)
    np.savez_compressed(os.path.join(out_dir, "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=label)
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    log(f"  {'PASS' if result['passed'] else 'FAIL'}")
    del ff
    return result


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--device", type=int, default=0)
    args = ap.parse_args()
    results = {}
    for label, diam_L, si, model_name, ref_subdir in ILLUSION_CASES:
        results[label] = run_one(args.device, label, diam_L, si, model_name, ref_subdir)

    log("\n=== Illusion Remaining Summary ===")
    for name, r in results.items():
        log(f"  {name}: DTW={r['dtw_sim']:.4f}, act_corr={[f'{c:.3f}' for c in r['action_corr']]} -> {'PASS' if r['passed'] else 'FAIL'}")


if __name__ == "__main__":
    main()
