from pycuda import gpuarray  # type: ignore
from pycuda.driver import Function  # type: ignore

from pydtnn.metrics.metric import Metric
from pydtnn.backends.pycuda.utils.tensor_array import TensorArray
from pydtnn.utils.constants import ArrayShape


class MetricPycuda(Metric[TensorArray]):
    """
    Extends a Metric class with the attributes and methods required by GPU Metrics.
    """

    def __init__(self, eps=1e-8):
        super().__init__(eps)
        # NOTE: The following attributes will be initializated later.
        self.grid = None
        self.block = None

    def _model_init(self) -> None:
        super()._model_init()
        self.kernel = self._kernel_init()

        self.grid = self.model.cuda_grid
        self.block = self.model.cuda_block

    def _kernel_init(self) -> Function:
        raise NotImplementedError()
