"""MPS smoke test for Exp 4 Phase 3 — run on the owner's M5 Mac, no assistant needed.

Everything the analysis needs is printed between the REPORT BEGIN/END banners —
copy that whole block back. Partial failures still produce a report (each stage
catches and records its own error).

Setup on a fresh Mac (Terminal):

    python3 -m venv ~/mps-test && source ~/mps-test/bin/activate
    pip install torch transformers pillow numpy
    python scripts/mps_smoke_test.py            # (or run from any folder)

If the report shows MPS ops missing, run once more with fallback to measure the
realistic mixed speed:

    PYTORCH_ENABLE_MPS_FALLBACK=1 python scripts/mps_smoke_test.py
"""

from __future__ import annotations

import os
import platform
import time
import traceback

import numpy as np
import torch
from PIL import Image

MODEL_ID = "tue-mps/eomt-dinov3-ade-semantic-large-512"
REVISION = "1d1f172700ab69371df7afe778fe133fff7dd521"  # keep in sync with segmentation.MODEL_REVISIONS


def synthetic_tile(size: int) -> Image.Image:
    rng = np.random.default_rng(0)
    return Image.fromarray(rng.integers(0, 255, (size, size, 3), dtype=np.uint8))


def _sync(device: str) -> None:
    if device == "mps":
        torch.mps.synchronize()


def time_forward(model, inputs, device, runs: int = 3) -> float:
    model = model.to(device).eval()
    moved = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
    with torch.inference_mode():
        model(**moved)  # warm-up
        _sync(device)
        times = []
        for _ in range(runs):
            t0 = time.perf_counter()
            model(**moved)
            _sync(device)
            times.append(time.perf_counter() - t0)
    return min(times)


def time_train_step(model, inputs, device) -> float:
    model = model.to(device).train()
    moved = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
    optim = torch.optim.AdamW(model.parameters(), lr=1e-5)
    t0 = time.perf_counter()
    out = model(**moved)
    loss = out.masks_queries_logits.float().mean() + out.class_queries_logits.float().mean()
    loss.backward()
    optim.step()
    optim.zero_grad(set_to_none=True)
    _sync(device)
    return time.perf_counter() - t0


def stage(report: dict, name: str, fn):
    try:
        report[name] = fn()
        print(f"{name}: {report[name]}")
    except Exception:
        report[name] = "FAILED"
        print(f"{name}: FAILED")
        print(traceback.format_exc())


def main() -> None:
    print("=" * 66)
    print("MPS SMOKE REPORT BEGIN — copy everything down to REPORT END")
    print("=" * 66)
    print(f"machine: {platform.machine()} | {platform.platform()}")
    print(f"python: {platform.python_version()} | torch: {torch.__version__}")
    print(f"PYTORCH_ENABLE_MPS_FALLBACK={os.getenv('PYTORCH_ENABLE_MPS_FALLBACK', '<unset>')}")
    print(f"mps available: {torch.backends.mps.is_available()} (built: {torch.backends.mps.is_built()})")

    report: dict = {}
    try:
        from transformers import AutoImageProcessor, AutoModelForUniversalSegmentation

        processor = AutoImageProcessor.from_pretrained(MODEL_ID, revision=REVISION)
        model = AutoModelForUniversalSegmentation.from_pretrained(MODEL_ID, revision=REVISION)
        print(f"model: {MODEL_ID} ({sum(p.numel() for p in model.parameters()) / 1e6:.0f}M params)")
    except Exception:
        print("MODEL LOAD FAILED:")
        print(traceback.format_exc())
        print("REPORT END")
        return

    inputs_512 = processor(images=synthetic_tile(512), return_tensors="pt")
    inputs_1280 = processor(images=synthetic_tile(1280), return_tensors="pt")

    stage(report, "cpu_forward_512_s", lambda: round(time_forward(model, inputs_512, "cpu"), 2))
    stage(report, "cpu_forward_1280_s", lambda: round(time_forward(model, inputs_1280, "cpu"), 2))

    if not torch.backends.mps.is_available():
        print("NO MPS DEVICE — training would need the cloud-GPU fallback.")
        print("REPORT END")
        return

    stage(report, "mps_forward_512_s", lambda: round(time_forward(model, inputs_512, "mps"), 2))
    stage(report, "mps_forward_1280_s", lambda: round(time_forward(model, inputs_1280, "mps"), 2))
    stage(report, "mps_train_step_512_s", lambda: round(time_train_step(model, inputs_512, "mps"), 2))
    # second step is the honest one (first includes optimizer-state allocation)
    stage(report, "mps_train_step_512_s_2nd", lambda: round(time_train_step(model, inputs_512, "mps"), 2))
    try:
        alloc = torch.mps.current_allocated_memory() / 1e9
        print(f"mps_allocated_gb: {alloc:.1f}")
    except Exception:
        pass

    step = report.get("mps_train_step_512_s_2nd")
    if isinstance(step, float):
        print(f"epoch_math: 2000 crops / batch 1 -> {2000 * step / 60:.0f} min/epoch at this step time")
        print("VERDICT: MPS training viable" if step < 30 else
              "VERDICT: MPS works but slow — consider cloud GPU for the real fine-tune")
    else:
        print("VERDICT: train step failed on MPS — retry with PYTORCH_ENABLE_MPS_FALLBACK=1; "
              "if it still fails, forward-only on MPS is usable and training goes to cloud GPU")
    print("=" * 66)
    print("REPORT END")
    print("=" * 66)


if __name__ == "__main__":
    main()