"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones.

MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we
mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv
{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise,
init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn
and post-mlp.
"""
from __future__ import annotations

import logging
import torch
import torch.nn as nn

from src.models.stripnet.model import StripNet, Block

LOGGER = logging.getLogger("caption_test.stripnet.adapters")


class ConvMona(nn.Module):
    """Single Conv-MONA adapter.

    Args:
        dim: input channel dim.
        bottleneck: bottleneck channel dim (e.g. 64).
        gamma_init: layer-scale init value (1e-6 for near-identity at start).
    """

    def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None:
        super().__init__()
        self.norm = nn.BatchNorm2d(dim)
        self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True)
        self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True)
        self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True)
        self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True)
        self.act = nn.GELU()
        self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True)
        # Channel-wise layer scale (γ), broadcast across H, W.
        self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        h = self.down(h)
        h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h
        h = self.act(h)
        h = self.up(h)
        return self.gamma.view(1, -1, 1, 1) * h


def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona):
    """Closure that wraps a Block.forward with two Conv-MONA residuals."""
    orig_attn = block.attn
    orig_mlp = block.mlp
    orig_norm1 = block.norm1
    orig_norm2 = block.norm2
    orig_drop = block.drop_path
    ls1 = block.layer_scale_1
    ls2 = block.layer_scale_2

    def forward(x: torch.Tensor) -> torch.Tensor:
        x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x)
        x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x)
        return x

    return forward


def inject_conv_mona_into_stripnet(
    model: StripNet,
    bottleneck: int = 64,
    last_n_stages: int = 2,
    use_bf16: bool = False,
) -> int:
    """Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet.

    Each Block in the targeted stages gets two adapters (post-attn, post-mlp).
    Returns the number of adapters injected.

    Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2`
    we adapt block3 and block4 — the deepest, semantically richest features.
    """
    n_stages = model.num_stages
    target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1))
    n_added = 0

    for stage_idx in target_stages:
        blocks: nn.ModuleList = getattr(model, f"block{stage_idx}")
        dim = model.embed_dims[stage_idx - 1]
        for blk_idx, block in enumerate(blocks):
            mona_a = ConvMona(dim=dim, bottleneck=bottleneck)
            mona_m = ConvMona(dim=dim, bottleneck=bottleneck)
            if use_bf16:
                mona_a.to(dtype=torch.bfloat16)
                mona_m.to(dtype=torch.bfloat16)
            # Register as submodules so they get moved by .to(device) / .train() etc.
            block.add_module(f"mona_attn", mona_a)
            block.add_module(f"mona_mlp", mona_m)
            block.forward = _patched_block_forward(block, mona_a, mona_m)
            n_added += 2

    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    LOGGER.info(
        "🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)",
        n_added, target_stages, n_trainable, bottleneck,
    )
    return n_added