"""Small environment smoke check; does not download data or train a PII model."""
import importlib.metadata
import json
import platform
from pathlib import Path

import onnxruntime as ort
import torch

result = {
    "platform": platform.platform(),
    "python": platform.python_version(),
    "versions": {name: importlib.metadata.version(name) for name in (
        "torch", "transformers", "datasets", "accelerate", "safetensors",
        "scikit-learn", "onnx", "onnxscript", "onnxruntime",
    )},
    "mps_built": torch.backends.mps.is_built(),
    "mps_available": torch.backends.mps.is_available(),
    "onnxruntime_providers": ort.get_available_providers(),
}
if result["mps_available"]:
    torch.manual_seed(0)
    model = torch.nn.Linear(8, 2).to("mps")
    optimizer = torch.optim.AdamW(model.parameters())
    before = model.weight.detach().clone()
    loss = model(torch.ones(4, 8, device="mps")).square().mean()
    loss.backward()
    finite_gradients = all(torch.isfinite(p.grad).all().item() for p in model.parameters())
    optimizer.step()
    torch.mps.synchronize()
    result["mps_smoke_check"] = {
        "finite_loss": torch.isfinite(loss).item(),
        "finite_gradients": finite_gradients,
        "weights_changed": not torch.equal(before, model.weight),
    }
    assert all(result["mps_smoke_check"].values()), result
out = Path(__file__).resolve().parents[1] / "artifacts" / "environment-check.json"
out.parent.mkdir(exist_ok=True)
out.write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))