"""Small pretrained models for independent tiles and continuous grids."""

import timm
from torch import nn

BACKBONE = "mobilenetv4_conv_medium.e250_r384_in12k_ft_in1k"


class GridModel(nn.Module):
    def __init__(self, tiles, classes, pretrained=True, backbone=BACKBONE):
        super().__init__()
        if tiles not in (9, 16):
            raise ValueError("Grid size must be 9 or 16 tiles")
        self.tiles = tiles
        if tiles == 9:
            self.backbone = timm.create_model(backbone, pretrained=pretrained)
            self.backbone.reset_classifier(0)
            channels = self.backbone.head_hidden_size
            self.head = nn.Linear(channels, classes)
        else:
            # The stride-16 stage: 448 px input gives a 28x28 map, 7x7 per cell.
            stage = (
                timm.create_model(backbone, features_only=True).feature_info.reduction().index(16)
            )
            self.backbone = timm.create_model(
                backbone, pretrained=pretrained, features_only=True, out_indices=(stage,)
            )
            channels = self.backbone.feature_info.channels()[0]
            self.head = nn.Sequential(
                nn.Conv2d(channels, 128, 3, padding=1),
                nn.ReLU(),
                nn.Conv2d(128, classes, 1),
                # Max, not mean: any part of the target selects the cell.
                nn.AdaptiveMaxPool2d((4, 4)),
            )

    def forward(self, images):
        features = self.backbone(images)
        return self.head(features if self.tiles == 9 else features[0])

    def freeze_backbone(self, frozen):
        for parameter in self.backbone.parameters():
            parameter.requires_grad_(not frozen)


def select_target(logits, targets):
    """Only the requested target is supervised; other targets are unknown."""
    import torch

    return logits[torch.arange(logits.shape[0], device=logits.device), targets]


def transfer_backbone(model, state):
    """Copy matching backbone weights, e.g. from a tile classifier into the 4x4 feature extractor."""
    source = {k.removeprefix("backbone."): v for k, v in state.items() if k.startswith("backbone.")}
    target = model.backbone.state_dict()
    matched = {k: v for k, v in source.items() if k in target and target[k].shape == v.shape}
    missing = sorted(set(target) - set(matched))
    if missing:
        raise ValueError(
            f"Backbone weights missing for {len(missing)} parameters, e.g. {missing[:3]}"
        )
    model.backbone.load_state_dict(matched)
    return len(matched)