"""Mock EUPE-ViT-B backbone matching the real model's output contract.

Used to unit-test speedups without loading the 86M-parameter real backbone.
Matches:
  - Input:  (B, 3, H, W) images at 640×640 (stride-16 backbone → 40×40 patches)
  - Output: dict with x_norm_patchtokens = (B, H*W, 768), x_norm_clstoken = (B, 768)
  - Spatial feature shape: (B, 768, 40, 40) after reshape
  - Intermediate block hooks: 12 blocks, each outputs (B, 201, 768) [CLS + 4 storage + 196 patches]
  - Storage-token placement: positions 1..4 in the token sequence
  - Dtype: float32 (real backbone uses bf16 under autocast, test doesn't need that)

Generates deterministic features given a seed, for reproducible tests.
"""
import torch
import torch.nn as nn


class MockEUPEBackbone(nn.Module):
    def __init__(self, seed=0):
        super().__init__()
        self.embed_dim = 768
        self.patch_size = 16
        self.n_storage_tokens = 4
        self.n_blocks = 12
        self._seed = seed
        # Register 12 dummy block modules so forward hooks can attach to
        # `backbone.blocks[i]` exactly like the real backbone.
        self.blocks = nn.ModuleList([nn.Identity() for _ in range(self.n_blocks)])
        # Dummy norm layer for the final LN.
        self.norm = nn.LayerNorm(self.embed_dim)
        # Cached outputs per (image_id, block_idx) so repeat hooks see stable values.
        self._feature_cache = {}

    def forward_features(self, x):
        B, C, H, W = x.shape
        assert C == 3, f"Expected 3-channel input, got {C}"
        Hp, Wp = H // self.patch_size, W // self.patch_size
        N = Hp * Wp
        total = 1 + self.n_storage_tokens + N  # CLS + storage + patches
        # Generate deterministic features per-image via a content-derived seed.
        g = torch.Generator(device=x.device).manual_seed(self._seed + int(x.sum().item() * 1e6) % (2 ** 31))
        tokens = torch.randn(B, total, self.embed_dim, generator=g, device=x.device)
        # Emit a forward pass through each "block" so hooks fire with full token seq.
        for i, blk in enumerate(self.blocks):
            tokens = blk(tokens)
        # Final LayerNorm
        normed = self.norm(tokens)
        return {
            "x_norm_clstoken": normed[:, 0],
            "x_storage_tokens": normed[:, 1:1 + self.n_storage_tokens],
            "x_norm_patchtokens": normed[:, 1 + self.n_storage_tokens:],
            "x_prenorm": tokens,
        }


def make_mock_features(B=4, H=40, W=40, feat_dim=768, device="cuda", seed=0):
    """Generate mock spatial features directly at (B, feat_dim, H, W) — skips the
    mock backbone forward pass when only the final patch features are needed."""
    torch.manual_seed(seed)
    return torch.randn(B, feat_dim, H, W, device=device)


def make_mock_boxes(B=4, n_boxes_per_image=8, resolution=640, device="cuda", seed=0):
    """Generate mock GT boxes and labels for testing the loss computation.
    Boxes are (x1, y1, x2, y2), labels are in [0, 80)."""
    torch.manual_seed(seed)
    boxes_list = []
    labels_list = []
    for _ in range(B):
        # Random center + size, clipped to [0, resolution]
        cx = torch.rand(n_boxes_per_image, device=device) * resolution
        cy = torch.rand(n_boxes_per_image, device=device) * resolution
        w = torch.rand(n_boxes_per_image, device=device) * (resolution / 3) + 16
        h = torch.rand(n_boxes_per_image, device=device) * (resolution / 3) + 16
        x1 = (cx - w / 2).clamp(0, resolution)
        y1 = (cy - h / 2).clamp(0, resolution)
        x2 = (cx + w / 2).clamp(0, resolution)
        y2 = (cy + h / 2).clamp(0, resolution)
        boxes = torch.stack([x1, y1, x2, y2], dim=1)
        labels = torch.randint(0, 80, (n_boxes_per_image,), device=device)
        boxes_list.append(boxes)
        labels_list.append(labels)
    return boxes_list, labels_list


if __name__ == "__main__":
    # Self-test: confirm mock backbone produces correctly-shaped outputs.
    import sys
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Mock backbone self-test on {device}")
    bb = MockEUPEBackbone().to(device).eval()
    x = torch.randn(2, 3, 640, 640, device=device)
    with torch.no_grad():
        out = bb.forward_features(x)
    assert out["x_norm_patchtokens"].shape == (2, 1600, 768)
    assert out["x_norm_clstoken"].shape == (2, 768)
    assert out["x_storage_tokens"].shape == (2, 4, 768)
    print("  forward_features shapes: OK")
    # Hook test
    captured = []
    def hook(m, i, o):
        captured.append(o.shape)
    bb.blocks[5].register_forward_hook(hook)
    with torch.no_grad():
        bb.forward_features(x)
    assert captured[-1] == (2, 201, 768), f"Expected (2, 201, 768), got {captured[-1]}"
    print("  block hook shape: OK")
    feats = make_mock_features(B=4, device=device)
    assert feats.shape == (4, 768, 40, 40)
    print("  make_mock_features shape: OK")
    boxes, labels = make_mock_boxes(B=4, n_boxes_per_image=8, device=device)
    assert len(boxes) == 4 and boxes[0].shape == (8, 4)
    assert len(labels) == 4 and labels[0].shape == (8,)
    print("  make_mock_boxes shape: OK")
    print("\nAll mock-backbone self-tests passed.")