# OID_analysis/scripts/collect_baseline_forces.py
"""
Quick collection of missing baseline forces and norm data.
For scenes where fields already exist but forces/norm are missing.

Usage:
    # pinball_baseline forces (needed for steady+illusion baselines)
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_baseline_forces.py \
        --scene pinball_baseline --device 3
"""
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
from OID_analysis.configs import get_scene, data_dir_for_scene, LEGACY_CFG_DIR
from OID_analysis.utils.cfd_interface import load_legacy_configs

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


def collect_pinball_forces(device_id: int, n_steps: int):
    """Collect forces.npz and norm.json for pinball_baseline."""
    cfg = get_scene("pinball_baseline")
    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    action_bias = cfg["action_bias"]
    out_dir = data_dir_for_scene("pinball_baseline")

    has_forces = os.path.isfile(os.path.join(out_dir, "forces.npz"))
    has_norm = os.path.isfile(os.path.join(out_dir, "norm.json"))
    if has_forces and has_norm:
        print(f"  Both forces.npz and norm.json exist, skipping")
        return

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    # Build env
    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor((40.0 * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)
    ff.add_cylinder((30.0 * L0, CENTER_Y, 0.0), L0 / 2.0)
    ff.add_cylinder((31.3 * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((31.3 * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)

    n_obj = 6
    n_stab = int(4 * ff.FIELD_SHAPE[0] / u0)
    print(f"  Stabilising {n_stab} steps...")
    ff.run(n_stab, np.zeros(n_obj, dtype=DATA_TYPE))

    if not has_norm:
        # Norm collection
        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()[0:12])

        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(),
            "action_bias": list(action_bias),
        }
        with open(os.path.join(out_dir, "norm.json"), "w") as f:
            json.dump(norm, f, indent=2)
        print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    if not has_forces:
        forces = []
        print(f"  Recording {n_steps} steps of forces...")
        for step in range(n_steps):
            ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
            obs_slice = ff.obs.copy()[0:12]
            forces.append(obs_slice[6:12])

        np.savez(os.path.join(out_dir, "forces.npz"),
                 forces=np.array(forces, dtype=np.float32))
        print(f"  Forces saved: {len(forces)} steps")

    del ff
    print(f"  Done")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, default="pinball_baseline")
    ap.add_argument("--device", type=int, default=3)
    ap.add_argument("--steps", type=int, default=500)
    args = ap.parse_args()

    if args.scene == "pinball_baseline":
        t0 = time.time()
        collect_pinball_forces(args.device, args.steps)
        print(f"Done in {time.time() - t0:.1f}s")
    else:
        print(f"Unknown scene: {args.scene}")


if __name__ == "__main__":
    main()
