"""Dataset container for NeuroWeave analyses."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator, Protocol

import numpy as np
import pandas as pd

from neuroweave.design.design_matrix import DesignMatrix
from neuroweave.design.lag import make_lagged_matrix
from neuroweave.design.lag_metadata import LagMetadata


class FeatureMatrix(Protocol):
    """Structural protocol for feature containers used by :class:`Dataset`."""

    X: np.ndarray

    @property
    def n_samples(self) -> int:
        """Number of time samples."""

    @property
    def n_features(self) -> int:
        """Number of features."""

    @property
    def feature_names(self) -> list[str]:
        """Names of feature columns."""

    def get_feature(self, name: str) -> np.ndarray:
        """Return a single named feature column."""


@dataclass
class _ArrayFeatureMatrix:
    """Private array-backed feature container used for subsetting."""

    X: np.ndarray
    _feature_names: list[str]

    @property
    def n_samples(self) -> int:
        """Number of time samples."""

        return self.X.shape[0]

    @property
    def n_features(self) -> int:
        """Number of features."""

        return self.X.shape[1]

    @property
    def feature_names(self) -> list[str]:
        """Feature names."""

        return list(self._feature_names)

    def get_feature(self, name: str) -> np.ndarray:
        """Return a named feature column."""

        try:
            index = self._feature_names.index(name)
        except ValueError as exc:
            raise KeyError(f"Unknown feature name: {name}") from exc
        return self.X[:, index].copy()


class Dataset:
    """Container for complete encoding-model analysis inputs.

    Parameters
    ----------
    neural : np.ndarray
        Neural data array with shape ``(time, channels)`` or ``(time,)``.
    features : FeatureMatrix
        Feature container describing the stimulus features.
    sampling_rate : float
        Sampling rate of the neural recordings.
    events : pd.DataFrame | None, optional
        Optional event table.
    metadata : dict | None, optional
        Optional dataset metadata.
    """

    def __init__(
        self,
        neural: np.ndarray,
        features: FeatureMatrix,
        sampling_rate: float,
        events: pd.DataFrame | None = None,
        metadata: dict | None = None,
    ) -> None:
        neural_array = np.asarray(neural, dtype=float)
        if neural_array.ndim == 1:
            neural_array = neural_array[:, np.newaxis]
        if neural_array.ndim != 2:
            raise ValueError("neural must be a 1D or 2D array.")
        if neural_array.shape[0] != features.n_samples:
            raise ValueError("neural and features must have the same number of samples.")

        self.neural = neural_array.copy()
        self.features = features
        self.sampling_rate = float(sampling_rate)
        self.events = events.copy() if events is not None else None
        self.metadata = dict(metadata) if metadata is not None else {}

    @property
    def n_samples(self) -> int:
        """Number of time samples."""

        return self.neural.shape[0]

    @property
    def n_channels(self) -> int:
        """Number of neural channels."""

        return self.neural.shape[1]

    @property
    def n_features(self) -> int:
        """Number of stimulus features."""

        return self.features.n_features

    @property
    def feature_names(self) -> list[str]:
        """Feature names."""

        return self.features.feature_names

    def get_feature(self, name: str) -> np.ndarray:
        """Return a named feature column.

        Parameters
        ----------
        name : str
            Feature name.

        Returns
        -------
        np.ndarray
            Feature time series.
        """

        return self.features.get_feature(name)

    def get_channel(self, index: int) -> np.ndarray:
        """Return a neural channel time series.

        Parameters
        ----------
        index : int
            Zero-based channel index.

        Returns
        -------
        np.ndarray
            Neural signal for the selected channel.
        """

        if not 0 <= index < self.n_channels:
            raise ValueError("Channel index is out of bounds.")
        return self.neural[:, index].copy()

    def events_of_type(self, event_type: str) -> pd.DataFrame:
        """Return events matching a specific type.

        Parameters
        ----------
        event_type : str
            Event type to filter.

        Returns
        -------
        pd.DataFrame
            Filtered event table.
        """

        if self.events is None:
            return pd.DataFrame(columns=["time", "event_type", "value"])
        if "event_type" not in self.events.columns:
            raise ValueError("events table must include an 'event_type' column.")
        return self.events.loc[self.events["event_type"] == event_type].copy()

    def subset_time(self, start: int, end: int) -> Dataset:
        """Return a time-restricted subset of the dataset.

        Parameters
        ----------
        start : int
            Inclusive start sample.
        end : int
            Exclusive end sample.

        Returns
        -------
        Dataset
            Dataset restricted to the requested time window.
        """

        neural_subset = self.neural[start:end, :].copy()
        feature_matrix = np.asarray(self.features.X, dtype=float)[start:end, :].copy()
        feature_subset = _ArrayFeatureMatrix(
            X=feature_matrix,
            _feature_names=self.feature_names,
        )

        events_subset: pd.DataFrame | None = None
        if self.events is not None and "time" in self.events.columns:
            mask = (self.events["time"] >= start) & (self.events["time"] < end)
            events_subset = self.events.loc[mask].copy()
            events_subset.loc[:, "time"] = events_subset["time"] - start
        elif self.events is not None:
            events_subset = self.events.copy()

        return Dataset(
            neural=neural_subset,
            features=feature_subset,
            sampling_rate=self.sampling_rate,
            events=events_subset,
            metadata=self.metadata,
        )

    def make_design_matrix(self, lags: tuple[int, int]) -> DesignMatrix:
        """Construct a lagged design matrix from the dataset features.

        Parameters
        ----------
        lags : tuple[int, int]
            Inclusive lag range ``(min_lag, max_lag)``.

        Returns
        -------
        DesignMatrix
            Lagged design matrix with associated metadata.
        """

        feature_array = np.asarray(self.features.X, dtype=float)
        lagged = make_lagged_matrix(feature_array, lags)
        metadata = LagMetadata(
            n_features=self.n_features,
            lags=lags,
            feature_names=self.feature_names,
        )
        return DesignMatrix(
            X=lagged,
            metadata=metadata,
            sampling_rate=self.sampling_rate,
        )

    def iterate_channels(self) -> Iterator[tuple[int, np.ndarray]]:
        """Yield neural channels one at a time.

        Yields
        ------
        tuple[int, np.ndarray]
            Channel index and channel signal.
        """

        for index in range(self.n_channels):
            yield index, self.neural[:, index].copy()

    def __repr__(self) -> str:
        """Return a readable representation."""

        return (
            f"Dataset(samples={self.n_samples}, "
            f"channels={self.n_channels}, "
            f"features={self.n_features})"
        )
