from abc import ABC, abstractmethod

from torch import Tensor


class WriterInterface(ABC):
    @abstractmethod
    def open(self):
        pass

    @abstractmethod
    def close(self):
        pass

    @abstractmethod
    def _write(self, data: Tensor):
        pass

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def write(self, data: Tensor, **kwargs):
        """data is a 2D Tensor of shape (N, 5) with rows (frame, x, y, z, n)"""
        if not isinstance(data, Tensor) or data.ndim != 2 or data.shape[1] != 5:
            raise ValueError(
                "data must be a 2D of shape (N, 5) with rows (frame, x, y, z, n)"
            )
        self._write(data, **kwargs)
