IO¶

@dataclass
class EphysPath(VTKPath):
    """Path class for ephys session-level data.

    Follows the naming convention::

        {root}/sessions/{session_id}/{dtype}_{session_id}[_probe{N}].{ext}

    The ``fpath`` property points to a file under ``sessions/``.
    Auxiliary properties ``raw_dir`` and ``nwb_path`` expose the ``raw/``
    and ``nwb/`` subdirectories for convenience.

    VTKPath fields ``subject``, ``task``, ``run``, ``desc``, and ``suffix``
    are inherited but not used by ``fpath``; use ``session_id`` and ``dtype``
    instead.

    Parameters
    ----------
    root : Path
        Root project directory (e.g. ``DB/ephys/MonkeyVision``).
    session_id : str, optional
        Full session identifier, e.g. ``"251024_FanFan_nsd1w_MSB"``.
        Format: ``{date}_{subject}_{paradigm}_{region}``.
    dtype : str, optional
        Data type.  Must be one of :data:`EPHYS_DTYPES`.
    probe : int, optional
        Probe index (0-based).  ``None`` for single-probe sessions.
    extension : str, optional
        File extension.  Must be one of :data:`EPHYS_EXTENSIONS`.
        Default is ``"h5"``.
    """

    session_id: str | None = None
    dtype: str | None = None
    modality: str = field(default="ephys", init=False)

    def __post_init__(self) -> None:
        """Validate dtype and extension."""
        super().__post_init__()
        if self.dtype is not None and self.dtype not in EPHYS_DTYPES:
            raise ValueError(f"Invalid dtype '{self.dtype}'. Must be one of {sorted(EPHYS_DTYPES)}")
        if self.extension is not None:
            ext = self.extension.lstrip(".")
            if ext not in EPHYS_EXTENSIONS:
                raise ValueError(f"Invalid extension '{ext}'. Must be one of {sorted(EPHYS_EXTENSIONS)}")

    # ------------------------------------------------------------------
    # Primary path: sessions/
    # ------------------------------------------------------------------

    @property
    def session_dir(self) -> Path:
        """Directory containing this session's analysis files.

        Returns
        -------
        Path
            ``{root}/sessions/{session_id}``

        Raises
        ------
        ValueError
            If ``session_id`` is not set.
        """
        if self.session_id is None:
            raise ValueError("session_id must be set to access session_dir")
        return self.root / "sessions" / self.session_id

    @property
    def fpath(self) -> Path:
        """Full path to the session-level analysis file.

        Returns
        -------
        Path
            ``{root}/sessions/{session_id}/{dtype}_{session_id}[_probe{N}].{ext}``

        Raises
        ------
        ValueError
            If ``session_id`` or ``dtype`` is not set.
        """
        if self.session_id is None:
            raise ValueError("session_id must be set to construct fpath")
        if self.dtype is None:
            raise ValueError("dtype must be set to construct fpath")
        probe_tag = f"_probe{self.probe}" if self.probe is not None else ""
        filename = f"{self.dtype}_{self.session_id}{probe_tag}"
        ext = (self.extension or "h5").lstrip(".")
        return self.session_dir / f"{filename}.{ext}"

    # ------------------------------------------------------------------
    # Auxiliary paths: raw/ and nwb/
    # ------------------------------------------------------------------

    @property
    def raw_dir(self) -> Path:
        """Raw data directory for this session.

        Returns
        -------
        Path
            ``{root}/raw/{session_id}``

        Raises
        ------
        ValueError
            If ``session_id`` is not set.
        """
        if self.session_id is None:
            raise ValueError("session_id must be set to access raw_dir")
        return self.root / "raw" / self.session_id

    @property
    def nwb_path(self) -> Path:
        """Path to the NWB intermediate file for this session.

        Returns
        -------
        Path
            ``{root}/nwb/{session_id}[_probe{N}].nwb``

        Raises
        ------
        ValueError
            If ``session_id`` is not set.
        """
        if self.session_id is None:
            raise ValueError("session_id must be set to access nwb_path")
        probe_tag = f"_probe{self.probe}" if self.probe is not None else ""
        return self.root / "nwb" / f"{self.session_id}{probe_tag}.nwb"

    # ------------------------------------------------------------------
    # Convenience constructor
    # ------------------------------------------------------------------

    @classmethod
    def from_components(
        cls,
        root: str | Path,
        date: str,
        subject: str,
        paradigm: str,
        region: str,
        dtype: str | None = None,
        probe: int | None = None,
        extension: str = "h5",
    ) -> EphysPath:
        """Build an EphysPath from individual session components.

        Parameters
        ----------
        root : str or Path
            Root project directory.
        date : str
            Session date, e.g. ``"251024"``.
        subject : str
            Subject name, e.g. ``"FanFan"``.
        paradigm : str
            Paradigm string including optional block/run marker,
            e.g. ``"nsd1w"`` (paradigm ``nsd``, block ``1w``).
        region : str
            Brain region, e.g. ``"MSB"``.
        dtype : str, optional
            Data type (see :data:`EPHYS_DTYPES`).
        probe : int, optional
            Probe index (0-based).
        extension : str
            File extension.  Default ``"h5"``.

        Returns
        -------
        EphysPath
        """
        session_id = f"{date}_{subject}_{paradigm}_{region}"
        return cls(
            root=Path(root),
            session_id=session_id,
            dtype=dtype,
            probe=probe,
            extension=extension,
        )

    def load(self, pre_load: bool = False) -> BaseData:
        """Load this Ephys session into a :class:`~vneurotk.neuro.base.BaseData`.

        Parameters
        ----------
        pre_load : bool
            If ``True``, eagerly load neuro data into memory before returning.

        Returns
        -------
        BaseData
        """
        from vneurotk.io.loader import _load_from_ephys

        bd = _load_from_ephys(self)
        return bd.load() if pre_load else bd

@dataclass
class MNEPath(VTKPath):
    """Path class for MNE data.

    Inherits all attributes from VTKPath.
    Sets modality to 'mne' by default.
    Constructs MNE-style file paths.
    """

    modality: str = field(default="mne", init=False)

    @property
    def fpath(self) -> Path:
        """Construct MNE-style file path.

        Returns
        -------
        Path
            Full file path in MNE format:
            root/sub-{subject}_ses-{session}_task-{task}_run-{run}_{suffix}.{extension}
        """
        parts = []
        if self.subject is not None:
            parts.append(f"sub-{self.subject}")
        if self.session is not None:
            parts.append(f"ses-{self.session}")
        if self.task is not None:
            parts.append(f"task-{self.task}")
        if self.run is not None:
            parts.append(f"run-{self.run}")

        filename = "_".join(parts) if parts else "data"

        if self.suffix is not None:
            filename = f"{filename}_{self.suffix}"

        ext = self.extension if self.extension else ".fif"
        if not ext.startswith("."):
            ext = f".{ext}"

        return self.root / f"{filename}{ext}"

    def load(self, pre_load: bool = False) -> BaseData:
        """Load this MNE recording into a :class:`~vneurotk.neuro.base.BaseData`.

        Parameters
        ----------
        pre_load : bool
            If ``True``, eagerly load neuro data into memory before returning.

        Returns
        -------
        BaseData
        """
        from vneurotk.io.loader import _load_from_mne

        bd = _load_from_mne(self)
        return bd.load() if pre_load else bd

@dataclass
class VTKPath:
    """Base path class for VneuroTK data sources.

    Attributes
    ----------
    root : Path
        Root directory for the data.
    session : str | None
        Session identifier.
    subject : str | None
        Subject identifier.
    task : str | None
        Task identifier.
    run : str | None
        Run identifier.
    desc : str | None
        Description identifier.
    probe : int | None
        Probe number (ephys only).
    suffix : str | None
        File suffix.
    extension : str | None
        File extension.
    modality : str | None
        Data modality (ephys, mne, bids).
    """

    root: Path
    session: str | None = None
    subject: str | None = None
    task: str | None = None
    run: str | None = None
    desc: str | None = None
    probe: int | None = None
    suffix: str | None = None
    extension: str | None = None
    modality: str | None = None

    def __post_init__(self) -> None:
        """Convert root to Path if needed."""
        if not isinstance(self.root, Path):
            self.root = Path(self.root)

    @property
    def fpath(self) -> Path:
        """Construct full file path.

        Returns
        -------
        Path
            Full file path constructed from attributes.
        """
        # Build filename from non-None attributes
        parts = []
        if self.subject is not None:
            parts.append(f"sub-{self.subject}")
        if self.session is not None:
            parts.append(f"ses-{self.session}")
        if self.task is not None:
            parts.append(f"task-{self.task}")
        if self.run is not None:
            parts.append(f"run-{self.run}")
        if self.desc is not None:
            parts.append(f"desc-{self.desc}")
        if self.probe is not None:
            parts.append(f"probe-{self.probe}")
        if self.suffix is not None:
            parts.append(self.suffix)

        filename = "_".join(parts) if parts else "data"

        # Add extension
        ext = self.extension if self.extension else ".h5"
        if not ext.startswith("."):
            ext = f".{ext}"

        return self.root / f"{filename}{ext}"

    def load(self, pre_load: bool = False) -> BaseData:
        """Load data described by this path into a :class:`BaseData`.

        For base :class:`VTKPath` instances pointing to a saved ``.h5`` file,
        loads via the internal HDF5 reader.  Typed subclasses override this
        method with format-specific loading logic.

        Raises
        ------
        NotImplementedError
            If this path does not point to a ``.h5`` file and has no typed
            loading strategy (i.e., it is a plain :class:`VTKPath`).
        """
        fpath = self.fpath
        if fpath.suffix == ".h5":
            from vneurotk.io.loader import _load_from_h5

            bd = _load_from_h5(fpath)
            return bd.load() if pre_load else bd
        raise NotImplementedError(
            f"{type(self).__name__} does not implement load() for non-HDF5 paths. "
            "Use a typed subclass (EphysPath, MNEPath, BIDSPath)."
        )

@dataclass
class BIDSPath(VTKPath):
    """Path class for BIDS data.

    Inherits all attributes from VTKPath.
    Sets modality to 'bids' by default.
    Wraps mne_bids.BIDSPath internally.
    """

    modality: str = field(default="bids", init=False)
    _bids_path: Any = field(default=None, init=False, repr=False)

    def __post_init__(self) -> None:
        """Initialize BIDS path wrapper."""
        super().__post_init__()
        try:
            from mne_bids import BIDSPath as MNEBIDSPath  # type: ignore

            self._bids_path = MNEBIDSPath(
                root=self.root,
                subject=self.subject,
                session=self.session,
                task=self.task,
                run=self.run,
                suffix=self.suffix,
                extension=self.extension,
            )
        except ImportError:
            logger.warning("mne_bids not available, BIDSPath functionality limited")
            self._bids_path = None

    @property
    def fpath(self) -> Path:
        """Get BIDS file path.

        Returns
        -------
        Path
            Full BIDS file path.
        """
        if self._bids_path is not None:
            return Path(self._bids_path.fpath)
        else:
            # Fallback to basic path construction
            return super().fpath

    @property
    def bids_path(self) -> Any:
        """Get underlying mne_bids.BIDSPath object.

        Returns
        -------
        mne_bids.BIDSPath
            The wrapped BIDS path object.
        """
        return self._bids_path

    def load(self, pre_load: bool = False) -> BaseData:
        """Load this BIDS recording into a :class:`~vneurotk.neuro.base.BaseData`.

        Parameters
        ----------
        pre_load : bool
            If ``True``, eagerly load neuro data into memory before returning.

        Returns
        -------
        BaseData
        """
        from vneurotk.io.loader import _load_from_bids

        bd = _load_from_bids(self)
        return bd.load() if pre_load else bd

class LazyNeuroLoader:
    """Deferred loader for a neuro array — invokes *loader_fn* on first call.

    Wraps any ``Callable[[], np.ndarray]`` and guarantees the underlying
    function is executed at most once.  The result is cached so subsequent
    calls return the same array without re-reading from disk.

    Satisfies the ``NeuroLoader = Callable[[], np.ndarray]`` contract, so it
    can be passed directly to :meth:`~vneurotk.neuro.base.BaseData.set_neuro_loader`.

    Parameters
    ----------
    loader_fn : Callable[[], np.ndarray]
        Zero-argument function that reads and returns the neuro array.

    Examples
    --------
    >>> loader = LazyNeuroLoader(lambda: np.load("data.npy"))
    >>> bd.set_neuro_loader(loader)  # called at most once on first bd.neuro access
    """

    def __init__(self, loader_fn: Callable[[], np.ndarray]) -> None:
        self._loader_fn: Callable[[], np.ndarray] | None = loader_fn
        self._data: object = _MISSING

    def __call__(self) -> np.ndarray:
        if self._data is _MISSING:
            self._data = self._loader_fn()  # ty: ignore[call-non-callable]
            self._loader_fn = None  # release closure captures
        return self._data  # ty: ignore[invalid-return-type]

    @property
    def is_loaded(self) -> bool:
        """``True`` after the first call has materialised the array."""
        return self._data is not _MISSING

    def __repr__(self) -> str:
        return f"LazyNeuroLoader(loaded={self.is_loaded})"

class LazyH5Dict(Mapping):
    """HDF5-backed read-only image dict that loads arrays on demand.

    The key index is built on first access (one lightweight pass over
    attribute metadata only).  Each ``__getitem__`` call opens the file,
    reads a single dataset, and closes it immediately.

    Parameters
    ----------
    path : Path or str
        Path to the HDF5 file.
    group : str
        HDF5 group name that contains the image datasets.
        Default ``"stimuli_db"``.
    """

    def __init__(self, path: Path | str, group: str = "stimuli_db") -> None:
        self._path = Path(path)
        self._group = group
        self._index: dict[Any, str] | None = None  # native_key → ds_key

    # ------------------------------------------------------------------
    # Index build (reads only attrs, not image data)
    # ------------------------------------------------------------------

    def _build_index(self) -> None:
        if self._index is not None:
            return
        idx: dict[Any, str] = {}
        with h5py.File(self._path, "r") as f:
            grp = f[self._group]
            for ds_key in grp:
                key_type = str(grp[ds_key].attrs.get("key_type", "str"))
                native: Any = ds_key
                if key_type == "int":
                    try:
                        native = int(ds_key)
                    except ValueError:
                        pass
                elif key_type == "float":
                    try:
                        native = float(ds_key)
                    except ValueError:
                        pass
                idx[native] = ds_key
        self._index = idx

    # ------------------------------------------------------------------
    # Mapping interface
    # ------------------------------------------------------------------

    def __getitem__(self, key: Any) -> np.ndarray:
        self._build_index()
        native = key.item() if hasattr(key, "item") else key
        ds_key = self._index[native]  # ty: ignore[not-subscriptable]
        with h5py.File(self._path, "r") as f:
            return self._decode_item(f[self._group][ds_key])

    @staticmethod
    def _decode_item(ds: Any) -> np.ndarray:
        """Decode a stored dataset to a numpy array based on its ``kind`` attribute.

        Parameters
        ----------
        ds : h5py.Dataset
            An open dataset from the stimuli_db group.

        Returns
        -------
        np.ndarray
        """
        kind = str(ds.attrs.get("kind", "array"))
        data = ds[()] if kind in ("path", "image_bytes") else ds[:]
        return _decode_image(data, kind)

    def __len__(self) -> int:
        self._build_index()
        return len(self._index)  # ty: ignore[invalid-argument-type]

    def __iter__(self):
        self._build_index()
        return iter(self._index)  # ty: ignore[no-matching-overload]

    def __repr__(self) -> str:
        self._build_index()
        return f"LazyH5Dict(n={len(self._index)}, path={self._path.name!r})"  # ty: ignore[invalid-argument-type]