import torch
import torch.nn as nn
from torchvision.ops import DeformConv2d

class DCNBlock(nn.Module):
    """
    StripNet-style block but uses deformable conv instead of rigid convs.
    """
    def __init__(self, in_ch, out_ch, hidden_ratio=0.25, modulation=True):
        super().__init__()
        hidden_ch = max(1, int(out_ch * hidden_ratio))

        # 1x1 reduce
        self.reduce = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(hidden_ch),
            nn.ReLU(inplace=True),
        )

        # Offset conv (predicts offsets and optional mask)
        offset_channels = 2 * 3 * 3 if not modulation else 3 * 3 * 3
        self.offset_conv = nn.Conv2d(hidden_ch, offset_channels,
                                     kernel_size=3, padding=1)

        # Deformable conv
        self.deform = DeformConv2d(hidden_ch, hidden_ch,
                                   kernel_size=3, padding=1, bias=False)
        self.bn = nn.BatchNorm2d(hidden_ch)
        self.act = nn.ReLU(inplace=True)

        # 1x1 expand
        self.expand = nn.Sequential(
            nn.Conv2d(hidden_ch, out_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_ch),
        )

        self.residual = (in_ch == out_ch)

    def forward(self, x):
        identity = x

        x = self.reduce(x)
        offset = self.offset_conv(x)

        if offset.shape[1] == 18:  # DCNv1
            x = self.deform(x, offset)
        else:  # DCNv2: last 9 channels are mask
            o, mask = offset.split([18, 9], dim=1)
            mask = mask.sigmoid()
            x = self.deform(x, o, mask)

        x = self.act(self.bn(x))
        x = self.expand(x)

        if self.residual:
            x = x + identity
        return x
