from typing import Optional, Tuple


try:
  import cupy as cp
  HAS_CUPY = True
except ImportError:
  HAS_CUPY = False


class CupyKernelCache:
  def __init__(self):
    self._cache: dict = {}

  def get_or_compile(self, name: str, source: str, options: tuple = ()) -> "cp.RawKernel":
    key = (name, source, options)
    if key not in self._cache:
      self._cache[key] = cp.RawKernel(source, name, options=options)
    return self._cache[key]

  def clear(self):
    self._cache.clear()


_default_cache = None


def _get_cache() -> CupyKernelCache:
  global _default_cache
  if _default_cache is None:
    _default_cache = CupyKernelCache()
  return _default_cache


class CupyKernel:
  def __init__(
    self,
    name: str,
    source: str,
    grid: Tuple,
    block: Tuple,
    options: tuple = (),
  ):
    if not HAS_CUPY:
      raise RuntimeError("cupy is not installed")
    self.name = name
    self.source = source
    self.grid = grid
    self.block = block
    self.options = options
    self._kernel: Optional["cp.RawKernel"] = None

  def _ensure_compiled(self):
    if self._kernel is None:
      self._kernel = _get_cache().get_or_compile(
        self.name, self.source, self.options
      )

  def __call__(self, *args):
    self._ensure_compiled()
    self._kernel(self.grid, self.block, args)

  @staticmethod
  def from_builder(
    builder,
    name: str,
    grid: Tuple,
    block: Tuple,
    options: tuple = (),
  ) -> "CupyKernel":
    source = builder.build_source()
    return CupyKernel(name, source, grid, block, options)


def to_cupy_array(data, dtype=None):
  if not HAS_CUPY:
    raise RuntimeError("cupy is not installed")
  import numpy as np
  arr = np.array(data, dtype=dtype)
  return cp.asarray(arr)


def cupy_available() -> bool:
  return HAS_CUPY
