import math
from typing import Literal

import torch
import torchvision.models as tv_models
from torch import Tensor, nn
from torch.nn import functional as F
from torchvision.models.densenet import DenseNet121_Weights, DenseNet161_Weights
from torchvision.models.resnet import (
    ResNet50_Weights,
    ResNet101_Weights,
    ResNeXt50_32X4D_Weights,
    ResNeXt101_32X8D_Weights,
)

from torch_uncertainty.layers.distributions import get_dist_conv_layer
from torch_uncertainty.models.utils import Backbone

resnet_feat_out_channels = [64, 256, 512, 1024, 2048]
resnet_feat_names = ["relu", "layer1", "layer2", "layer3", "layer4"]
densenet_feat_names = [
    "relu0",
    "pool0",
    "transition1",
    "transition2",
    "norm5",
]

bts_backbones = [
    "densenet121",
    "densenet161",
    "resnet50",
    "resnet101",
    "resnext50",
    "resnext101",
]


class AtrousBlock2d(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        dilation: int,
        norm_first: bool = True,
        norm_momentum: float = 0.1,
        **factory_kwargs,
    ):
        """Atrous block with 1x1 and 3x3 convolutions.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            dilation (int): Dilation rate for the 3x3 convolution.
            norm_first (bool): Whether to apply normalization before the 1x1 convolution.
                Defaults to True.
            norm_momentum (float): Momentum for the normalization layer. Defaults to 0.1.
            factory_kwargs: Additional arguments for the PyTorch layers.
        """
        super().__init__()

        self.norm_first = norm_first
        if norm_first:
            self.first_norm = nn.BatchNorm2d(in_channels, momentum=norm_momentum, **factory_kwargs)

        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels * 2,
            bias=False,
            kernel_size=1,
            stride=1,
            padding=0,
            **factory_kwargs,
        )
        self.norm = nn.BatchNorm2d(out_channels * 2, momentum=norm_momentum, **factory_kwargs)
        self.conv2 = nn.Conv2d(
            in_channels=out_channels * 2,
            out_channels=out_channels,
            bias=False,
            kernel_size=3,
            stride=1,
            padding=(dilation, dilation),
            dilation=dilation,
            **factory_kwargs,
        )

    def forward(self, x: Tensor) -> Tensor:
        if self.norm_first:
            x = self.first_norm(x)
        out = F.relu(self.conv1(x))
        out = F.relu(self.norm(out))
        return self.conv2(out)


class UpConv2d(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        ratio: int = 2,
        **factory_kwargs,
    ):
        """Upsampling convolution.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            ratio (int): Upsampling ratio.
            factory_kwargs: Additional arguments for the convolution layer.
        """
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            bias=False,
            kernel_size=3,
            stride=1,
            padding=1,
            **factory_kwargs,
        )
        self.ratio = ratio

    def forward(self, x: Tensor) -> Tensor:
        out = F.interpolate(x, scale_factor=self.ratio, mode="nearest")
        return F.elu(self.conv(out))


class Reduction1x1(nn.Module):
    def __init__(
        self,
        num_in_filters: int,
        num_out_filters: int,
        max_depth: float,
        is_final: bool = False,
        **factory_kwargs,
    ):
        super().__init__()
        self.max_depth = max_depth
        self.is_final = is_final
        self.reduc = torch.nn.Sequential()

        while num_out_filters >= 4:
            if num_out_filters < 8:
                if self.is_final:
                    self.reduc.add_module(
                        "final",
                        torch.nn.Sequential(
                            nn.Conv2d(
                                num_in_filters,
                                out_channels=1,
                                bias=False,
                                kernel_size=1,
                                stride=1,
                                padding=0,
                                **factory_kwargs,
                            ),
                            nn.Sigmoid(),
                        ),
                    )
                else:
                    self.reduc.add_module(
                        "plane_params",
                        torch.nn.Conv2d(
                            num_in_filters,
                            out_channels=3,
                            bias=False,
                            kernel_size=1,
                            stride=1,
                            padding=0,
                            **factory_kwargs,
                        ),
                    )
                break

            self.reduc.add_module(
                f"inter_{num_in_filters}_{num_out_filters}",
                torch.nn.Sequential(
                    nn.Conv2d(
                        in_channels=num_in_filters,
                        out_channels=num_out_filters,
                        bias=False,
                        kernel_size=1,
                        stride=1,
                        padding=0,
                        **factory_kwargs,
                    ),
                    nn.ELU(),
                ),
            )

            num_in_filters = num_out_filters
            num_out_filters = num_out_filters // 2

    def forward(self, x: Tensor) -> Tensor:
        x = self.reduc.forward(x)
        if not self.is_final:
            theta = F.sigmoid(x[:, 0, :, :]) * math.pi / 3
            phi = F.sigmoid(x[:, 1, :, :]) * math.pi * 2
            dist = F.sigmoid(x[:, 2, :, :]) * self.max_depth
            x = torch.cat(
                [
                    torch.mul(torch.sin(theta), torch.cos(phi)).unsqueeze(1),
                    torch.mul(torch.sin(theta), torch.sin(phi)).unsqueeze(1),
                    torch.cos(theta).unsqueeze(1),
                    dist.unsqueeze(1),
                ],
                dim=1,
            )
        return x


class LocalPlanarGuidance(nn.Module):
    def __init__(self, up_ratio: int) -> None:
        super().__init__()
        self.register_buffer("u", torch.arange(up_ratio).reshape([1, 1, up_ratio]))
        self.register_buffer("v", torch.arange(up_ratio).reshape([1, up_ratio, 1]))
        self.up_ratio = up_ratio

    def forward(self, x: Tensor) -> Tensor:
        x_expanded = torch.repeat_interleave(
            torch.repeat_interleave(x, self.up_ratio, 2), self.up_ratio, 3
        )

        u = self.u.repeat(
            x.size(0),
            x.size(2) * self.up_ratio,
            x.size(3),
        )
        u = (u - (self.up_ratio - 1) * 0.5) / self.up_ratio

        v = self.v.repeat(
            x.size(0),
            x.size(2),
            x.size(3) * self.up_ratio,
        )
        v = (v - (self.up_ratio - 1) * 0.5) / self.up_ratio

        return x_expanded[:, 3, :, :] / (
            x_expanded[:, 0, :, :] * u + x_expanded[:, 1, :, :] * v + x_expanded[:, 2, :, :]
        )


class BTSBackbone(Backbone):
    def __init__(self, backbone_name: str, pretrained: bool) -> None:
        """BTS backbone.

        Args:
            backbone_name (str): Name of the backbone.
            pretrained (bool): Use a pretrained backbone.
        """
        if backbone_name == "densenet121":
            model = tv_models.densenet121(
                weights=DenseNet121_Weights.DEFAULT if pretrained else None
            ).features
            feat_names = densenet_feat_names
            self.feat_out_channels = [64, 64, 128, 256, 1024]
        elif backbone_name == "densenet161":
            model = tv_models.densenet161(
                weights=DenseNet161_Weights.DEFAULT if pretrained else None
            ).features
            feat_names = densenet_feat_names
            self.feat_out_channels = [96, 96, 192, 384, 2208]
        elif backbone_name == "resnet50":
            model = tv_models.resnet50(
                weights=ResNet50_Weights.IMAGENET1K_V2 if pretrained else None
            )
        elif backbone_name == "resnet101":
            model = tv_models.resnet101(
                weights=ResNet101_Weights.IMAGENET1K_V2 if pretrained else None
            )
        elif backbone_name == "resnext50":
            model = tv_models.resnext50_32x4d(
                weights=ResNeXt50_32X4D_Weights.IMAGENET1K_V2 if pretrained else None
            )
        else:  # backbone_name == "resnext101":
            model = tv_models.resnext101_32x8d(
                weights=ResNeXt101_32X8D_Weights.IMAGENET1K_V2 if pretrained else None
            )
        if "res" in backbone_name:  # remove classification heads from ResNets
            feat_names = resnet_feat_names
            self.feat_out_channels = resnet_feat_out_channels
            model.avgpool = nn.Identity()
            model.fc = nn.Identity()
        super().__init__(model=model, feat_names=feat_names)


class BTSDecoder(nn.Module):
    """BTS decoder.

    Args:
        max_depth (float): The maximum predicted depth.
        feat_out_channels (list[int]): The number of output channels from the backbone.
        num_features (int): The number of features to use in the decoder.
        dist_family (str | None, optional): The distribution family name. ``None`` means point-wise
            prediction. Defaults to ``None``.
    """

    def __init__(
        self,
        max_depth: float,
        feat_out_channels: list[int],
        num_features: int,
        dist_family: str | None = None,
    ):
        super().__init__()
        self.max_depth = max_depth

        self.upconv5 = UpConv2d(in_channels=feat_out_channels[4], out_channels=num_features)
        self.bn5 = nn.BatchNorm2d(num_features, momentum=0.01, affine=True)

        self.conv5 = nn.Conv2d(
            num_features + feat_out_channels[3],
            num_features,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False,
        )
        self.upconv4 = UpConv2d(in_channels=num_features, out_channels=num_features // 2)
        self.bn4 = nn.BatchNorm2d(num_features // 2, momentum=0.01, affine=True)
        self.conv4 = nn.Conv2d(
            num_features // 2 + feat_out_channels[2],
            num_features // 2,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False,
        )
        self.bn4_2 = nn.BatchNorm2d(num_features // 2, momentum=0.01, affine=True)

        self.daspp_3 = AtrousBlock2d(
            num_features // 2,
            num_features // 4,
            3,
            norm_first=False,
            norm_momentum=0.01,
        )
        self.daspp_6 = AtrousBlock2d(
            num_features // 2 + num_features // 4 + feat_out_channels[2],
            num_features // 4,
            6,
            norm_momentum=0.01,
        )
        self.daspp_12 = AtrousBlock2d(
            num_features + feat_out_channels[2],
            num_features // 4,
            12,
            norm_momentum=0.01,
        )
        self.daspp_18 = AtrousBlock2d(
            num_features + num_features // 4 + feat_out_channels[2],
            num_features // 4,
            18,
            norm_momentum=0.01,
        )
        self.daspp_24 = AtrousBlock2d(
            num_features + num_features // 2 + feat_out_channels[2],
            num_features // 4,
            24,
            norm_momentum=0.01,
        )
        self.daspp_conv = torch.nn.Sequential(
            nn.Conv2d(
                num_features + num_features // 2 + num_features // 4,
                num_features // 4,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            ),
            nn.ELU(),
        )
        self.reduc8x8 = Reduction1x1(num_features // 4, num_features // 4, self.max_depth)
        self.lpg8x8 = LocalPlanarGuidance(8)

        self.upconv3 = UpConv2d(in_channels=num_features // 4, out_channels=num_features // 4)
        self.bn3 = nn.BatchNorm2d(num_features // 4, momentum=0.01, affine=True)
        self.conv3 = nn.Conv2d(
            num_features // 4 + feat_out_channels[1] + 1,
            num_features // 4,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False,
        )
        self.reduc4x4 = Reduction1x1(num_features // 4, num_features // 8, self.max_depth)
        self.lpg4x4 = LocalPlanarGuidance(4)

        self.upconv2 = UpConv2d(in_channels=num_features // 4, out_channels=num_features // 8)
        self.bn2 = nn.BatchNorm2d(num_features // 8, momentum=0.01, affine=True)
        self.conv2 = nn.Conv2d(
            num_features // 8 + feat_out_channels[0] + 1,
            num_features // 8,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False,
        )

        self.reduc2x2 = Reduction1x1(num_features // 8, num_features // 16, self.max_depth)
        self.lpg2x2 = LocalPlanarGuidance(2)

        self.upconv1 = UpConv2d(in_channels=num_features // 8, out_channels=num_features // 16)
        self.reduc1x1 = Reduction1x1(
            num_features // 16,
            num_features // 32,
            self.max_depth,
            is_final=True,
        )
        self.conv1 = nn.Conv2d(num_features // 16 + 4, num_features // 16, 3, 1, 1, bias=False)
        self.output_channels = 1

        if dist_family is not None:
            dist_layer_class = get_dist_conv_layer(dist_family)
            self.depth = dist_layer_class(
                base_layer=nn.Conv2d,
                event_dim=self.output_channels,
                in_channels=num_features // 16,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            )
        else:
            self.depth = nn.Conv2d(num_features // 16, self.output_channels, 3, 1, 1, bias=False)

    def feat_forward(self, features: list[Tensor]) -> Tensor:
        dense_features = F.relu(features[4])
        upconv5 = self.bn5(self.upconv5(dense_features))  # H/16
        iconv5 = F.elu(self.conv5(torch.cat([upconv5, features[3]], dim=1)))

        upconv4 = self.bn4(self.upconv4(iconv5))  # H/8
        concat4 = torch.cat([upconv4, features[2]], dim=1)
        iconv4 = self.bn4_2(F.elu(self.conv4(concat4)))

        daspp_3 = self.daspp_3(iconv4)
        concat4_2 = torch.cat([concat4, daspp_3], dim=1)
        daspp_6 = self.daspp_6(concat4_2)
        concat4_3 = torch.cat([concat4_2, daspp_6], dim=1)
        daspp_12 = self.daspp_12(concat4_3)
        concat4_4 = torch.cat([concat4_3, daspp_12], dim=1)
        daspp_18 = self.daspp_18(concat4_4)
        daspp_24 = self.daspp_24(torch.cat([concat4_4, daspp_18], dim=1))
        concat4_daspp = torch.cat([iconv4, daspp_3, daspp_6, daspp_12, daspp_18, daspp_24], dim=1)
        daspp_feat = self.daspp_conv(concat4_daspp)

        reduc8x8 = self.reduc8x8(daspp_feat)
        plane_normal_8x8 = reduc8x8[:, :3, :, :]
        plane_normal_8x8 = F.normalize(plane_normal_8x8, p=2, dim=1)
        plane_dist_8x8 = reduc8x8[:, 3, :, :]
        plane_eq_8x8 = torch.cat([plane_normal_8x8, plane_dist_8x8.unsqueeze(1)], 1)
        depth_8x8 = self.lpg8x8(plane_eq_8x8)
        depth_8x8_scaled = depth_8x8.unsqueeze(1) / self.max_depth
        depth_8x8_scaled_ds = F.interpolate(depth_8x8_scaled, scale_factor=0.25, mode="nearest")

        upconv3 = self.bn3(self.upconv3(daspp_feat))  # H/4
        concat3 = torch.cat([upconv3, features[1], depth_8x8_scaled_ds], dim=1)
        iconv3 = F.elu(self.conv3(concat3))

        reduc4x4 = self.reduc4x4(iconv3)
        plane_normal_4x4 = reduc4x4[:, :3, :, :]
        plane_normal_4x4 = F.normalize(plane_normal_4x4, p=2, dim=1)
        plane_dist_4x4 = reduc4x4[:, 3, :, :]
        plane_eq_4x4 = torch.cat([plane_normal_4x4, plane_dist_4x4.unsqueeze(1)], 1)
        depth_4x4 = self.lpg4x4(plane_eq_4x4)
        depth_4x4_scaled = depth_4x4.unsqueeze(1) / self.max_depth
        depth_4x4_scaled_ds = F.interpolate(depth_4x4_scaled, scale_factor=0.5, mode="nearest")

        upconv2 = self.bn2(self.upconv2(iconv3))  # H/2
        iconv2 = F.elu(self.conv2(torch.cat([upconv2, features[0], depth_4x4_scaled_ds], dim=1)))

        reduc2x2 = self.reduc2x2(iconv2)
        plane_normal_2x2 = reduc2x2[:, :3, :, :]
        plane_normal_2x2 = F.normalize(plane_normal_2x2, p=2, dim=1)
        plane_dist_2x2 = reduc2x2[:, 3, :, :]
        plane_eq_2x2 = torch.cat([plane_normal_2x2, plane_dist_2x2.unsqueeze(1)], 1)
        depth_2x2 = self.lpg2x2(plane_eq_2x2)
        depth_2x2_scaled = depth_2x2.unsqueeze(1) / self.max_depth

        upconv1 = self.upconv1(iconv2)
        reduc1x1 = self.reduc1x1(upconv1)
        concat1 = torch.cat(
            [
                upconv1,
                reduc1x1,
                depth_2x2_scaled,
                depth_4x4_scaled,
                depth_8x8_scaled,
            ],
            dim=1,
        )
        return F.elu(self.conv1(concat1))

    def forward(self, features: list[Tensor]) -> Tensor | dict[str, Tensor]:
        """Forward pass.

        Args:
            features (list[Tensor]): list of the features from the backbone.

        Note:
            Depending of the :attr:`dist_family` of the backbone, the output can
            be a dictionnary of distribution parameters or a single tensor.
        """
        # TODO: handle focal
        out = self.depth(self.feat_forward(features))
        if isinstance(out, dict):
            if "loc" not in out or "scale" not in out:
                raise ValueError(
                    "Expected 'loc' and 'scale' in the output dictionary.",
                    "Consider raising an issue on the repository if you need support ",
                    "for distributions that do not use location and scale.",
                )
            out["loc"] = self.max_depth * F.sigmoid(out["loc"])
            out["scale"] = self.max_depth * out["scale"]
        else:
            out = self.max_depth * F.sigmoid(out)
        return out


class _BTS(nn.Module):
    def __init__(
        self,
        backbone_name: Literal[
            "densenet121",
            "densenet161",
            "resnet50",
            "resnet101",
            "resnext50",
            "resnext101",
        ],
        max_depth: float,
        bts_size: int = 512,
        dist_family: str | None = None,
        pretrained_backbone: bool = True,
    ) -> None:
        """BTS model.

        Args:
            backbone_name (str): Name of the encoding backbone.
            max_depth (float): Maximum predicted depth.
            bts_size (int): BTS feature size. Defaults to 512.
            dist_family (str): Distribution family name. Defaults to None.
            estimation. Defaults to nn.Identity.
            pretrained_backbone (bool): Use a pretrained backbone. Defaults to True.

        Reference:
            From Big to Small: Multi-Scale Local Planar Guidance for Monocular Depth Estimation.
            Jin Han Lee, Myung-Kyu Han, Dong Wook Ko, Il Hong Suh. ArXiv.
        """
        super().__init__()
        self.max_depth = max_depth

        self.backbone = BTSBackbone(backbone_name, pretrained_backbone)
        self.decoder = BTSDecoder(max_depth, self.backbone.feat_out_channels, bts_size, dist_family)

    # TODO: Handle focal
    def forward(self, x: Tensor, focal: float | None = None) -> Tensor:
        """Forward pass.

        Args:
            x (Tensor): Input tensor.
            focal (float): Focal length for API consistency.
        """
        return self.decoder(self.backbone(x))


def _bts(
    backbone_name: str,
    max_depth: float,
    bts_size: int = 512,
    dist_family: str | None = None,
    pretrained_backbone: bool = True,
) -> _BTS:
    if backbone_name not in bts_backbones:
        raise ValueError(f"Unsupported backbone. Got {backbone_name}.")
    return _BTS(backbone_name, max_depth, bts_size, dist_family, pretrained_backbone)


def bts_resnet50(
    max_depth: float,
    bts_size: int = 512,
    dist_family: str | None = None,
    pretrained_backbone: bool = True,
) -> _BTS:
    """BTS model with ResNet-50 backbone.

    Args:
        max_depth (float): Maximum predicted depth.
        bts_size (int): BTS feature size. Defaults to 512.
        dist_family (str): Distribution family name. Defaults to None.
        pretrained_backbone (bool): Use a pretrained backbone. Defaults to True.
    """
    return _bts(
        "resnet50",
        max_depth,
        bts_size=bts_size,
        dist_family=dist_family,
        pretrained_backbone=pretrained_backbone,
    )


def bts_resnet101(
    max_depth: float,
    bts_size: int = 512,
    dist_family: str | None = None,
    pretrained_backbone: bool = True,
) -> _BTS:
    """BTS model with ResNet-101 backbone.

    Args:
        max_depth (float): Maximum predicted depth.
        bts_size (int): BTS feature size. Defaults to 512.
        dist_family (str): Distribution family name. Defaults to None.
        pretrained_backbone (bool): Use a pretrained backbone. Defaults to True.
    """
    return _bts(
        "resnet101",
        max_depth,
        bts_size=bts_size,
        dist_family=dist_family,
        pretrained_backbone=pretrained_backbone,
    )
