#!/usr/bin/env python3
"""
MagicDrive-V2 推理脚本（纯推理，不包含任何下载/安装逻辑）
前提：已运行 scripts/setup_space.py 完成环境搭建

用法:
    # 224x400 快速预览（17帧，~22GB 显存）
    python scripts/infer.py --scene s8 --fast

    # 424x800 正式生成（129帧=10.75s，~44GB 显存）
    python scripts/infer.py --scene s8

    # 指定输出路径
    python scripts/infer.py --scene s25 --output outputs/my_video

场景:
    s8   停车场 21车 (车在正前方)
    s25  停车场+天桥 25车 (车最近)
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

# 场景配置
SCENES = {
    "s8": {
        "pkl": "data/nuscenes_mmdet3d-12Hz/ghost_peek_s8.pkl",
        "desc": "停车场 21车 (车在正前方 -20m)",
    },
    "s25": {
        "pkl": "data/nuscenes_mmdet3d-12Hz/ghost_peek_s25.pkl",
        "desc": "停车场+天桥 25车 (车最近 6.9m)",
    },
}

# 分辨率配置
PRESETS = {
    "fast": {
        "config": "fullx224x400_stdit3_CogVAE_boxTDS_wCT_xCE_wSST",
        "num_frames": 17,
        "desc": "低分辨率预览 (224x400, 17帧, ~22GB)",
    },
    "full": {
        "config": "fullx424x800_stdit3_CogVAE_boxTDS_wCT_xCE_wSST",
        "num_frames": 129,
        "desc": "正式生成 (424x800, 129帧=10.75s, ~44GB)",
    },
}


def check_ready() -> bool:
    """检查环境是否就绪"""
    required = [
        ROOT / "pretrained" / "CogVideoX-2b" / "vae",
        ROOT / "pretrained" / "t5-v1_1-xxl",
        ROOT / "ckpts" / "MagicDriveDiT-stage3-40k-ft",
    ]
    all_ok = True
    for p in required:
        if not p.exists():
            print(f"[MISS] {p.relative_to(ROOT)}")
            all_ok = False
    if not all_ok:
        print("\n请先运行: python scripts/setup_space.py")
    return all_ok


def main():
    parser = argparse.ArgumentParser(description="MagicDrive-V2 推理")
    parser.add_argument("--scene", choices=list(SCENES.keys()), default="s8",
                        help="选择场景")
    parser.add_argument("--fast", action="store_true",
                        help="低分辨率快速预览 (224x400, 17帧)")
    parser.add_argument("--num_frames", type=int, default=None,
                        help="自定义帧数 (覆盖预设)")
    parser.add_argument("--output", type=str, default=None,
                        help="输出目录")
    parser.add_argument("--cpu_offload", action="store_true", default=True,
                        help="CPU offload (省显存)")
    parser.add_argument("--seed", type=int, default=42,
                        help="随机种子")
    args = parser.parse_args()

    os.chdir(str(ROOT))
    sys.path.insert(0, str(ROOT))

    if not check_ready():
        sys.exit(1)

    scene = SCENES[args.scene]
    preset = PRESETS["fast"] if args.fast else PRESETS["full"]
    num_frames = args.num_frames or preset["num_frames"]

    print("=" * 60)
    print(f"场景: {args.scene} - {scene['desc']}")
    print(f"配置: {preset['desc']}")
    print(f"帧数: {num_frames}")
    print(f"输出: {args.output or 'auto'}")
    print("=" * 60)

    # 构建推理命令
    # 注意: parse_args 需要的是配置文件路径（位置参数）+ --cfg-options key=val ...
    config_path = str(ROOT / "configs" / "magicdrive" / "inference" / f"{preset['config']}.py")
    cmd = [
        "torchrun", "--standalone", "--nproc_per_node=1",
        str(ROOT / "scripts" / "inference_magicdrive.py"),
        config_path,
        "--cfg-options",
        f"num_frames={num_frames}",
        f"seed={args.seed}",
    ]

    if args.cpu_offload:
        cmd.append("cpu_offload=true")

    # 用 dataset_cfg_overrides 切换 PKL
    override = f"dataset_cfg_overrides=[('dataset.data.val.ann_file', '{scene['pkl']}')]"
    cmd.append(override)

    if args.output:
        cmd.append(f"outputs={args.output}")

    print(f"\n运行: {' '.join(cmd)}")
    print("-" * 60)

    result = subprocess.run(
        cmd,
        cwd=str(ROOT),
        env={**os.environ, "PYTHONPATH": str(ROOT)},
    )

    if result.returncode == 0:
        print("\n推理完成！输出在 outputs/ 目录下")
    else:
        print(f"\n推理失败 (exit code: {result.returncode})")
        sys.exit(result.returncode)


if __name__ == "__main__":
    main()