import torch
import torch.nn as nn
from DCNv4 import DCNv4  # ← your installed OpenGVLab version
import math

class DCNBlockV4(nn.Module):
    """
    StripNet-style block that uses OpenGVLab DCNv4 instead of torchvision DCNv2.
    """

    def __init__(self, in_ch, out_ch, hidden_ratio=0.25,
                 kernel_size=3, stride=1, dilation=1, group=4,
                 offset_scale=1.0, use_bias=False):
        super().__init__()
        assert kernel_size in (3, 5, 7)
        pad = (kernel_size // 2) * dilation

        # Hidden channels — must satisfy (hidden_ch // group) % 16 == 0
        hidden_ch = max(16, int(out_ch * hidden_ratio))
        hidden_ch = math.ceil(hidden_ch / 16) * 16
        # increase until kernel constraint satisfied
        while (hidden_ch // group) % 16 != 0:
            hidden_ch += 16

        #print(f"[DCNv4] adjusted hidden_ch={hidden_ch}, group={group}")


        self.reduce = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, 1, bias=False),
            self.make_gn(hidden_ch),
            nn.ReLU(inplace=True),
        )

        # DCNv4 core
        self.dcn = DCNv4(
            channels=hidden_ch,
            kernel_size=kernel_size,
            stride=stride,
            pad=pad,
            dilation=dilation,
            group=group,
            offset_scale=offset_scale,
            dw_kernel_size=None,
            center_feature_scale=False,
            remove_center=False,
            output_bias=True,
            without_pointwise=False,
        )


        self.expand = nn.Sequential(
            nn.Conv2d(hidden_ch, out_ch, 1, bias=False),
            self.make_gn(out_ch),

        )

        self.residual = (in_ch == out_ch and stride == 1)
        self.stride = stride

    def make_gn(self, num_channels):
        num_groups = max(1, num_channels // 16)
        return nn.GroupNorm(num_groups, num_channels)

    def forward(self, x):
        """
        Input:  [B, C, H, W]
        Output: [B, C, H', W']
        """
        identity = x

        x = self.reduce(x)
        B, C, H, W = x.shape

        # Flatten for DCNv4
        x_seq = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)

        # --- clamp activations to avoid huge values ---
        x_seq = torch.clamp(x_seq, -50.0, 50.0)

        x_seq = self.dcn(x_seq)
        x_seq = torch.nan_to_num(x_seq, nan=0.0, posinf=1e4, neginf=-1e4)

        # --- back to (B, C, H, W) ---
        x = x_seq.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()

        # --- expand + residual ---
        x = self.expand(x)
        if self.residual:
            x = x + identity
        return x


def test_dcnblock_v4():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    block = DCNBlockV4(
        in_ch=128,  # must match input
        out_ch=128,
        hidden_ratio=0.25,
        kernel_size=3,
        stride=1,
        dilation=1,
        group=5,
        offset_scale=1.0,
    ).to(device)

    # ✅ internal feature map (not RGB)
    x = torch.randn(2, 128, 128, 128, device=device, requires_grad=True)

    with torch.cuda.amp.autocast_mode.autocast(enabled=torch.cuda.is_available()):
        y = block(x)

    print(f"Input  shape : {x.shape}")
    print(f"Output shape : {y.shape}")
    assert y.shape == x.shape

    loss = y.mean()
    loss.backward()
    print("✅ DCNBlockV4 test passed.\n")



if __name__ == "__main__":
    test_dcnblock_v4()
