from torch import nn, Tensor
import torch.nn.functional as F
import timm
from typing import Union, Optional

from ..utils import BasicBlock, Bottleneck, make_resnet_layers
from ..utils import _init_weights


model_configs = {
    "resnet18.tv_in1k": {
        "decoder_channels": [512, 256, 128],
    },
    "resnet34.tv_in1k": {
        "decoder_channels": [512, 256, 128],
    },
    "resnet50.tv_in1k": {
        "decoder_channels": [512, 256, 256, 128],
    },
    "resnet101.tv_in1k": {
        "decoder_channels": [512, 512, 256, 256, 128],
    },
    "resnet152.tv_in1k": {
        "decoder_channels": [512, 512, 512, 256, 256, 128],
    },
}


class ResNet(nn.Module):
    def __init__(
        self,
        decoder_block: Union[BasicBlock, Bottleneck],
        backbone: str = "resnet34.tv_in1k",
        reduction: Optional[int] = None,
    ) -> None:
        super().__init__()
        assert backbone in model_configs.keys(), f"Backbone should be in {model_configs.keys()}"
        config = model_configs[backbone]
        encoder = timm.create_model(backbone, pretrained=True, features_only=True, out_indices=(-1,))
        encoder_reduction = encoder.feature_info.reduction()[-1]

        if reduction <= 16:
            if "resnet18" in backbone or "resnet34" in backbone:
                encoder.layer4[0].conv1.stride = (1, 1)
                encoder.layer4[0].downsample[0].stride = (1, 1)
            else:
                encoder.layer4[0].conv2.stride = (1, 1)
                encoder.layer4[0].downsample[0].stride = (1, 1)
            encoder_reduction = encoder_reduction // 2

        self.encoder = encoder
        self.encoder_reduction = encoder_reduction

        encoder_out_channels = self.encoder.feature_info.channels()[-1]

        decoder_channels = config["decoder_channels"]
        self.decoder = make_resnet_layers(
            block=decoder_block,
            cfg=decoder_channels,
            in_channels=encoder_out_channels,
            dilation=1,
            expansion=1,
        )
        self.decoder.apply(_init_weights)

        self.reduction = self.encoder_reduction if reduction is None else reduction
        self.channels = decoder_channels[-1]

    def forward(self, x: Tensor) -> Tensor:
        x = self.encoder(x)[-1]
        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 resnet18(reduction: int = 32) -> ResNet:
    return ResNet(decoder_block=BasicBlock, backbone="resnet18.tv_in1k", reduction=reduction)


def resnet34(reduction: int = 32) -> ResNet:
    return ResNet(decoder_block=BasicBlock, backbone="resnet34.tv_in1k", reduction=reduction)


def resnet50(reduction: int = 32) -> ResNet:
    return ResNet(decoder_block=Bottleneck, backbone="resnet50.tv_in1k", reduction=reduction)


def resnet101(reduction: int = 32) -> ResNet:
    return ResNet(decoder_block=Bottleneck, backbone="resnet101.tv_in1k", reduction=reduction)


def resnet152(reduction: int = 32) -> ResNet:
    return ResNet(decoder_block=Bottleneck, backbone="resnet152.tv_in1k", reduction=reduction)