# OID_analysis/scripts/collect_all_data.py
"""
Batch data collection for all OID scenes.
Handles baseline + controlled + target scenes efficiently.

Usage:
    # Check all data (no GPU needed)
    python3 src/OID_analysis/scripts/collect_all_data.py --check

    # Collect all data (GPU required)
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_all_data.py \
        --device-steady 3 --device-karman 1 --device-illusion 3
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time

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 OID_analysis.configs import (  # noqa: E402
    get_scene, data_dir_for_scene, model_path_for_scene,
    SCENES, DATA_DIR,
)


SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


def check_data() -> bool:
    """Check which data files exist."""
    all_ok = True
    required_files = {
        "empty_channel": ["fields.npz", "sensors.npz"],
        "pinball_baseline": ["fields.npz", "sensors.npz", "forces.npz", "norm.json"],
        "disturbance_only": ["fields.npz", "sensors.npz", "target.npz"],
        "karman_blk": ["fields.npz", "sensors.npz", "forces.npz", "norm.json"],
        "karman_re100": ["fields.npz", "controlled.npz", "norm.json"],
        "steady_cloak": ["fields.npz", "sensors.npz", "forces.npz"],
        "illusion_0.75L": ["fields.npz", "controlled.npz", "norm.json"],
        "illusion_1.0L": ["fields.npz", "controlled.npz", "norm.json"],
        "illusion_1.5L": ["fields.npz", "controlled.npz", "norm.json"],
        "target_cylinder_0.75L": ["fields.npz", "target.npz", "target_harmonics.json"],
        "target_cylinder_1.0L": ["fields.npz", "target.npz", "target_harmonics.json"],
        "target_cylinder_1.5L": ["fields.npz", "target.npz", "target_harmonics.json"],
    }

    for scene_name, files in required_files.items():
        dd = data_dir_for_scene(scene_name)
        missing = [f for f in files if not os.path.isfile(os.path.join(dd, f))]
        if missing:
            print(f"  MISSING [{scene_name}]: {missing}")
            all_ok = False
        else:
            sizes = {f: os.path.getsize(os.path.join(dd, f)) for f in files}
            total_mb = sum(sizes.values()) / 1e6
            print(f"  OK [{scene_name:30s}] {total_mb:.0f} MB  {list(sizes.keys())}")

    return all_ok


def run_script(script_name: str, args: list, description: str = ""):
    """Run a collection script."""
    script_path = os.path.join(SCRIPT_DIR, script_name)
    cmd = [sys.executable, script_path] + [str(a) for a in args]
    print(f"\n{'='*60}")
    print(f"  {description}")
    print(f"  {' '.join(cmd)}")
    print(f"{'='*60}")
    t0 = time.time()
    result = subprocess.run(cmd, capture_output=False, text=True)
    elapsed = time.time() - t0
    if result.returncode != 0:
        print(f"  FAILED (return code {result.returncode}, {elapsed:.0f}s)")
    else:
        print(f"  DONE ({elapsed:.0f}s)")
    return result.returncode


def collect_all(device_steady: int, device_karman: int, device_illusion: int):
    """Run all collection scripts."""
    errors = 0

    # === STEADY GROUP ===
    # Step 1: Empty channel (q_in for steady/illusion)
    dd = data_dir_for_scene("empty_channel")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_empty_channel.py",
                             ["--device", device_steady, "--steps", "200"],
                             "Empty channel (q_in for steady/illusion)")

    # Step 2: Pinball baseline (q_blk for steady/illusion)
    dd = data_dir_for_scene("pinball_baseline")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_pinball_baseline.py",
                             ["--device", device_steady, "--steps", "500"],
                             "Pinball baseline (q_blk for steady/illusion)")

    # === KARMAN GROUP ===
    # Step 3: Disturbance only (q_in for Karman)
    dd = data_dir_for_scene("disturbance_only")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_disturbance_only.py",
                             ["--device", device_karman, "--steps", "500"],
                             "Disturbance only (q_in for Karman)")

    # Step 4: Karman blk (dist+pinball zero, q_blk for Karman)
    dd = data_dir_for_scene("karman_blk")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_karman_blk.py",
                             ["--device", device_karman, "--steps", "500"],
                             "Karman blk (dist+pinball zero, q_blk for Karman)")

    # Step 5: Karman controlled (PPO, q_ctl)
    dd = data_dir_for_scene("karman_re100")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_fields_replay.py",
                             ["--scene", "karman_re100", "--device", device_karman],
                             "Karman PPO replay (q_ctl)")

    # === STEADY CLOAK ===
    dd = data_dir_for_scene("steady_cloak")
    if not os.path.isfile(os.path.join(dd, "fields.npz")):
        errors += run_script("collect_steady_cloak.py",
                             ["--device", device_steady, "--steps", "500", "--omega-rear", "5.1"],
                             "Steady cloak (q_ctl)")

    # === ILLUSION GROUP ===
    for diam_name in ["illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]:
        # Target cylinder
        tgt_cfg = SCENES.get(f"target_cylinder_{diam_name.replace('illusion_', '')}")
        tgt_dd = data_dir_for_scene(f"target_cylinder_{diam_name.replace('illusion_', '')}")
        
        # Check if target data exists in CCD
        ccd_target = os.path.join(_SRC, "CCD_analysis", "data", "target_cylinder",
                                  f"target_cylinder_{diam_name.replace('illusion_', '')}")
        if not os.path.isfile(os.path.join(tgt_dd, "fields.npz")) and \
           not os.path.isfile(os.path.join(tgt_dd, "fields_aligned.npz")):
            errors += run_script("collect_target_cylinder.py",
                                 ["--diameter", str(SCENES[diam_name]["target_diameter"]),
                                  "--device", device_illusion, "--steps", "500"],
                                 f"Target cylinder {diam_name}")

        # PPO fields replay
        dd = data_dir_for_scene(diam_name)
        if not os.path.isfile(os.path.join(dd, "fields.npz")):
            errors += run_script("collect_fields_replay.py",
                                 ["--scene", diam_name, "--device", device_illusion],
                                 f"Illusion {diam_name} PPO replay (q_ctl)")

    # === CHECK ===
    print(f"\n{'='*60}")
    print(f"All collections complete. Errors: {errors}")
    print(f"{'='*60}")
    check_data()

    return errors


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--check", action="store_true", help="Only check data, no collection")
    ap.add_argument("--device-steady", type=int, default=3)
    ap.add_argument("--device-karman", type=int, default=1)
    ap.add_argument("--device-illusion", type=int, default=3)
    args = ap.parse_args()

    if args.check:
        ok = check_data()
        print(f"\nAll data present: {ok}")
        return 0 if ok else 1

    return collect_all(args.device_steady, args.device_karman, args.device_illusion)


if __name__ == "__main__":
    sys.exit(main())
