"""Prove that an MLX distributed group forms, and time one collective on it.

Purpose: the first thing to run on a new cluster, before any model is involved. It
    joins the distributed group, reports the rank and size every process sees, runs
    an all-sum whose result is known in advance so a wrong answer is obvious, and
    then times repeated all-sums at several array sizes so that the ring backend and
    the JACCL backend can be compared on the same cable. Rank 0 appends one JSON line
    to the lab notebook.
Platform: mac (Track M) for the two-Mac path; the same script runs anywhere MLX runs,
    including two ranks on one machine, which is the single-machine path.
Minimum memory: 32 GB per Mac for the primary path; 24 GB for the single-Mac path.
    This script itself allocates a few hundred megabytes at the largest size.
Assumes: mlx installed in the python that mlx.launch starts, and this file present at
    the SAME absolute path on every machine. Launched by mlx.launch, never directly:
    run with plain python it reports a group of size one, which is correct and dull.

Usage:
    mlx.launch -n 2 -- python check-group.py --label single-machine
    mlx.launch --backend ring --hostfile hosts.json -- \
        /path/to/python check-group.py --label two-macs-ring
    mlx.launch --backend jaccl --hostfile hosts.json -- \
        /path/to/python check-group.py --label two-macs-jaccl
    python3 check-group.py --print            (group of one; writes nothing)
"""

from __future__ import annotations

import argparse
import json
import platform
import time
from datetime import date, datetime, timezone
from pathlib import Path

import mlx.core as mx

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

# Array sizes to time, in float32 elements. The smallest is latency-dominated and
# the largest is bandwidth-dominated, which is exactly the difference between the
# ring backend and an RDMA backend that the lesson describes.
DEFAULT_SIZES = [1024, 262144, 4194304, 33554432]


def parse_args():
    p = argparse.ArgumentParser(description="MLX distributed group check")
    p.add_argument(
        "--backend",
        default="any",
        choices=["any", "ring", "jaccl", "mpi", "nccl"],
        help="Backend to request from mx.distributed.init(). Default: any.",
    )
    p.add_argument(
        "--sizes",
        default=",".join(str(s) for s in DEFAULT_SIZES),
        help="Comma-separated float32 element counts to time.",
    )
    p.add_argument("--repeat", type=int, default=20, help="All-sums per size.")
    p.add_argument("--warmup", type=int, default=5, help="Untimed all-sums per size.")
    p.add_argument(
        "--label",
        default="",
        help="A name for this run, for example two-macs-jaccl. Goes in the record.",
    )
    p.add_argument("--labbook", default=str(HERE / "labbook.md"))
    p.add_argument(
        "--print",
        dest="print_only",
        action="store_true",
        help="Show the record and write nothing.",
    )
    return p.parse_args()


def time_all_sum(elements: int, repeat: int, warmup: int) -> dict:
    """Average seconds per all-sum at this size, and the bytes each one moves."""
    x = mx.ones(elements, dtype=mx.float32)
    for _ in range(warmup):
        mx.eval(mx.distributed.all_sum(x))
    start = time.perf_counter()
    for _ in range(repeat):
        mx.eval(mx.distributed.all_sum(x))
    elapsed = time.perf_counter() - start
    return {
        "elements": elements,
        "bytes_per_array": elements * 4,
        "repeats": repeat,
        "seconds_per_all_sum": elapsed / repeat,
    }


def main() -> int:
    args = parse_args()

    world = mx.distributed.init(backend=args.backend)
    rank = world.rank()
    size = world.size()

    # A result every process can check without trusting the network: an array of
    # ones summed across `size` processes must be exactly `size` everywhere.
    probe = mx.distributed.all_sum(mx.ones(8, dtype=mx.float32))
    mx.eval(probe)
    expected = float(size)
    observed = [float(v) for v in probe.tolist()]
    correct = all(abs(v - expected) < 1e-6 for v in observed)

    print(f"rank {rank} of {size}: all_sum(ones) = {observed[0]}, expected {expected}")
    if not correct:
        print(f"rank {rank}: WRONG RESULT. The group formed but the maths did not.")
        return 1

    if size == 1:
        print(
            "Group size is 1, so every collective was a noop. That is the expected "
            "result for plain python; launch with mlx.launch to get a real group."
        )

    sizes = [int(s) for s in args.sizes.split(",") if s.strip()]
    timings = [time_all_sum(n, args.repeat, args.warmup) for n in sizes]

    record = {
        "lab": "part-21/two-mac-cluster-over-thunderbolt-5",
        "record": "group-check",
        "date": date.today().isoformat(),
        "recorded_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "label": args.label,
        "backend_requested": args.backend,
        "group_size": size,
        "all_sum_correct": correct,
        "machine": platform.node().split(".")[0],
        "platform": platform.platform(),
        "timings": timings,
    }

    if rank != 0:
        return 0

    text = json.dumps(record, sort_keys=True)
    if args.print_only:
        print(text)
        return 0
    with open(args.labbook, "a", encoding="utf-8") as fh:
        fh.write(text + "\n")
    print(text)
    print(f"Appended to {args.labbook}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())