import torch
import torch.nn as nn
import torch.nn.functional as F
from ..utils.misc import NestedTensor, nested_tensor_from_tensor_list
import torchvision.transforms as transforms
from .backbone import build_backbone
from .deformable_transformer import build_deforamble_transformer


class BasicLayer(nn.Module):
    """
    Basic Convolutional Layer: Conv2d -> BatchNorm -> ReLU
    """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, bias=False):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, padding=padding, stride=stride, dilation=dilation, bias=bias
            ),
            nn.BatchNorm2d(out_channels, affine=False),
            nn.ReLU(inplace=False),
        )

    def forward(self, x):
        return self.layer(x)


class RDD_Descriptor(nn.Module):
    def __init__(self, backbone, transformer, num_feature_levels):
        super().__init__()
        self.transformer = transformer
        self.hidden_dim = transformer.d_model
        self.num_feature_levels = num_feature_levels

        self.matchibility_head = nn.Sequential(
            BasicLayer(256, 128, 1, padding=0), BasicLayer(128, 64, 1, padding=0), nn.Conv2d(64, 1, 1), nn.Sigmoid()
        )

        if num_feature_levels > 1:
            num_backbone_outs = len(backbone.strides)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.num_channels[_]
                input_proj_list.append(
                    nn.Sequential(
                        nn.Conv2d(in_channels, self.hidden_dim, kernel_size=1),
                        nn.GroupNorm(32, self.hidden_dim),
                    )
                )
            for _ in range(num_feature_levels - num_backbone_outs):
                input_proj_list.append(
                    nn.Sequential(
                        nn.Conv2d(in_channels, self.hidden_dim, kernel_size=3, stride=2, padding=1),
                        nn.GroupNorm(32, self.hidden_dim),
                    )
                )
                in_channels = self.hidden_dim
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            self.input_proj = nn.ModuleList(
                [
                    nn.Sequential(
                        nn.Conv2d(backbone.num_channels[0], self.hidden_dim, kernel_size=1),
                        nn.GroupNorm(32, self.hidden_dim),
                    )
                ]
            )
        self.backbone = backbone
        self.stride = backbone.strides[0]
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)

    def forward(self, samples: NestedTensor):

        if not isinstance(samples, NestedTensor):
            samples = nested_tensor_from_tensor_list(samples)

        features, pos = self.backbone(samples)

        srcs = []
        masks = []
        for l, feat in enumerate(features):
            src, mask = feat.decompose()
            srcs.append(self.input_proj[l](src))
            masks.append(mask)
            assert mask is not None
        if self.num_feature_levels > len(srcs):
            _len_srcs = len(srcs)
            for l in range(_len_srcs, self.num_feature_levels):
                if l == _len_srcs:
                    src = self.input_proj[l](features[-1].tensors)
                else:
                    src = self.input_proj[l](srcs[-1])
                m = samples.mask
                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
                srcs.append(src)
                masks.append(mask)
                pos.append(pos_l)

        flatten_feats, spatial_shapes, level_start_index = self.transformer(srcs, masks, pos)
        # Reshape the flattened features back to the original spatial shapes
        feats = []
        level_start_index = torch.cat(
            (level_start_index, torch.tensor([flatten_feats.shape[1] + 1]).to(level_start_index.device))
        )
        for i, shape in enumerate(spatial_shapes):
            assert len(shape) == 2
            temp = flatten_feats[:, level_start_index[i] : level_start_index[i + 1], :]
            feats.append(temp.transpose(1, 2).view(-1, self.hidden_dim, *shape))

        # Sum up the features from different levels
        final_feature = feats[0]
        for feat in feats[1:]:
            final_feature = final_feature + F.interpolate(
                feat, size=final_feature.shape[-2:], mode="bilinear", align_corners=True
            )

        matchibility = self.matchibility_head(final_feature)

        return final_feature, matchibility


def build_descriptor(config):
    backbone = build_backbone(config)
    transformer = build_deforamble_transformer(config)
    return RDD_Descriptor(backbone, transformer, config["num_feature_levels"])
