src/data_real.py

"""
Loaders for The Well data.

CONTRACT (see provenance.py for shared error types):
  - get_dataset() loads REAL data only. It NEVER silently substitutes
    synthetic data. On any failure it raises DataLoadError or
    SchemaValidationError. This is a deliberate design decision:
    retraining on the wrong data (or on synthetic data believed to be
    real) is worse than a loud failure.
  - Synthetic data is only available via the separate, explicitly-named
    get_synthetic_dataset() — callers must opt in on purpose.
  - Tensor layout is never guessed from shape magnitude. Callers must
    declare the expected channel_layout; a mismatch is a hard error,
    not a silent transpose.

Setup for real data:
  1. Run:  the-well-download --base-path ./data/well --dataset active_matter --split train
  2. Or place .hdf5 files under data/real/ with a known, declared schema.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Optional, List, Literal

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

from .provenance import DataLoadError, SchemaValidationError

ChannelLayout = Literal["channels_first", "channels_last"]


class LocalWellHDF5(Dataset):

    KNOWN_KEYS = ("t0_fields", "fields", "data", "x", "trajectory")

    def __init__(
        self,
        root: str,
        max_samples: int = 256,
        n_steps: Optional[int] = None,
        expected_channels: int = 2,
        channel_layout: ChannelLayout = "channels_first",
        strict: bool = True,
    ):
        self.root = Path(root)
        self.files = sorted(self.root.rglob("*.hdf5")) + sorted(self.root.rglob("*.h5"))
        if not self.files:
            raise DataLoadError(
                f"no .hdf5/.h5 files found under {root}",
                outcome_code="NO_FILES_FOUND",
            )
        self.max_samples = max_samples
        self.n_steps = n_steps
        self.expected_channels = expected_channels
        self.channel_layout = channel_layout
        self.strict = strict
        self.samples: List[torch.Tensor] = []
        self._load()

    def _load(self):
        try:
            import h5py
        except ImportError as e:
            raise DataLoadError(
                "h5py is required to load local Well HDF5 files but is not installed",
                outcome_code="MISSING_DEPENDENCY",
            ) from e

        count = 0
        rejected = []
        for fp in self.files:
            if count >= self.max_samples:
                break
            with h5py.File(fp, "r") as f:
                arr = None
                for k in self.KNOWN_KEYS:
                    if k in f:
                        arr = f[k]
                        break
                if arr is None and list(f.keys()):
                    arr = f[list(f.keys())[0]]
                if arr is None:
                    rejected.append((str(fp), "no recognizable dataset key"))
                    continue

                data = np.array(arr)
                try:
                    if data.ndim == 5:  # N, T, C/H, H/C, W or similar
                        for i in range(min(data.shape[0], self.max_samples - count)):
                            traj = self._validate_and_orient(data[i], fp)
                            self.samples.append(torch.from_numpy(traj).float())
                            count += 1
                    elif data.ndim == 4:
                        traj = self._validate_and_orient(data, fp)
                        self.samples.append(torch.from_numpy(traj).float())
                        count += 1
                    else:
                        rejected.append((str(fp), f"unsupported ndim={data.ndim}"))
                except SchemaValidationError as e:
                    rejected.append((str(fp), e.detail))
                    if self.strict:
                        raise

        if rejected and not self.strict:
            print(f"[LocalWellHDF5] WARNING: {len(rejected)} file(s) rejected: {rejected}")

        if not self.samples:
            raise SchemaValidationError(
                f"no valid trajectories loaded from {self.root}; "
                f"rejected files: {rejected}",
                outcome_code="NO_VALID_TRAJECTORIES",
            )
        print(f"[LocalWellHDF5] loaded {len(self.samples)} trajectories "
              f"from {len(self.files)} files (schema={self.channel_layout}, "
              f"channels={self.expected_channels})")

    def _validate_and_orient(self, data: np.ndarray, fp: Path) -> np.ndarray:
        if data.ndim != 4:
            raise SchemaValidationError(
                f"{fp.name}: expected 4D (T,C,H,W)-like array, got ndim={data.ndim}",
                outcome_code="WRONG_NDIM",
            )
        T, A, B, C_ = data.shape
        if self.channel_layout == "channels_first":
            channel_axis_size = A
            oriented = data
        elif self.channel_layout == "channels_last":
            channel_axis_size = C_
            oriented = np.transpose(data, (0, 3, 1, 2))
        else:
            raise SchemaValidationError(
                f"unknown channel_layout '{self.channel_layout}'",
                outcome_code="INVALID_LAYOUT_SPEC",
            )

        if channel_axis_size != self.expected_channels:
            raise SchemaValidationError(
                f"{fp.name}: declared channel_layout='{self.channel_layout}' "
                f"implies {channel_axis_size} channels, but "
                f"expected_channels={self.expected_channels}. "
                f"Refusing to guess a different layout for this file — "
                f"pass the correct channel_layout/expected_channels explicitly.",
                outcome_code="CHANNEL_COUNT_MISMATCH",
            )
        return oriented

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

    def __getitem__(self, idx):
        traj = self.samples[idx]
        if self.n_steps is not None and traj.size(0) > self.n_steps:
            traj = traj[: self.n_steps]
        return {"fields": traj, "idx": idx}


def get_dataset(
    max_samples: int = 128,
    n_steps: int = 14,
    expected_channels: int = 2,
    channel_layout: ChannelLayout = "channels_first",
    search_roots: Optional[List[str]] = None,
):
    roots = search_roots or ["./data/real", "./data/well"]
    existing_dirs = [r for r in roots if os.path.isdir(r)]
    if not existing_dirs:
        raise DataLoadError(
            f"no real-data directory found among candidates: {roots}. "
            f"Run `the-well-download ...` or place .hdf5 files under one "
            f"of these paths before training.",
            outcome_code="NO_DATA_DIRECTORY",
        )

    last_error = None
    for root in existing_dirs:
        try:
            ds = LocalWellHDF5(
                root,
                max_samples=max_samples,
                n_steps=n_steps,
                expected_channels=expected_channels,
                channel_layout=channel_layout,
                strict=True,
            )
            print(f"[data] using REAL local Well data from {root} ({len(ds)} trajs)")
            return ds, "REAL_LOCAL"
        except (DataLoadError, SchemaValidationError) as e:
            last_error = e
            continue

    raise last_error or DataLoadError(
        f"no usable real data found in {existing_dirs}",
        outcome_code="NO_DATA_DIRECTORY",
    )


def get_synthetic_dataset(max_samples: int = 128, n_steps: int = 14):
    """
    Explicit, opt-in synthetic data. Never called implicitly by get_dataset().
    Returns: (dataset, provenance) where provenance == "SYNTHETIC".
    """
    from .synthetic_fields import SyntheticWellLike
    print("[data] EXPLICIT synthetic mode requested — not real Well data")
    ds = SyntheticWellLike(n_samples=max_samples, n_steps=n_steps, height=32, width=32)
    return ds, "SYNTHETIC"