#!/usr/bin/env python3
"""Profile the NCCL all-reduce paths consumed by the Qwen235B experiment."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import statistics
import sys
from datetime import datetime, timezone
from pathlib import Path

import torch
import torch.distributed as dist


DEFAULT_TOKENS = [
    1,
    2,
    3,
    4,
    6,
    8,
    12,
    16,
    24,
    32,
    48,
    63,
    64,
    65,
    96,
    127,
    128,
    129,
    192,
    255,
    256,
    257,
    384,
    511,
    512,
    513,
    768,
    1023,
    1024,
    1025,
    1536,
    2047,
    2048,
    2049,
    3072,
    4095,
    4096,
    4097,
    6144,
    8191,
    8192,
    8193,
    12288,
    16383,
    16384,
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--hidden-size", type=int, default=4096)
    parser.add_argument("--world-sizes", type=int, nargs="+", default=[4, 8])
    parser.add_argument("--tokens", type=int, nargs="+", default=DEFAULT_TOKENS)
    parser.add_argument("--warmup-iterations", type=int, default=5)
    parser.add_argument("--measured-iterations", type=int, default=30)
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def profile_one_size(
    *,
    group: dist.ProcessGroup,
    group_size: int,
    num_tokens: int,
    hidden_size: int,
    warmup_iterations: int,
    measured_iterations: int,
) -> list[float]:
    tensor = torch.ones(
        (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
    )
    for _ in range(warmup_iterations):
        dist.all_reduce(tensor, group=group)
    dist.barrier(group=group)
    torch.cuda.synchronize()

    local_samples: list[float] = []
    for _ in range(measured_iterations):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        dist.all_reduce(tensor, group=group)
        end.record()
        end.synchronize()
        local_samples.append(float(start.elapsed_time(end)))

    samples_by_rank: list[list[float] | None] = [None] * group_size
    dist.all_gather_object(samples_by_rank, local_samples, group=group)
    del tensor

    # Collective completion is limited by the slowest rank. Preserve that
    # critical-path statistic instead of reporting only rank 0.
    return [
        max(float(rank_samples[index]) for rank_samples in samples_by_rank if rank_samples)
        for index in range(measured_iterations)
    ]


def profile_group(
    *,
    group: dist.ProcessGroup,
    group_ranks: list[int],
    args: argparse.Namespace,
) -> list[dict[str, object]]:
    global_rank = dist.get_rank()
    if global_rank not in group_ranks:
        return []

    group_size = len(group_ranks)
    group_rank = dist.get_rank(group=group)
    rows: list[dict[str, object]] = []
    trial_orders = (("forward", args.tokens), ("reverse", list(reversed(args.tokens))))
    for trial_id, (order_name, tokens) in enumerate(trial_orders):
        for num_tokens in tokens:
            samples = profile_one_size(
                group=group,
                group_size=group_size,
                num_tokens=num_tokens,
                hidden_size=args.hidden_size,
                warmup_iterations=args.warmup_iterations,
                measured_iterations=args.measured_iterations,
            )
            if group_rank != 0:
                continue
            ordered = sorted(samples)
            p95_index = int(0.95 * (len(ordered) - 1))
            rows.append(
                {
                    "time_stats.all_reduce.min": min(samples),
                    "time_stats.all_reduce.max": max(samples),
                    "time_stats.all_reduce.mean": statistics.fmean(samples),
                    "time_stats.all_reduce.median": statistics.median(samples),
                    "time_stats.all_reduce.std": statistics.stdev(samples),
                    "time_stats.all_reduce.p95": ordered[p95_index],
                    "rank": 0,
                    "num_workers": group_size,
                    "size": num_tokens * args.hidden_size * 2,
                    "collective": "all_reduce",
                    "devices_per_node": group_size,
                    "max_devices_per_node": dist.get_world_size(),
                    "profiling_precision": "BF16",
                    "measurement_type": "CUDA_EVENT",
                    "backend": "nccl",
                    "num_tokens": num_tokens,
                    "hidden_size": args.hidden_size,
                    "trial_id": trial_id,
                    "order": order_name,
                    "warmup_iterations": args.warmup_iterations,
                    "measured_iterations": args.measured_iterations,
                }
            )
    return rows


def main() -> None:
    args = parse_args()
    if args.hidden_size <= 0:
        raise ValueError("hidden-size must be positive")
    if args.warmup_iterations <= 0 or args.measured_iterations < 2:
        raise ValueError("warmup must be positive and measured iterations must be >= 2")
    if sorted(set(args.tokens)) != sorted(args.tokens) or min(args.tokens) <= 0:
        raise ValueError("tokens must be unique positive integers")

    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    global_rank = dist.get_rank()
    global_world_size = dist.get_world_size()
    if max(args.world_sizes) > global_world_size:
        raise ValueError(
            f"requested world size {max(args.world_sizes)} exceeds torchrun world size "
            f"{global_world_size}"
        )

    groups: dict[int, tuple[dist.ProcessGroup, list[int]]] = {}
    for world_size in sorted(set(args.world_sizes)):
        ranks = list(range(world_size))
        group = dist.group.WORLD if world_size == global_world_size else dist.new_group(ranks)
        groups[world_size] = (group, ranks)

    all_rows: list[dict[str, object]] = []
    for world_size in args.world_sizes:
        group, ranks = groups[world_size]
        all_rows.extend(profile_group(group=group, group_ranks=ranks, args=args))
        dist.barrier()

    if global_rank == 0:
        expected_rows = len(set(args.world_sizes)) * len(args.tokens) * 2
        if len(all_rows) != expected_rows:
            raise RuntimeError(
                f"profile row count mismatch: expected={expected_rows}, actual={len(all_rows)}"
            )
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", newline="", encoding="utf-8") as output_file:
            writer = csv.DictWriter(output_file, fieldnames=list(all_rows[0]))
            writer.writeheader()
            writer.writerows(all_rows)

        manifest = {
            "generated_at_utc": datetime.now(timezone.utc).isoformat(),
            "command": [sys.executable, *sys.argv],
            "output": str(args.output.resolve()),
            "output_sha256": sha256(args.output),
            "rows": len(all_rows),
            "world_sizes": args.world_sizes,
            "tokens": args.tokens,
            "hidden_size": args.hidden_size,
            "dtype": "torch.bfloat16",
            "backend": dist.get_backend(),
            "torch_version": torch.__version__,
            "cuda_version": torch.version.cuda,
            "nccl_version": list(torch.cuda.nccl.version()),
            "gpu": torch.cuda.get_device_name(0),
            "warmup_iterations": args.warmup_iterations,
            "measured_iterations": args.measured_iterations,
            "sample_semantics": "per-iteration maximum CUDA-event time across participating ranks",
        }
        manifest_path = args.output.with_suffix(args.output.suffix + ".manifest.json")
        manifest_path.write_text(
            json.dumps(manifest, indent=2) + "\n", encoding="utf-8"
        )
        print(json.dumps(manifest, indent=2))

    dist.destroy_process_group()


if __name__ == "__main__":
    main()