"""Synthetic GLORYS12-shaped data for smoke tests and local development."""

import glob
import torch
from torch.utils.data import Dataset

try:
    import h5py
except ImportError:
    h5py = None


class SyntheticOceanDataset(Dataset):
    def __init__(self, samples, channels, grid, input_steps=2, output_steps=1, seed=42, data_dir=None):
        generator = torch.Generator().manual_seed(seed)
        files = sorted(glob.glob(f"{data_dir}/data/*.h5")) if data_dir else []
        if files and h5py is not None:
            with h5py.File(files[0], "r") as handle:
                fields = torch.from_numpy(handle["fields"][:]).float()
            total = min(samples, fields.shape[0] - input_steps - output_steps + 1)
            self.x = torch.stack([fields[i:i + input_steps] for i in range(total)])
            self.y = torch.stack([fields[i + input_steps:i + input_steps + output_steps] for i in range(total)])
        else:
            self.x = torch.randn(samples, input_steps, channels, *grid, generator=generator)
            self.y = torch.randn(samples, output_steps, channels, *grid, generator=generator)

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]