"""Check imports and optional tiny CUDA operation, never download models."""

import argparse
import importlib.metadata
import json
import platform


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cuda", action="store_true", help="Allocate a tiny tensor on logical cuda:0")
    args = parser.parse_args()
    import torch
    from trl import SFTConfig, SFTTrainer
    from peft import LoraConfig
    from transformers import AutoModelForSequenceClassification

    assert SFTTrainer is not None and AutoModelForSequenceClassification is not None
    assert "completion_only_loss" in SFTConfig.__dataclass_fields__
    assert LoraConfig(r=16, lora_alpha=32, task_type="CAUSAL_LM").r == 16
    result = {
        "python": platform.python_version(),
        "packages": {name: importlib.metadata.version(name) for name in (
            "torch", "transformers", "trl", "peft", "accelerate", "datasets"
        )},
        "cuda_available": torch.cuda.is_available(),
        "cuda_runtime": torch.version.cuda,
        "cuda_operation": "not_requested",
        "model_training": "not_run",
    }
    if args.cuda:
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA requested but unavailable")
        x = torch.ones((2, 2), device="cuda:0", requires_grad=True)
        y = (x @ x).sum()
        y.backward()
        torch.cuda.synchronize()
        assert y.item() == 8 and torch.all(x.grad == 4).item()
        result["cuda_operation"] = "tiny_matmul_backward_passed"
        result["device"] = torch.cuda.get_device_name(0)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()