"""Torch dataset over an OriginIndex with feature-set masking (ablations) and a padded collate."""
from __future__ import annotations

import numpy as np
import torch
from torch.utils.data import Dataset

from ..data.windows import OriginIndex, BAR_FEATURES

FEATURE_SETS = {
    "full": None,
    "returns_only": ["ret", "missing", "boundary", "gap", "minute_frac", "time_to_close", "weekday", "early_close"],
    "ohlcv": [f for f in BAR_FEATURES if f != "relvol"],
    "time_only": ["missing", "boundary", "gap", "minute_frac", "time_to_close", "weekday", "early_close"],
}


class OriginDataset(Dataset):
    def __init__(self, index: OriginIndex, feature_set: str = "full", shuffle_minutes: bool = False,
                 shuffle_seed: int = 0):
        self.index, self.shuffle_minutes = index, shuffle_minutes
        keep = FEATURE_SETS[feature_set]
        self.fmask = None if keep is None else np.array([f in keep for f in BAR_FEATURES], dtype=np.float32)
        self.rng = np.random.default_rng(shuffle_seed)

    def __len__(self):
        return len(self.index)

    def __getitem__(self, k: int) -> dict[str, np.ndarray]:
        ex = self.index.example(k)
        if self.fmask is not None:
            for b in ("xA", "xB", "xC"):
                ex[b] = ex[b] * self.fmask
            if "ret" not in (FEATURE_SETS and []):
                pass
        if self.shuffle_minutes:   # falsification: destroy temporal order inside the 1-minute window
            perm = self.rng.permutation(ex["xA"].shape[0])
            ex["xA"] = ex["xA"][perm]; ex["mA"] = ex["mA"][perm]
        return ex


def collate(items: list[dict[str, np.ndarray]]) -> dict[str, torch.Tensor]:
    out = {}
    for k in items[0]:
        arr = np.stack([it[k] for it in items])
        out[k] = torch.from_numpy(arr)
    return out