from __future__ import annotations

"""Contract tests shared by all fusion implementations."""

import torch

from fuse_proj.data.types import FusionViewOutput, ViewBatch, ViewName
from fuse_proj.models.fusion.base import FusionModelBase


class DummyFusion(FusionModelBase):
    """Minimal parameter-free implementation used only to test the base contract."""

    def encode_view(self, batch: ViewBatch, view: ViewName) -> FusionViewOutput:
        """Convert RGB channel means into a deterministic normalized descriptor."""
        del view
        rgb_summary = batch["rgb"].mean(dim=(-2, -1))  # [B, 3]
        repeats = (self.descriptor_dim + rgb_summary.shape[1] - 1) // rgb_summary.shape[1]
        raw = rgb_summary.repeat(1, repeats)[:, : self.descriptor_dim]  # [B, D]
        descriptor = self.normalize_descriptor(raw)
        contributions = torch.zeros(raw.shape[0], 3, dtype=raw.dtype, device=raw.device)
        return {
            "descriptor": descriptor,
            "rgb_descriptor": descriptor,
            "modality_contributions": contributions,
            "diagnostics": {},
        }


def _make_batch(batch_size: int) -> ViewBatch:
    """Create a valid synthetic multimodal batch."""
    return {
        "rgb": torch.rand(batch_size, 3, 256, 256),
        "geometry": torch.rand(batch_size, 1, 256, 256),
        "segmentation": torch.zeros(batch_size, 1, 256, 256, dtype=torch.uint8),
        "geometry_valid": torch.ones(batch_size, 1, 256, 256, dtype=torch.bool),
        "segmentation_valid": torch.ones(batch_size, 1, 256, 256, dtype=torch.bool),
        "text_valid": torch.ones(batch_size, dtype=torch.bool),
        "text": ["urban scene"] * batch_size,
        "sample_id": [f"sample-{index}" for index in range(batch_size)],
    }


def test_pair_output_shapes() -> None:
    """The base contract must return one normalized descriptor per view."""
    model = DummyFusion(descriptor_dim=1024)
    output = model(_make_batch(4), _make_batch(4))

    for view in ("satellite", "uav"):
        descriptor = output[view]["descriptor"]
        assert descriptor.shape == (4, 1024)
        assert torch.allclose(descriptor.norm(dim=-1), torch.ones(4), atol=1e-5)


def test_batch_size_one() -> None:
    """Fusion implementations must not squeeze away the batch dimension."""
    model = DummyFusion(descriptor_dim=1024)
    output = model(_make_batch(1), _make_batch(1))
    assert output["satellite"]["descriptor"].shape == (1, 1024)
    assert output["uav"]["descriptor"].shape == (1, 1024)