"""StripNet (small) backbone — adapted from Strip-R-CNN (HVision-NKU).

Self-contained: no external utils. State-dict naming follows the upstream
ImageNet-pretrained checkpoint (`conv_spatial1/2` for the strip kernels).
"""
from __future__ import annotations

import logging
import math
from functools import partial
from pathlib import Path
from typing import List

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


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


def _to_2tuple(x):
    if isinstance(x, (tuple, list)):
        return tuple(x)
    return (x, x)


def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
    if drop_prob == 0.0 or not training:
        return x
    keep = 1.0 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
    rand = x.new_empty(shape).bernoulli_(keep)
    if keep > 0:
        rand.div_(keep)
    return x * rand


class DropPath(nn.Module):
    def __init__(self, p: float = 0.0) -> None:
        super().__init__()
        self.p = p

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return drop_path(x, self.p, self.training)


class DWConv(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dwconv(x)


class Mlp(nn.Module):
    def __init__(self, in_features: int, hidden_features: int, drop: float = 0.0) -> None:
        super().__init__()
        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
        self.dwconv = DWConv(hidden_features)
        self.act = nn.GELU()
        self.fc2 = nn.Conv2d(hidden_features, in_features, 1)
        self.drop = nn.Dropout(drop)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = self.dwconv(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class StripGatingUnit(nn.Module):
    """Strip spatial gating: 5x5 DWConv -> (1, k2) -> (k2, 1) -> 1x1 -> gate."""

    def __init__(self, dim: int, k1: int, k2: int) -> None:
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        # Names match upstream pretrained checkpoint: conv_spatial1 / conv_spatial2.
        self.conv_spatial1 = nn.Conv2d(dim, dim, kernel_size=(k1, k2), stride=1,
                                       padding=(k1 // 2, k2 // 2), groups=dim)
        self.conv_spatial2 = nn.Conv2d(dim, dim, kernel_size=(k2, k1), stride=1,
                                       padding=(k2 // 2, k1 // 2), groups=dim)
        self.conv1 = nn.Conv2d(dim, dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn = self.conv0(x)
        attn = self.conv_spatial1(attn)
        attn = self.conv_spatial2(attn)
        attn = self.conv1(attn)
        return x * attn


class StripAttention(nn.Module):
    def __init__(self, dim: int, k1: int, k2: int) -> None:
        super().__init__()
        self.proj_1 = nn.Conv2d(dim, dim, 1)
        self.activation = nn.GELU()
        self.spatial_gating_unit = StripGatingUnit(dim, k1, k2)
        self.proj_2 = nn.Conv2d(dim, dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.proj_1(x)
        x = self.activation(x)
        x = self.spatial_gating_unit(x)
        x = self.proj_2(x)
        return x + residual


class Block(nn.Module):
    def __init__(self, dim: int, mlp_ratio: float, k1: int, k2: int, drop: float, drop_path: float) -> None:
        super().__init__()
        self.norm1 = nn.BatchNorm2d(dim)
        self.norm2 = nn.BatchNorm2d(dim)
        self.attn = StripAttention(dim, k1, k2)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop)
        ls_init = 1e-2
        self.layer_scale_1 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
        self.layer_scale_2 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.drop_path(
            self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x))
        )
        x = x + self.drop_path(
            self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))
        )
        return x


class OverlapPatchEmbed(nn.Module):
    def __init__(self, patch_size: int, stride: int, in_chans: int, embed_dim: int) -> None:
        super().__init__()
        ph, pw = _to_2tuple(patch_size)
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(ph, pw), stride=stride,
                              padding=(ph // 2, pw // 2))
        self.norm = nn.BatchNorm2d(embed_dim)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
        x = self.proj(x)
        _, _, H, W = x.shape
        x = self.norm(x)
        return x, H, W


class StripNet(nn.Module):
    """Strip-R-CNN backbone: 4-stage hierarchical CNN with strip-shaped DWConv attention.

    Output: list of [B, C_i, H/s_i, W/s_i] per stage. Use `forward_last_features` for
    the deepest stage only.
    """

    def __init__(
        self,
        embed_dims: List[int] = [64, 128, 320, 512],
        mlp_ratios: List[int] = [8, 8, 4, 4],
        k1s: List[int] = [1, 1, 1, 1],
        k2s: List[int] = [19, 19, 19, 19],
        depths: List[int] = [2, 2, 4, 2],
        drop_rate: float = 0.1,
        drop_path_rate: float = 0.15,
        in_chans: int = 3,
    ) -> None:
        super().__init__()
        self.depths = depths
        self.num_stages = len(depths)
        self.embed_dims = embed_dims
        self.gradient_checkpointing = False

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
        cur = 0
        for i in range(self.num_stages):
            patch_embed = OverlapPatchEmbed(
                patch_size=7 if i == 0 else 3,
                stride=4 if i == 0 else 2,
                in_chans=in_chans if i == 0 else embed_dims[i - 1],
                embed_dim=embed_dims[i],
            )
            block = nn.ModuleList([
                Block(dim=embed_dims[i], mlp_ratio=mlp_ratios[i], k1=k1s[i], k2=k2s[i],
                      drop=drop_rate, drop_path=dpr[cur + j])
                for j in range(depths[i])
            ])
            norm = nn.LayerNorm(embed_dims[i], eps=1e-6)
            cur += depths[i]
            setattr(self, f"patch_embed{i + 1}", patch_embed)
            setattr(self, f"block{i + 1}", block)
            setattr(self, f"norm{i + 1}", norm)

    def set_gradient_checkpointing(self, enable: bool = True) -> None:
        """Recompute Strip blocks during backward to lower VRAM (full fine-tune)."""
        self.gradient_checkpointing = bool(enable)

    def forward_features(self, x: torch.Tensor) -> List[torch.Tensor]:
        B = x.shape[0]
        outs: List[torch.Tensor] = []
        for i in range(self.num_stages):
            patch_embed = getattr(self, f"patch_embed{i + 1}")
            block = getattr(self, f"block{i + 1}")
            norm = getattr(self, f"norm{i + 1}")
            x, H, W = patch_embed(x)
            for blk in block:
                if self.gradient_checkpointing and self.training:
                    x = torch.utils.checkpoint.checkpoint(
                        blk, x, use_reentrant=False,
                    )
                else:
                    x = blk(x)
            x = x.flatten(2).transpose(1, 2)
            x = norm(x)
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
            outs.append(x)
        return outs

    def forward_last_features(self, x: torch.Tensor) -> torch.Tensor:
        return self.forward_features(x)[-1]


def get_stripnet_small() -> StripNet:
    return StripNet(
        embed_dims=[64, 128, 320, 512],
        mlp_ratios=[8, 8, 4, 4],
        k1s=[1, 1, 1, 1],
        k2s=[19, 19, 19, 19],
        depths=[2, 2, 4, 2],
        drop_rate=0.1,
        drop_path_rate=0.15,
    )


def load_stripnet_small_pretrained(checkpoint_path: str | Path) -> StripNet:
    """Build StripNet-small and load ImageNet-pretrained weights.

    Strips the classification `head.*` keys. Tolerates missing/extra keys
    (norm{N}.* are LayerNorm here vs BatchNorm in some forks — we keep LN).
    """
    LOGGER.info("📐 Loading StripNet-small from %s", checkpoint_path)
    model = get_stripnet_small()
    raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw

    # Drop classification head + the BN-form norm{N} keys if present (we use LN here).
    drop_prefixes = ("head.",)
    cleaned = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop_prefixes)}

    # The pretrained checkpoint stores norm{N} as BatchNorm2d (running_mean/var/num_batches_tracked).
    # Our code uses LayerNorm at this position. Strip BN running stats if found; copy weight/bias.
    for n in (1, 2, 3, 4):
        for suffix in ("running_mean", "running_var", "num_batches_tracked"):
            cleaned.pop(f"norm{n}.{suffix}", None)

    missing, unexpected = model.load_state_dict(cleaned, strict=False)
    if missing:
        LOGGER.info("StripNet missing keys (expected for newly-init layers): %d", len(missing))
    if unexpected:
        LOGGER.info("StripNet unexpected keys (ignored): %d", len(unexpected))
    LOGGER.info("📐 StripNet-small loaded: %d params", sum(p.numel() for p in model.parameters()))
    return model
