import torch
from torch import nn, Tensor
import torch.nn.functional as F

from typing import List, Optional

from ..utils import _init_weights
from .csrnet import CSRNet, csrnet, csrnet_bn

EPS = 1e-6


class ContextualModule(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int = 512,
        sizes: List[int] = [1, 2, 3, 6],
    ) -> None:
        super().__init__()
        self.scales = nn.ModuleList([self.__make_scale__(in_channels, size) for size in sizes])
        self.bottleneck = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1)
        self.relu = nn.ReLU(inplace=True)
        self.weight_net = nn.Conv2d(in_channels, in_channels, kernel_size=1)

    def __make_weight__(self, feature: Tensor, scale_feature: Tensor) -> Tensor:
        weight_feature = feature - scale_feature
        weight_feature = self.weight_net(weight_feature)
        return F.sigmoid(weight_feature)
    
    def __make_scale__(self, channels: int, size: int) -> nn.Module:
        return nn.Sequential(
            nn.AdaptiveAvgPool2d(output_size=(size, size)),
            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
        )

    def forward(self, feature: Tensor) -> Tensor:
        h, w = feature.shape[-2:]
        multi_scales = [F.interpolate(input=scale(feature), size=(h, w), mode="bilinear") for scale in self.scales]
        weights = [self.__make_weight__(feature, scale_feature) for scale_feature in multi_scales]
        multi_scales = sum([multi_scales[i] * weights[i] for i in range(len(weights))]) / (sum(weights) + EPS)
        overall_features = torch.cat([multi_scales, feature], dim=1)
        overall_features = self.bottleneck(overall_features)
        overall_features = self.relu(overall_features)
        return overall_features


class CANNet(nn.Module):
    def __init__(
        self,
        csrnet: CSRNet,
        sizes: List[int] = [1, 2, 3, 6],
        reduction: Optional[int] = 8,
    ) -> None:
        super().__init__()
        assert isinstance(csrnet, CSRNet), f"csrnet should be an instance of CSRNet, got {type(csrnet)}."
        assert isinstance(sizes, (tuple, list)), f"sizes should be a list or tuple, got {type(sizes)}."
        assert len(sizes) > 0, f"Expected at least one size, got {len(sizes)}."
        assert all([isinstance(size, int) for size in sizes]), f"Expected all size to be int, got {sizes}."
        self.sizes = sizes
        self.encoder_reduction = csrnet.encoder_reduction
        self.reduction = self.encoder_reduction if reduction is None else reduction

        self.features = csrnet.features
        self.decoder = csrnet.decoder
        self.decoder.apply(_init_weights)
        self.context = ContextualModule(512, 512, self.sizes)
        self.context.apply(_init_weights)

        self.channels = csrnet.channels

    def forward(self, x: Tensor) -> Tensor:
        x = self.features(x)
        x = self.context(x)
        if self.encoder_reduction != self.reduction:
            x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
        x = self.decoder(x)
        return x


def cannet(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
    return CANNet(csrnet(), sizes=sizes, reduction=reduction)

def cannet_bn(sizes=[1, 2, 3, 6], reduction: int = 8) -> CANNet:
    return CANNet(csrnet_bn(), sizes=sizes, reduction=reduction)