"""Benchmark ragged H3 Ulysses transport and attention under torchrun."""

import argparse
import json
import os
import statistics
import time
from pathlib import Path

import torch
import torch.distributed as dist

from h3_blackwell_runtime.attention import run_attention, run_flash4_attention_bshd, run_sol_attention_bshd
from h3_blackwell_runtime.distributed import SequenceParallelContext


def synchronize(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def run_backend(backend: str, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    if backend == "flash4":
        return run_flash4_attention_bshd(q, k, v, is_causal=False)
    if backend == "sol_attn":
        return run_sol_attention_bshd(q, k, v, is_causal=False)
    return run_attention(
        q.transpose(1, 2).contiguous(),
        k.transpose(1, 2).contiguous(),
        v.transpose(1, 2).contiguous(),
        backend=backend,
        is_causal=False,
    ).transpose(1, 2)


parser = argparse.ArgumentParser()
parser.add_argument("--sequence", type=int, default=20000)
parser.add_argument("--heads", type=int, default=56)
parser.add_argument("--head-dim", type=int, default=128)
parser.add_argument("--backend", default="sdpa")
parser.add_argument("--warmup", type=int, default=3)
parser.add_argument("--iterations", type=int, default=10)
parser.add_argument("--output", type=Path)
args = parser.parse_args()

local_rank = int(os.environ.get("LOCAL_RANK", "0"))
use_cuda = torch.cuda.is_available()
device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
if use_cuda:
    torch.cuda.set_device(device)
dist.init_process_group(backend="nccl" if use_cuda else "gloo", device_id=device if use_cuda else None)
rank = dist.get_rank()
world_size = dist.get_world_size()
context = SequenceParallelContext.create(args.sequence, args.heads, args.head_dim)

generator = torch.Generator(device=device).manual_seed(440420 + rank)
shape = (1, context.local_token_length, args.heads, args.head_dim)
dtype = torch.bfloat16 if use_cuda else torch.float32
q = torch.randn(shape, generator=generator, dtype=dtype, device=device)
k = torch.randn(shape, generator=generator, dtype=dtype, device=device)
v = torch.randn(shape, generator=generator, dtype=dtype, device=device)


def iteration() -> tuple[float, float, float, torch.Tensor]:
    dist.barrier()
    synchronize(device)
    started = time.perf_counter()
    full_q, full_k, full_v = context.seq_to_heads(q, k, v)
    synchronize(device)
    after_forward = time.perf_counter()
    head_output = run_backend(args.backend, full_q, full_k, full_v)
    synchronize(device)
    after_attention = time.perf_counter()
    local_output = context.heads_to_seq(head_output)
    synchronize(device)
    finished = time.perf_counter()
    return (
        after_forward - started,
        after_attention - after_forward,
        finished - after_attention,
        local_output,
    )


for _ in range(args.warmup):
    *_timings, output = iteration()
    del output

rank_timings = []
for _ in range(args.iterations):
    forward, attention, inverse, output = iteration()
    rank_timings.append((forward, attention, inverse, forward + attention + inverse))
    del output

timings = torch.tensor(rank_timings, dtype=torch.float64, device=device)
gathered = [torch.empty_like(timings) for _ in range(world_size)]
dist.all_gather(gathered, timings)

if use_cuda:
    peak_memory = torch.tensor([torch.cuda.max_memory_allocated(device)], dtype=torch.int64, device=device)
else:
    peak_memory = torch.tensor([0], dtype=torch.int64, device=device)
memory_by_rank = [torch.empty_like(peak_memory) for _ in range(world_size)]
dist.all_gather(memory_by_rank, peak_memory)

if rank == 0:
    stacked = torch.stack(gathered).cpu()
    stage_names = ("forward_all_to_all", "attention", "inverse_all_to_all", "total")
    stages = {}
    for index, name in enumerate(stage_names):
        maximum_rank = stacked[:, :, index].amax(dim=0).tolist()
        stages[name] = {
            "median_seconds": statistics.median(maximum_rank),
            "minimum_seconds": min(maximum_rank),
            "maximum_seconds": max(maximum_rank),
        }
    element_size = q.element_size()
    report = {
        "world_size": world_size,
        "backend": args.backend,
        "device": torch.cuda.get_device_name(device) if use_cuda else "cpu",
        "torch": torch.__version__,
        "sequence": args.sequence,
        "heads": args.heads,
        "head_dim": args.head_dim,
        "token_lengths": list(context.token_lengths),
        "head_lengths": list(context.head_lengths),
        "dtype": str(dtype),
        "iterations": args.iterations,
        "aggregate_transport_bytes_per_iteration": {
            "forward_qkv": 3 * args.sequence * args.heads * args.head_dim * element_size,
            "inverse_output": args.sequence * args.heads * args.head_dim * element_size,
        },
        "stages": stages,
        "peak_allocated_bytes_by_rank": [int(value.item()) for value in memory_by_rank],
    }
    serialized = json.dumps(report, indent=2)
    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(serialized + "\n", encoding="utf-8")
    print(serialized)

dist.destroy_process_group()