import torch
import torch.nn as nn

from src.models.sofia_v1.stripnet_blocks_dcn import DCNBlock
from src.models.sofia_v1.stripnet_blocks_dcn_new import DCNBlockV4
from src.models.stripnet.model import OverlapPatchEmbed


class StripNetDCN(nn.Module):
    """4-stage hierarchical CNN backbone: OverlapPatchEmbed + DCN blocks per stage.

    Output: list of [B, C_i, H_i, W_i] features at each stage. For 256x256 input
    with default downsampling (4, 2, 2, 2), final stage is 8x8.

    DCN variant:
      - "v2" (default): torchvision `DeformConv2d` — stable, no memory leaks.
      - "v4": OpenGVLab DCNv4 — faster on CUDA but has a known C++ extension
        memory leak (~9 MB per forward call) that causes OOM in long training
        runs. Only use if you have a patched DCNv4 build.
    """

    def __init__(
        self,
        in_chans: int = 3,
        embed_dims: list[int] = [64, 128, 256, 512],
        depths: list[int] = [3, 4, 6, 3],
        dcn_variant: str = "v2",
    ) -> None:
        super().__init__()
        if dcn_variant not in ("v2", "v4"):
            raise ValueError(f"dcn_variant must be 'v2' or 'v4', got {dcn_variant!r}")
        self.num_stages = len(embed_dims)
        self.embed_dims = embed_dims
        self.dcn_variant = dcn_variant

        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 self.embed_dims[i - 1],
                embed_dim=self.embed_dims[i],
            )

            block_cls = DCNBlockV4 if dcn_variant == "v4" else DCNBlock
            block = nn.ModuleList([
                block_cls(
                    in_ch=self.embed_dims[i],
                    out_ch=self.embed_dims[i],
                ) for _ in range(depths[i])
            ])

            setattr(self, f"patch_embed{i + 1}", patch_embed)
            setattr(self, f"block{i + 1}", block)

    def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]:
        outs = []
        for i in range(self.num_stages):
            patch_embed = getattr(self, f"patch_embed{i + 1}")
            block = getattr(self, f"block{i + 1}")
            x, H, W = patch_embed(x)
            for blk in block:
                x = blk(x)
            outs.append(x)
        return outs

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

def get_stripnet_dcn_small(dcn_variant: str = "v2") -> StripNetDCN:
    return StripNetDCN(
        in_chans=3,
        embed_dims=[64, 128, 320, 512],
        depths=[2, 2, 4, 2],
        dcn_variant=dcn_variant,
    )


def get_stripnet_dcn_small_v2(dcn_variant: str = "v2") -> StripNetDCN:
    return StripNetDCN(
        in_chans=3,
        embed_dims=[64, 128, 256, 384],
        depths=[2, 2, 4, 2],
        dcn_variant=dcn_variant,
    )


def get_stripnet_dcn_tiny(dcn_variant: str = "v2") -> StripNetDCN:
    return StripNetDCN(
        in_chans=3,
        embed_dims=[32, 64, 128, 256],
        depths=[3, 3, 5, 2],
        dcn_variant=dcn_variant,
    )


def get_stripnet_dcn_tiny_tiny(dcn_variant: str = "v2") -> StripNetDCN:
    return StripNetDCN(
        in_chans=3,
        embed_dims=[16, 32, 80, 128],
        depths=[3, 3, 5, 2],
        dcn_variant=dcn_variant,
    )


VARIANT_MAP = {
    "small": get_stripnet_dcn_small,
    "small_v2": get_stripnet_dcn_small_v2,
    "tiny": get_stripnet_dcn_tiny,
    "tiny_tiny": get_stripnet_dcn_tiny_tiny,
}


def build_stripnet_dcn(variant: str = "small", dcn_variant: str = "v2") -> StripNetDCN:
    """Factory: variant in {tiny_tiny, tiny, small, small_v2}, dcn_variant in {v2, v4}."""
    if variant not in VARIANT_MAP:
        raise ValueError(f"Unknown variant {variant!r}. Available: {list(VARIANT_MAP)}")
    return VARIANT_MAP[variant](dcn_variant=dcn_variant)