#!/usr/bin/env python3
"""Launch the released four-GPU Q-Prefer D2 training recipe."""

from __future__ import annotations

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

from qprefer_reward.constants import BASE_MODEL_ID, BASE_MODEL_REVISION


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--media-root", type=Path)
    parser.add_argument(
        "--path-prefix-map",
        action="append",
        default=[],
        metavar="OLD=NEW",
        help="Relocate absolute paths from the released manifest",
    )
    parser.add_argument("--base-model", default=BASE_MODEL_ID)
    parser.add_argument("--base-revision", default=BASE_MODEL_REVISION)
    parser.add_argument("--initial-special-embeddings", type=Path)
    parser.add_argument("--gpus", default="0,1,2,3", help="CUDA indices, comma separated")
    parser.add_argument("--master-port", type=int, default=29541)
    parser.add_argument("--effective-batch-size", type=int, default=32)
    parser.add_argument("--per-device-batch-size", type=int, default=2)
    parser.add_argument("--max-steps", type=int, default=1203)
    parser.add_argument("--max-train-samples", type=int)
    parser.add_argument("--dataloader-num-workers", type=int, default=6)
    parser.add_argument("--no-deepspeed", action="store_true")
    parser.add_argument("--no-check-media", action="store_true")
    parser.add_argument("--resume-from-checkpoint")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    gpus = [value.strip() for value in args.gpus.split(",") if value.strip()]
    if not gpus:
        raise ValueError("--gpus must contain at least one CUDA index")
    denominator = len(gpus) * args.per_device_batch_size
    if args.effective_batch_size % denominator:
        raise ValueError(
            "effective batch size must be divisible by world size * per-device batch size"
        )
    accumulation = args.effective_batch_size // denominator

    repository = Path(__file__).resolve().parents[1]
    command = [
        sys.executable,
        "-m",
        "torch.distributed.run",
        f"--nproc_per_node={len(gpus)}",
        f"--master_port={args.master_port}",
        "-m",
        "qprefer_reward.training.train",
        "--train-manifest",
        str(args.manifest),
        "--output-dir",
        str(args.output),
        "--model-name-or-path",
        args.base_model,
        "--base-revision",
        args.base_revision,
        "--max-steps",
        str(args.max_steps),
        "--per-device-train-batch-size",
        str(args.per_device_batch_size),
        "--gradient-accumulation-steps",
        str(accumulation),
        "--dataloader-num-workers",
        str(args.dataloader_num_workers),
    ]
    if args.media_root:
        command.extend(("--media-root", str(args.media_root)))
    for mapping in args.path_prefix_map:
        command.extend(("--path-prefix-map", mapping))
    if args.initial_special_embeddings:
        command.extend(
            ("--initial-special-embeddings", str(args.initial_special_embeddings))
        )
    if args.max_train_samples:
        command.extend(("--max-train-samples", str(args.max_train_samples)))
    if args.no_check_media:
        command.append("--no-check-media")
    if args.resume_from_checkpoint:
        command.extend(("--resume-from-checkpoint", args.resume_from_checkpoint))
    if not args.no_deepspeed:
        command.extend(
            (
                "--deepspeed",
                str(repository / "training/configs/deepspeed_zero2_no_offload.json"),
            )
        )

    environment = dict(os.environ)
    environment["CUDA_VISIBLE_DEVICES"] = ",".join(gpus)
    environment.setdefault("FORCE_QWENVL_VIDEO_READER", "decord")
    environment.setdefault("VIDEO_IO_MAX_RETRIES", "0")
    environment.setdefault("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", "1800")
    print("Launching:", " ".join(command), flush=True)
    subprocess.run(command, cwd=repository, env=environment, check=True)


if __name__ == "__main__":
    main()