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

from model.clip import build_model
from .layers import TransformerDecoder, CoordConv, conv_layer


# =========================
#  Language-conditioned FiLM
# =========================
class FeatureWiseAffine(nn.Module):
    def __init__(self, in_channels, out_channels, use_affine_level=True):
        super().__init__()
        self.use_affine_level = use_affine_level
        self.MLP = nn.Sequential(
            nn.Linear(in_channels, in_channels * 2),
            nn.LeakyReLU(),
            nn.Linear(in_channels * 2, out_channels * (1 + self.use_affine_level))
        )

    def forward(self, x, text_embed):
        # text_embed: (B, D)
        batch = x.shape[0]
        text_embed = text_embed.unsqueeze(1)  # (B, 1, D)
        if self.use_affine_level:
            gamma, beta = self.MLP(text_embed).view(batch, -1, 1, 1).chunk(2, dim=1)
            x = (1 + gamma) * x + beta
        return x


class ResidualBlock(nn.Module):
    """A residual block (note: original code uses BN(in_channels) - keep unchanged)"""

    def __init__(self, in_channels, out_channels, kernel_size=3):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding=1)
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size, padding=1)
        self.bn2 = nn.BatchNorm2d(in_channels)

    def forward(self, x_in):
        x = self.bn1(self.conv1(x_in))
        x = F.relu(x)
        x = self.bn2(self.conv2(x))
        return x + x_in


# =========================
#  RFB blocks
# =========================
class BasicConv(nn.Module):
    def __init__(
        self, in_planes, out_planes, kernel_size, stride=1, padding=0,
        dilation=1, groups=1, relu=True, bn=True, bias=False
    ):
        super().__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(
            in_planes, out_planes,
            kernel_size=kernel_size, stride=stride, padding=padding,
            dilation=dilation, groups=groups, bias=bias
        )
        self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU(inplace=True) if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class RFB(nn.Module):
    def __init__(self, in_planes, out_planes, stride=1, scale=0.1, map_reduce=8):
        super().__init__()
        self.scale = scale
        self.out_channels = out_planes
        inter_planes = in_planes // map_reduce

        self.branch0 = nn.Sequential(
            BasicConv(in_planes, 2 * inter_planes, kernel_size=1, stride=stride),
            BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=1, relu=False)
        )
        self.branch1 = nn.Sequential(
            BasicConv(in_planes, inter_planes, kernel_size=1, stride=1),
            BasicConv(inter_planes, 2 * inter_planes, kernel_size=3, stride=stride, padding=1),
            BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=3, dilation=3, relu=False)
        )
        self.branch2 = nn.Sequential(
            BasicConv(in_planes, inter_planes, kernel_size=1, stride=1),
            BasicConv(inter_planes, (inter_planes // 2) * 3, kernel_size=3, stride=1, padding=1),
            BasicConv((inter_planes // 2) * 3, 2 * inter_planes, kernel_size=3, stride=stride, padding=1),
            BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=5, dilation=5, relu=False)
        )
        self.branch3 = nn.Sequential(
            BasicConv(in_planes, inter_planes, kernel_size=1, stride=1),
            BasicConv(inter_planes, (inter_planes // 2) * 3, kernel_size=(1, 7), stride=1, padding=(0, 3)),
            BasicConv((inter_planes // 2) * 3, 2 * inter_planes, kernel_size=(7, 1), stride=stride, padding=(3, 0)),
            BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=7, dilation=7, relu=False)
        )

        self.ConvLinear = BasicConv(8 * inter_planes, out_planes, kernel_size=1, stride=1, relu=False)
        self.shortcut = BasicConv(in_planes, out_planes, kernel_size=1, stride=stride, relu=False)
        self.relu = nn.ReLU(inplace=False)

    def forward(self, x):
        x0 = self.branch0(x)
        x1 = self.branch1(x)
        x2 = self.branch2(x)
        x3 = self.branch3(x)

        out = torch.cat((x0, x1, x2, x3), 1)
        out = self.ConvLinear(out)
        short = self.shortcut(x)
        out = out * self.scale + short
        out = self.relu(out)
        return out


# =========================
#  MDAFN with FiLM
# =========================
class MDAFN(nn.Module):
    def __init__(self, in_channels, text_dim, reduction=16):
        super().__init__()
        self.pixel_conv = nn.Conv2d(in_channels, 1, kernel_size=3, padding=1)
        self.film = FeatureWiseAffine(text_dim, in_channels)

        self.channel_gap = nn.AdaptiveAvgPool2d(1)
        self.channel_fc1 = nn.Linear(in_channels, in_channels // reduction)
        self.channel_fc2 = nn.Linear(in_channels // reduction, in_channels)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, deep_features, shallow_features, text_embed):
        # concat features
        x = torch.cat([deep_features, shallow_features], dim=1)
        # FiLM
        x = self.film(x, text_embed)

        # Pixel attention
        pixel_attn_map = torch.sigmoid(self.pixel_conv(x))

        # Channel attention
        b, c, _, _ = x.size()
        y = self.channel_gap(x).view(b, c)
        y = self.channel_fc1(y)
        y = self.relu(y)
        y = self.channel_fc2(y)
        channel_attn_map = self.sigmoid(y).view(b, c, 1, 1)

        out = x * pixel_attn_map * channel_attn_map
        return out


class PixelShuffleUpsampling(nn.Module):
    def __init__(self, in_channels, out_channels, scale=2, padding=1):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels * (scale ** 2), kernel_size=3, padding=padding)
        self.pixel_shuffle = nn.PixelShuffle(scale)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.pixel_shuffle(x)
        x = self.relu(x)
        return x


# =========================
#  CORE (no instance mask)
# =========================
class CORE(nn.Module):
    """
    Outputs (no instance mask):
      - grasp_qua_pred, grasp_sin_pred, grasp_cos_pred, grasp_wid_pred
    """

    def __init__(self, cfg):
        super().__init__()
        self.use_pretrained_clip = cfg.use_pretrained_clip

        # Downsampling blocks
        self.conv1 = nn.Sequential(
            nn.Conv2d(cfg.input_channels, cfg.channel_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size),
            nn.ReLU()
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(cfg.channel_size, cfg.channel_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size),
            nn.ReLU()
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(cfg.channel_size, cfg.channel_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size),
            nn.ReLU()
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(cfg.channel_size, cfg.channel_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size),
            nn.ReLU()
        )
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv5 = nn.Sequential(
            nn.Conv2d(cfg.channel_size, cfg.channel_size * 2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 2),
            nn.ReLU()
        )
        self.conv6 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 2, cfg.channel_size * 2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 2),
            nn.ReLU()
        )
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv7 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 2, cfg.channel_size * 4, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 4),
            nn.ReLU()
        )
        self.conv8 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 4, cfg.channel_size * 4, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 4),
            nn.ReLU()
        )
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv9 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 4, cfg.channel_size * 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 8),
            nn.ReLU()
        )
        self.conv10 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 8, cfg.channel_size * 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 8),
            nn.ReLU()
        )
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv11 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 8, cfg.channel_size * 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 16),
            nn.ReLU()
        )
        self.conv12 = nn.Sequential(
            nn.Conv2d(cfg.channel_size * 16, cfg.channel_size * 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(cfg.channel_size * 16),
            nn.ReLU()
        )

        # Residual blocks
        self.res1 = ResidualBlock(cfg.channel_size * 16, cfg.channel_size * 16)
        self.res2 = ResidualBlock(cfg.channel_size * 16, cfg.channel_size * 16)
        self.res3 = ResidualBlock(cfg.channel_size * 16, cfg.channel_size * 16)

        # RFB
        self.rfb = RFB(cfg.channel_size * 16, cfg.channel_size * 16)

        # MDAFN + Upsampling
        self.mdafn1 = MDAFN(in_channels=(cfg.channel_size * 16 + cfg.channel_size * 16), text_dim=cfg.word_dim)
        self.upsample1 = PixelShuffleUpsampling(in_channels=cfg.channel_size * 32, out_channels=cfg.channel_size * 8, scale=2)

        self.mdafn2 = MDAFN(in_channels=(cfg.channel_size * 8 + cfg.channel_size * 8), text_dim=cfg.word_dim)
        self.upsample2 = PixelShuffleUpsampling(in_channels=cfg.channel_size * 16, out_channels=cfg.channel_size * 4, scale=2)

        self.mdafn3 = MDAFN(in_channels=(cfg.channel_size * 4 + cfg.channel_size * 4), text_dim=cfg.word_dim)
        self.upsample3 = PixelShuffleUpsampling(in_channels=cfg.channel_size * 8, out_channels=cfg.channel_size * 2, scale=2)

        self.mdafn4 = MDAFN(in_channels=(cfg.channel_size * 2 + cfg.channel_size * 2), text_dim=cfg.word_dim)
        self.upsample4 = PixelShuffleUpsampling(in_channels=cfg.channel_size * 4, out_channels=cfg.channel_size, scale=2)

        # CoordConv head
        self.coordconv = nn.Sequential(
            CoordConv(cfg.channel_size, cfg.channel_size, 3, 1),
            conv_layer(cfg.channel_size, cfg.channel_size, 3, 1)
        )

        # ====== OUTPUT HEADS (NO INSTANCE MASK) ======
        # from *5 -> *4 (qua, sin, cos, wid)
        self.vis_preproj = nn.Conv2d(cfg.channel_size, cfg.channel_size * 4, 1)

        # dynamic conv params from text
        self.text_preproj = nn.Linear(cfg.word_dim, 1 * cfg.channel_size * 3 * 3 + 1)

        # global FiLM on bottleneck features
        self.film = FeatureWiseAffine(cfg.word_dim, cfg.channel_size * 16)

        # init
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                nn.init.xavier_uniform_(m.weight, gain=1)

        # CLIP text encoder (JIT)
        clip_model = torch.jit.load(cfg.clip_pretrain, map_location="cpu").eval()
        self.backbone = build_model(clip_model.state_dict(), cfg.word_len, self.use_pretrained_clip).float()

        # Decoder
        self.decoder = TransformerDecoder(
            num_layers=cfg.num_layers,
            d_model=cfg.vis_dim,
            nhead=cfg.num_head,
            dim_ffn=cfg.dim_ffn,
            dropout=cfg.dropout,
            return_intermediate=cfg.intermediate
        )

    def forward(
        self,
        img,
        word,
        grasp_qua_mask=None,
        grasp_sin_mask=None,
        grasp_cos_mask=None,
        grasp_wid_mask=None
    ):
        """
        img: (B, C, H, W)
        word: token ids (B, L)
        grasp_*_mask: (B, 1, H, W) during training
        """

        # Downsampling block 1
        x = self.conv1(img)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.pool1(x)

        # Downsampling block 2
        x = self.conv5(x)
        shallow_feature_1 = self.conv6(x)
        x = self.pool2(shallow_feature_1)

        # Processing block
        x = self.conv7(x)
        shallow_feature_2 = self.conv8(x)
        x = self.pool3(shallow_feature_2)

        # More downsampling
        x = self.conv9(x)
        shallow_feature_3 = self.conv10(x)
        x = self.pool4(shallow_feature_3)

        x = self.conv11(x)
        shallow_feature_4 = self.conv12(x)

        # Bottleneck
        x = self.res1(shallow_feature_4)
        x = self.res2(x)
        x = self.res3(x)

        # RFB
        x = self.rfb(x)

        # padding mask for decoder
        pad_mask = torch.zeros_like(word).masked_fill_(word == 0, 1).bool()

        # CLIP text encode
        word_feat, state = self.backbone.encode_text(word)  # word_feat: (B,L,D), state:(B,D)

        # Global FiLM on visual bottleneck
        x = self.film(x, state)

        # Decoder fusion
        b, c, h, w = x.size()
        fq = self.decoder(x, word_feat, pad_mask)
        fq = fq.reshape(b, c, h, w)

        # MDAFN + upsample path (language-conditioned)
        f1 = self.mdafn1(fq, shallow_feature_4, state)
        f1_up = self.upsample1(f1)

        f2 = self.mdafn2(f1_up, shallow_feature_3, state)
        f2_up = self.upsample2(f2)

        f3 = self.mdafn3(f2_up, shallow_feature_2, state)
        f3_up = self.upsample3(f3)

        f4 = self.mdafn4(f3_up, shallow_feature_1, state)
        f4_up = self.upsample4(f4)

        feat_map = self.coordconv(f4_up)

        # ====== 4-head output ======
        out = self.vis_preproj(feat_map)                 # (B, 4*C, H, W)
        out = torch.tensor_split(out, 4, dim=1)          # 4 tensors of (B, C, H, W)

        grasp_qua_x = out[0]
        grasp_sin_x = out[1]
        grasp_cos_x = out[2]
        grasp_wid_x = out[3]

        B, Cc, Hh, Ww = grasp_qua_x.size()

        # group conv trick (same as your original)
        grasp_qua_x = grasp_qua_x.reshape(1, B * Cc, Hh, Ww)
        grasp_sin_x = grasp_sin_x.reshape(1, B * Cc, Hh, Ww)
        grasp_cos_x = grasp_cos_x.reshape(1, B * Cc, Hh, Ww)
        grasp_wid_x = grasp_wid_x.reshape(1, B * Cc, Hh, Ww)

        state_trans = self.text_preproj(state)  # (B, C*3*3 + 1)
        weight, bias = state_trans[:, :-1], state_trans[:, -1]
        weight = weight.reshape(B, Cc, 3, 3)

        grasp_qua_out = F.conv2d(grasp_qua_x, weight, padding=1, groups=weight.size(0), bias=bias)
        grasp_sin_out = F.conv2d(grasp_sin_x, weight, padding=1, groups=weight.size(0), bias=bias)
        grasp_cos_out = F.conv2d(grasp_cos_x, weight, padding=1, groups=weight.size(0), bias=bias)
        grasp_wid_out = F.conv2d(grasp_wid_x, weight, padding=1, groups=weight.size(0), bias=bias)

        grasp_qua_pred = grasp_qua_out.transpose(0, 1)
        grasp_sin_pred = grasp_sin_out.transpose(0, 1)
        grasp_cos_pred = grasp_cos_out.transpose(0, 1)
        grasp_wid_pred = grasp_wid_out.transpose(0, 1)

        if self.training:
            # resize targets to pred resolution
            if grasp_qua_pred.shape[-2:] != grasp_qua_mask.shape[-2:]:
                grasp_qua_mask = F.interpolate(grasp_qua_mask, grasp_qua_pred.shape[-2:], mode='nearest').detach()
                grasp_sin_mask = F.interpolate(grasp_sin_mask, grasp_sin_pred.shape[-2:], mode='nearest').detach()
                grasp_cos_mask = F.interpolate(grasp_cos_mask, grasp_cos_pred.shape[-2:], mode='nearest').detach()
                grasp_wid_mask = F.interpolate(grasp_wid_mask, grasp_wid_pred.shape[-2:], mode='nearest').detach()

            # grasp losses only
            grasp_qua_loss = F.smooth_l1_loss(grasp_qua_pred, grasp_qua_mask)
            grasp_sin_loss = F.smooth_l1_loss(grasp_sin_pred, grasp_sin_mask)
            grasp_cos_loss = F.smooth_l1_loss(grasp_cos_pred, grasp_cos_mask)
            grasp_wid_loss = F.smooth_l1_loss(grasp_wid_pred, grasp_wid_mask)

            total_loss = grasp_qua_loss + grasp_sin_loss + grasp_cos_loss + grasp_wid_loss

            loss_dict = {
                "m_qua": grasp_qua_loss.item(),
                "m_sin": grasp_sin_loss.item(),
                "m_cos": grasp_cos_loss.item(),
                "m_wid": grasp_wid_loss.item(),
            }

            preds = (
                grasp_qua_pred.detach(),
                grasp_sin_pred.detach(),
                grasp_cos_pred.detach(),
                grasp_wid_pred.detach(),
            )
            targets = (
                grasp_qua_mask,
                grasp_sin_mask,
                grasp_cos_mask,
                grasp_wid_mask,
            )
            return preds, targets, total_loss, loss_dict

        else:
            preds = (
                grasp_qua_pred.detach(),
                grasp_sin_pred.detach(),
                grasp_cos_pred.detach(),
                grasp_wid_pred.detach(),
            )
            targets = (
                grasp_qua_mask,
                grasp_sin_mask,
                grasp_cos_mask,
                grasp_wid_mask,
            )
            return preds, targets