from __future__ import annotations

"""MONA and LoRA adapters for DINOv3 and DGTRS-CLIP.

MONA (Multi-Cognitive One-Shot Nested Adaptation, CVPR 2025):
  Vision-specific adapter with multi-scale depthwise convolutions.
  Applied to DINOv3 ViT-L/16 (after MSA and MLP in each block).
  Original: github.com/Leiyi-Hu/mona (MIT license).

LoRA (Low-Rank Adaptation):
  Low-rank matrices on Q/V attention projections.
  Applied to DGTRS-CLIP text encoder (all 12 transformer blocks).
"""

import logging
import math

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

LOGGER = logging.getLogger("caption_test.adapters")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")


# ---------------------------------------------------------------------------
# MONA adapter (for DINOv3 ViT — 2D vision tokens)
# ---------------------------------------------------------------------------

class MonaOp(nn.Module):
    """Multi-cognitive visual filter: parallel depthwise convs (3×3, 5×5, 7×7)."""

    def __init__(self, channels: int) -> None:
        super().__init__()
        self.conv3 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels)
        self.conv5 = nn.Conv2d(channels, channels, kernel_size=5, padding=2, groups=channels)
        self.conv7 = nn.Conv2d(channels, channels, kernel_size=7, padding=3, groups=channels)
        self.projector = nn.Conv2d(channels, channels, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        x = (self.conv3(x) + self.conv5(x) + self.conv7(x)) / 3.0 + identity
        return x + self.projector(x)


class MonaAdapter(nn.Module):
    """MONA adapter module for ViT blocks.

    ScaledLayerNorm → Down(dim→bottleneck) → reshape 2D → MonaOp → reshape 1D
    → GELU → Dropout → Up(bottleneck→dim) + residual

    Args:
        dim: Input feature dimension (e.g. 1024 for ViT-L).
        bottleneck: Bottleneck dimension for down/up projections.
        dropout: Dropout rate.
    """

    def __init__(self, dim: int = 1024, bottleneck: int = 64, dropout: float = 0.1) -> None:
        super().__init__()
        self.down = nn.Linear(dim, bottleneck)
        self.up = nn.Linear(bottleneck, dim)
        self.mona_op = MonaOp(bottleneck)
        self.norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(p=dropout)
        # Scaled init: gamma ≈ 0 at start (near-identity), gammax ≈ 1.
        self.gamma = nn.Parameter(torch.ones(dim) * 1e-6)
        self.gammax = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor:
        """Apply MONA adapter (runs in fp32 to avoid AMP overflow).

        Args:
            x: Token features [B, N, D] where N includes CLS + register + patch tokens.
            hw: (H, W) spatial shape of patch tokens.

        Returns:
            Adapted features [B, N, D] (same shape, residual connection).
        """
        with torch.amp.autocast("cuda", enabled=False):
            return self._forward_fp32(x.float(), hw)

    def _forward_fp32(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor:
        identity = x
        # Scaled LayerNorm.
        x = self.norm(x) * self.gamma + x * self.gammax

        x = self.down(x)  # [B, N, bottleneck]

        B, N, C = x.shape
        H, W = hw
        n_special = N - H * W  # CLS + register tokens

        # Separate special tokens (CLS, registers) from patch tokens.
        special = x[:, :n_special]  # [B, n_special, C]
        patches = x[:, n_special:]  # [B, H*W, C]

        # Reshape patches to 2D for convolutions.
        patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2)  # [B, C, H, W]
        patches = self.mona_op(patches)
        patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C)  # [B, H*W, C]

        # Recombine.
        x = torch.cat([special, patches], dim=1)  # [B, N, C]

        x = F.gelu(x)
        x = self.dropout(x)
        x = self.up(x)  # [B, N, D]

        return identity + x


# ---------------------------------------------------------------------------
# LoRA adapter (for DGTRS-CLIP text encoder — 1D text tokens)
# ---------------------------------------------------------------------------

class LoRALinear(nn.Module):
    """LoRA adapter for a linear layer: output = W·x + (B·A)·x.

    Args:
        in_features: Input dimension.
        out_features: Output dimension.
        rank: LoRA rank (low-rank decomposition dimension).
        alpha: LoRA scaling factor (effective scale = alpha / rank).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 4,
        alpha: float = 1.0,
    ) -> None:
        super().__init__()
        self.rank = rank
        self.scale = alpha / rank
        # A: down-projection, B: up-projection. Init: A ~ N(0, 1/rank), B = 0.
        self.A = nn.Parameter(torch.randn(rank, in_features) / math.sqrt(rank))
        self.B = nn.Parameter(torch.zeros(out_features, rank))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Compute LoRA delta: (B @ A) @ x * scale. Always fp32."""
        with torch.amp.autocast("cuda", enabled=False):
            x = x.float()
            return (x @ self.A.t() @ self.B.t()) * self.scale


# ---------------------------------------------------------------------------
# Injection functions — apply adapters to existing frozen models
# ---------------------------------------------------------------------------

def inject_mona_into_dinov3(
    model: nn.Module,
    bottleneck: int = 64,
    dropout: float = 0.1,
) -> int:
    """Inject MONA adapters into a frozen DINOv3ViT model.

    Adds two MonaAdapter modules per block (after attention, after MLP).
    Only adapter parameters are trainable.

    Args:
        model: DINOv3ViT model (already frozen).
        bottleneck: MONA bottleneck dimension.
        dropout: MONA dropout rate.

    Returns:
        Number of trainable adapter parameters added.
    """
    dim = model.embed_dim
    n_adapters = 0

    for block in model.layer:
        # Create adapters.
        block.mona_attn = MonaAdapter(dim, bottleneck, dropout)
        block.mona_mlp = MonaAdapter(dim, bottleneck, dropout)

        # Wrap the original forward to inject adapters.
        original_forward = block.forward

        def make_mona_forward(blk, orig_fwd):
            def mona_forward(x: torch.Tensor) -> torch.Tensor:
                B, N, C = x.shape
                # DINOv3: [CLS, 4 registers, H*W patches]
                n_special = 1 + 4  # CLS + registers
                n_patches = N - n_special
                H = W = int(math.sqrt(n_patches))

                # Original block forward (attention + MLP with layer scale).
                x = x + blk.layer_scale1(blk.attention(blk.norm1(x)))
                x = blk.mona_attn(x, (H, W))
                x = x + blk.layer_scale2(blk.mlp(blk.norm2(x)))
                x = blk.mona_mlp(x, (H, W))
                return x
            return mona_forward

        block.forward = make_mona_forward(block, original_forward)
        n_adapters += 2

    n_params = sum(p.numel() for n, p in model.named_parameters() if "mona" in n)
    LOGGER.info(
        "🔧 MONA injected: %d adapters (%d per block), %s trainable params (bottleneck=%d)",
        n_adapters, 2, f"{n_params:,}", bottleneck,
    )
    return n_params


def inject_lora_into_dgtrs(
    model: nn.Module,
    rank: int = 4,
    alpha: float = 1.0,
) -> int:
    """Inject LoRA adapters into DGTRS-CLIP text encoder attention layers.

    Adds LoRA to Q and V projections in each ResidualAttentionBlock.
    Original weights stay frozen, LoRA params are trainable.

    Args:
        model: DGTRSTextEncoder model.
        rank: LoRA rank.
        alpha: LoRA scaling factor.

    Returns:
        Number of trainable LoRA parameters added.
    """
    n_adapted = 0

    for i, block in enumerate(model.transformer.resblocks):
        attn = block.attn
        # nn.MultiheadAttention uses in_proj_weight [3*dim, dim] for Q, K, V.
        embed_dim = attn.embed_dim

        # LoRA on Q and V (not K — standard practice).
        block.lora_q = LoRALinear(embed_dim, embed_dim, rank, alpha)
        block.lora_v = LoRALinear(embed_dim, embed_dim, rank, alpha)

        # Wrap attention forward to add LoRA deltas.
        original_attention = block.attention

        def make_lora_attention(blk):
            def lora_attention(x: torch.Tensor) -> torch.Tensor:
                # x shape: [T, B, D] (sequence-first for nn.MultiheadAttention)
                attn_mod = blk.attn
                attn_mask = blk.attn_mask
                if attn_mask is not None:
                    attn_mask = attn_mask.to(dtype=x.dtype, device=x.device)

                # Compute QKV using original in_proj_weight.
                T, B, D = x.shape

                # Original QKV.
                qkv = F.linear(x, attn_mod.in_proj_weight, attn_mod.in_proj_bias)
                q, k, v = qkv.chunk(3, dim=-1)

                # Add LoRA deltas to Q and V.
                # LoRA expects [B, T, D], but x is [T, B, D].
                x_btd = x.permute(1, 0, 2)  # [B, T, D]
                q = q + blk.lora_q(x_btd).permute(1, 0, 2)
                v = v + blk.lora_v(x_btd).permute(1, 0, 2)

                # Multi-head attention.
                num_heads = attn_mod.num_heads
                head_dim = D // num_heads
                q = q.reshape(T, B * num_heads, head_dim).transpose(0, 1)
                k = k.reshape(T, B * num_heads, head_dim).transpose(0, 1)
                v = v.reshape(T, B * num_heads, head_dim).transpose(0, 1)

                attn_output = F.scaled_dot_product_attention(
                    q, k, v, attn_mask=attn_mask,
                )
                attn_output = attn_output.transpose(0, 1).reshape(T, B, D)
                attn_output = attn_mod.out_proj(attn_output)
                return attn_output
            return lora_attention

        block.attention = make_lora_attention(block)
        n_adapted += 1

    n_params = sum(
        p.numel() for n, p in model.named_parameters() if "lora" in n
    )
    LOGGER.info(
        "🔧 LoRA injected: %d blocks, rank=%d, %s trainable params",
        n_adapted, rank, f"{n_params:,}",
    )
    return n_params