"""Data loading and validation for the local SEEDS data contract."""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, Optional

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset, DistributedSampler


class SEEDSDataset(Dataset):
    def __init__(self, path: str | Path, channels: int, faces: int, height: int, width: int, seed_count: int) -> None:
        self.path = Path(path)
        if not self.path.exists():
            raise FileNotFoundError(f"dataset does not exist: {self.path}")
        archive = np.load(self.path, allow_pickle=False)
        self.seeds = np.asarray(archive["seeds"], dtype=np.float32)
        self.targets = np.asarray(archive["targets"], dtype=np.float32)
        self.climate = np.asarray(archive["climate"], dtype=np.float32)
        expected = (channels, faces, height, width)
        if self.seeds.ndim != 6 or tuple(self.seeds.shape[2:]) != expected or self.seeds.shape[1] != seed_count:
            raise ValueError(f"seeds in {self.path} must have shape [N, {seed_count}, {expected}]")
        if self.targets.ndim != 5 or tuple(self.targets.shape[1:]) != expected:
            raise ValueError(f"targets in {self.path} must have shape [N, {expected}]")
        if self.climate.ndim != 5 or tuple(self.climate.shape[1:]) != expected:
            raise ValueError(f"climate in {self.path} must have shape [N, {expected}]")
        if len(self.seeds) != len(self.targets) or len(self.targets) != len(self.climate):
            raise ValueError("seeds, targets and climate must have the same sample count")
        if not np.isfinite(self.seeds).all() or not np.isfinite(self.targets).all() or not np.isfinite(self.climate).all():
            raise ValueError(f"dataset contains NaN or Inf: {self.path}")

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

    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
        return {
            "seeds": torch.from_numpy(self.seeds[index]),
            "targets": torch.from_numpy(self.targets[index]),
            "climate": torch.from_numpy(self.climate[index]),
        }


def build_dataloader(
    path: str | Path,
    channels: int,
    faces: int,
    height: int,
    width: int,
    seed_count: int,
    batch_size: int,
    shuffle: bool,
    num_workers: int = 0,
    max_batches: Optional[int] = None,
    sampler=None,
    dataset: Optional[SEEDSDataset] = None,
) -> Iterable[Dict[str, torch.Tensor]]:
    if dataset is None:
        dataset = SEEDSDataset(path, channels, faces, height, width, seed_count)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle if sampler is None else False,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
    )
    if max_batches is None:
        return loader
    return (batch for batch_index, batch in enumerate(loader) if batch_index < max_batches)