from typing import override

import logging
import numpy as np

from kedro.io import AbstractDataset
from pathlib import Path
from numpy.typing import NDArray


logger = logging.getLogger(__name__)


class NumpyDataset(AbstractDataset[NDArray, NDArray]):
    def __init__(self, filepath: str):
        super().__init__()
        # parse the path and protocol (e.g. file, http, s3, etc.)
        self._filepath = Path(filepath)

    @override
    def _describe(self):
        return dict(filepath=self._filepath)

    @override
    def _exists(self):
        return self._filepath.exists()

    @override
    def load(self):
        logger.info(f"loading numpy array from {self._filepath}")
        return np.load(self._filepath)

    @override
    def save(self, data: NDArray):
        # TODO: Support multiple array and npz format with np.savez
        logger.info(f"writing numpy array to {self._filepath}")
        self._filepath.parent.mkdir(exist_ok=True, parents=True)
        np.save(self._filepath, data)