"""Deep Evolved Detection Head.

10-layer MLP with interleaved depthwise 3x3 convolutions operating on
92 evolutionarily-selected feature dimensions. The dimension selection
was performed via GPU-batched evolutionary search (200 gen/s). The MLP
was trained on the selected dimensions with the backbone frozen.

182K params, 10.6 mAP, mAP@0.75 = 10.8.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

NUM_CLASSES = 80


def cofiber_decompose(f, n_scales):
    cofibers = []; residual = f
    for _ in range(n_scales - 1):
        omega = F.avg_pool2d(residual, 2)
        sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False)
        cofibers.append(residual - sigma_omega); residual = omega
    cofibers.append(residual); return cofibers


class SpatialDWConv(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, 3, padding=1, groups=channels)

    def forward(self, x, B, H, W):
        if x.dim() == 4:
            x = x.permute(0, 3, 1, 2)
            x = self.conv(x)
            x = x.permute(0, 2, 3, 1)
        return x


class EvolvedDeepHead(nn.Module):
    name = "evolved_deep"
    needs_intermediates = False

    def __init__(self, evolved_dims, hidden=128, n_layers=10, n_scales=3):
        super().__init__()
        self.evolved_dims = evolved_dims
        self.n_scales = n_scales
        K = len(evolved_dims)
        self.dim_idx = nn.Parameter(torch.tensor(evolved_dims, dtype=torch.long), requires_grad=False)
        self.scale_norms = nn.ModuleList([nn.LayerNorm(768) for _ in range(n_scales)])

        layers = []
        in_dim = K
        for i in range(n_layers):
            layers.append(nn.Linear(in_dim, hidden))
            layers.append(nn.GELU())
            if i % 2 == 1:
                layers.append(SpatialDWConv(hidden))
            in_dim = hidden
        self.backbone = nn.Sequential(*layers)

        self.cls_head = nn.Linear(hidden, NUM_CLASSES)
        self.reg_head = nn.Linear(hidden, 4)
        self.ctr_head = nn.Linear(hidden, 1)
        self.scale_params = nn.Parameter(torch.ones(n_scales))

    def forward(self, spatial, inter=None):
        cofibers = cofiber_decompose(spatial, self.n_scales)
        cls_l, reg_l, ctr_l = [], [], []
        for i, cof in enumerate(cofibers):
            B, C, H, W = cof.shape
            f = self.scale_norms[i](cof.permute(0, 2, 3, 1).reshape(-1, C))
            f_sel = f[:, self.dim_idx].reshape(B, H, W, -1)
            h = self._forward_with_spatial(f_sel, B, H, W)
            cls = self.cls_head(h.reshape(-1, h.shape[-1])).reshape(B, H, W, -1).permute(0, 3, 1, 2)
            reg_raw = (self.reg_head(h.reshape(-1, h.shape[-1])) * self.scale_params[i]).clamp(-10, 10)
            reg = reg_raw.exp().reshape(B, H, W, 4).permute(0, 3, 1, 2)
            ctr = self.ctr_head(h.reshape(-1, h.shape[-1])).reshape(B, H, W, 1).permute(0, 3, 1, 2)
            cls_l.append(cls); reg_l.append(reg); ctr_l.append(ctr)
        return cls_l, reg_l, ctr_l

    def _forward_with_spatial(self, x, B, H, W):
        for layer in self.backbone:
            if isinstance(layer, SpatialDWConv):
                x = layer(x, B, H, W)
            else:
                x = layer(x)
        return x