"""Optional PyTorch Dataset wrapper for therminal weather data.

Provides a torch.utils.data.Dataset that yields feature tensors
for each date. User supplies their own target via a callback or
by subclassing.

Requires: pip install therminal-py[torch]
"""

from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import Any

try:
    import torch
    from torch.utils.data import Dataset
except ImportError:
    raise ImportError(
        "PyTorch is required for TherminalDataset. "
        "Install with: pip install therminal-py[torch]"
    )


from therminal.client import TherminalClient
from therminal.ml.features import WeatherFeatures, _to_dates


class TherminalDataset(Dataset):  # type: ignore[type-arg]
    """PyTorch Dataset backed by WeatherFeatures.

    Each ``__getitem__`` returns a feature tensor for the corresponding date.
    If a ``target_fn`` is provided, returns ``(features, target)`` tuple.

    Parameters
    ----------
    dates : array-like of dates
        Target dates for the dataset.
    station : str
        Therminal station code.
    target_fn : callable or None
        Optional function ``(date_str) -> float`` that returns the target value
        for a given date. If None, ``__getitem__`` returns features only.
    sources : list of str, default ["omo", "metar"]
        Data sources to include.
    lookback_hours : int, default 24
        Feature lookback window.
    client : TherminalClient or None
        Custom client instance.
    **kwargs
        Additional keyword arguments passed to WeatherFeatures.

    Examples
    --------
    >>> ds = TherminalDataset(dates=dates, station="ATL")
    >>> loader = torch.utils.data.DataLoader(ds, batch_size=32, shuffle=True)
    >>> for batch in loader:
    ...     features = batch  # shape (32, n_features)
    """

    def __init__(
        self,
        dates: Any,
        station: str = "ATL",
        target_fn: Callable[[str], float] | None = None,
        sources: Sequence[str] = ("omo", "metar"),
        lookback_hours: int = 24,
        client: TherminalClient | None = None,
        **kwargs: Any,
    ):
        self._dates = _to_dates(dates)
        self._date_strs = [d.strftime("%Y-%m-%d") for d in self._dates]
        self._target_fn = target_fn

        self._wf = WeatherFeatures(
            station=station,
            sources=sources,
            lookback_hours=lookback_hours,
            client=client,
            **kwargs,
        )
        # Prefetch all data
        self._wf.fit(dates)
        self._X = self._wf.transform(dates)

    def __len__(self) -> int:
        return len(self._dates)

    def __getitem__(self, idx: int) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        features = torch.tensor(self._X[idx], dtype=torch.float32)

        if self._target_fn is not None:
            target = torch.tensor(self._target_fn(self._date_strs[idx]), dtype=torch.float32)
            return features, target

        return features

    @property
    def feature_names(self) -> list[str]:
        """Return feature column names."""
        return list(self._wf.get_feature_names_out())

    @property
    def n_features(self) -> int:
        """Return number of features."""
        return self._wf.n_features_
