import random
from abc import ABC, abstractmethod
from functools import lru_cache
from typing import NamedTuple, Optional

import torch

IS_ROCM = torch.version.hip is not None


class DeviceCapability(NamedTuple):
    major: int
    minor: int

    def as_version_str(self) -> str:
        return f"{self.major}.{self.minor}"

    def to_int(self) -> int:
        """
        Express device capability as an integer ``<major><minor>``.

        It is assumed that the minor version is always a single digit.
        """
        assert 0 <= self.minor < 10
        return self.major * 10 + self.minor


class Platform(ABC):
    simple_compile_backend: str = "inductor"

    @classmethod
    @abstractmethod
    def get_device_name(cls, device_id: int = 0) -> str: ...

    @abstractmethod
    def is_cuda_alike(self) -> bool:
        """Stateless version of :func:`torch.cuda.is_available`."""
        ...

    @abstractmethod
    def is_rocm(self): ...

    @classmethod
    def seed_everything(cls, seed: Optional[int] = None) -> None:
        """
        Set the seed of each random module.
        `torch.manual_seed` will set seed on all devices.

        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
        """
        if seed is not None:
            random.seed(seed)
            torch.manual_seed(seed)
            try:
                import numpy as np

                np.random.seed(seed)
            except ImportError:
                pass


class CudaPlatform(Platform):
    @classmethod
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)

    @classmethod
    @lru_cache(maxsize=8)
    def get_device_name(cls, device_id: int = 0) -> str:
        return torch.cuda.get_device_name(0)

    def is_cuda_alike(self) -> bool:
        return True

    def is_rocm(self):
        return False


class RocmPlatform(Platform):
    @classmethod
    @lru_cache(maxsize=8)
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)

    @classmethod
    @lru_cache(maxsize=8)
    def get_device_name(cls, device_id: int = 0) -> str:
        return torch.cuda.get_device_name(device_id)

    def is_cuda_alike(self) -> bool:
        return True

    def is_rocm(self):
        return True


current_platform = RocmPlatform() if IS_ROCM else CudaPlatform()