"""
Adapter between the real `the_well` package's WellDataset sample format
and this project's `{"fields": (T, C, H, W)}` contract.

VERIFIED against the actual installed `the_well` package source this
session (not guessed): WellDataset._construct_sample() returns a dict
with `input_fields`/`output_fields` keys, shaped `(T, H, W, C)` --
channels LAST -- confirmed both by the library's own docstring comment
and by reading `_postprocess_data()`, which flattens all field types
(scalar/vector/tensor-order) into one combined channel axis via
`.unsqueeze(-1).flatten(...)` then `torch.concatenate(..., dim=-1)`,
i.e. the new combined channel dimension is the LAST dimension.

Every consumer in this project (FieldNormalizer, MultiScaleEncoder,
ReplayBuffer, the collate functions in run_full.py/continual_demo.py)
assumes `{"fields": (T, C, H, W)}` -- channels-FIRST, one contiguous
trajectory, not split into separate input/output tensors. Without this
adapter, real streamed Well data would either crash on first use or
(worse) silently broadcast-mismatch through FieldNormalizer's
`(-1, 1, 1)` per-channel view, since channels-last (T,H,W,C) sliced per
timestep gives (H,W,C), which a channels-first-only broadcast can
sometimes "succeed" against incorrectly if H, W, or C happen to share a
size -- silent wrong results, not a crash. This adapter exists
specifically to prevent that: hard-fail on any shape it doesn't
recognize, never guess a layout.
"""
from __future__ import annotations
from typing import Any, Dict, Optional

import torch
from torch.utils.data import Dataset

from .provenance import SchemaValidationError


def well_sample_to_fields(sample: Dict[str, Any], include_output: bool = True) -> torch.Tensor:
    if "input_fields" not in sample:
        raise SchemaValidationError(
            f"WellDataset sample is missing 'input_fields'. Keys present: "
            f"{sorted(sample.keys())}. This adapter was built against "
            f"the_well's documented input_fields/output_fields contract; "
            f"if that has changed, this needs updating -- not guessing "
            f"a different key.",
            outcome_code="WELL_SAMPLE_MISSING_INPUT_FIELDS",
        )

    parts = [sample["input_fields"]]
    if include_output:
        if "output_fields" not in sample:
            raise SchemaValidationError(
                f"include_output=True but sample has no 'output_fields'. "
                f"Keys present: {sorted(sample.keys())}.",
                outcome_code="WELL_SAMPLE_MISSING_OUTPUT_FIELDS",
            )
        parts.append(sample["output_fields"])

    for i, part in enumerate(parts):
        if not torch.is_tensor(part) or part.dim() != 4:
            raise SchemaValidationError(
                f"WellDataset field tensor #{i} has shape "
                f"{tuple(part.shape) if torch.is_tensor(part) else type(part)}, "
                f"expected 4D (T, H, W, C). Refusing to guess how to "
                f"reinterpret it.",
                outcome_code="WELL_SAMPLE_UNEXPECTED_SHAPE",
            )

    if len(parts) > 1:
        hw_c_shapes = {tuple(p.shape[1:]) for p in parts}
        if len(hw_c_shapes) > 1:
            raise SchemaValidationError(
                f"input_fields and output_fields have mismatched (H,W,C): "
                f"{hw_c_shapes}. Cannot concatenate along time without "
                f"either dropping data or inventing values -- refusing both.",
                outcome_code="WELL_SAMPLE_SHAPE_MISMATCH",
            )

    combined = torch.cat(parts, dim=0)  # (T, H, W, C)
    fields = combined.permute(0, 3, 1, 2).contiguous()  # -> (T, C, H, W)
    return fields


class WellStreamAdapter(Dataset):
    def __init__(self, well_dataset, include_output: bool = True):
        self.well_dataset = well_dataset
        self.include_output = include_output
        self.provenance = getattr(well_dataset, "provenance", None)

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

    def __getitem__(self, idx):
        raw = self.well_dataset[idx]
        fields = well_sample_to_fields(raw, include_output=self.include_output)
        return {"fields": fields, "idx": idx}