import collections.abc
import math
import numbers
import warnings
from itertools import repeat

import torch
from einops import rearrange
from torch import nn as nn
from torch.nn import functional as F
from torch.nn import init as init
from torch.nn.modules.batchnorm import _BatchNorm


@torch.no_grad()  # @ is a decorator, which takes the decorated function as the argument of the function after @
def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
    """Initialize network weights.

    Args:
        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
        scale (float): Scale initialized weights, especially for residual
            blocks. Default: 1.
        bias_fill (float): The value to fill bias. Default: 0
        kwargs (dict): Other arguments for initialization function.
    """
    if not isinstance(module_list, list):
        module_list = [module_list]
    for module in module_list:
        for m in module.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, _BatchNorm):
                init.constant_(m.weight, 1)
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)


def make_layer(basic_block, num_basic_block, **kwarg):
    """Make layers by stacking the same blocks.

    Args:
        basic_block (nn.module): nn.module class for basic block.
        num_basic_block (int): number of blocks.

    Returns:
        nn.Sequential: Stacked blocks in nn.Sequential.
    """
    layers = []
    for _ in range(num_basic_block):
        layers.append(basic_block(**kwarg))
    return nn.Sequential(*layers)


class ResidualBlockNoBN(nn.Module):
    """Residual block without BN.

    It has a style of:
        ---Conv-ReLU-Conv-+-
         |________________|

    Args:
        num_feat (int): Channel number of intermediate features.
            Default: 64.
        res_scale (float): Residual scale. Default: 1.
        pytorch_init (bool): If set to True, use pytorch default init,
            otherwise, use default_init_weights. Default: False.
    """

    def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
        super(ResidualBlockNoBN, self).__init__()
        self.res_scale = res_scale
        self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
        self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
        self.relu = nn.ReLU(inplace=True)

        if not pytorch_init:
            default_init_weights([self.conv1, self.conv2], 0.1)

    def forward(self, x):
        identity = x
        out = self.conv2(self.relu(self.conv1(x)))
        return identity + out * self.res_scale


class Upsample(nn.Sequential):
    """Upsample module.

    Args:
        scale (int): Scale factor. Supported scales: 2^n and 3.
        num_feat (int): Channel number of intermediate features.
    """

    def __init__(self, scale, num_feat):
        m = []
        if (scale & (scale - 1)) == 0:  # scale = 2^n
            for _ in range(int(math.log(scale, 2))):
                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
                m.append(nn.PixelShuffle(2))
        elif scale == 3:
            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
            m.append(nn.PixelShuffle(3))
        else:
            raise ValueError(
                f"scale {scale} is not supported. Supported scales: 2^n and 3."
            )
        super(Upsample, self).__init__(*m)


def flow_warp(
    x, flow, interp_mode="bilinear", padding_mode="zeros", align_corners=True
):
    """Warp an image or feature map with optical flow.

    Args:
        x (Tensor): Tensor with size (n, c, h, w).
        flow (Tensor): Tensor with size (n, h, w, 2), normal value.
        interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
        padding_mode (str): 'zeros' or 'border' or 'reflection'.
            Default: 'zeros'.
        align_corners (bool): Before pytorch 1.3, the default value is
            align_corners=True. After pytorch 1.3, the default value is
            align_corners=False. Here, we use the True as default.

    Returns:
        Tensor: Warped image or feature map.
    """
    assert x.size()[-2:] == flow.size()[1:3]
    _, _, h, w = x.size()
    # create mesh grid
    grid_y, grid_x = torch.meshgrid(
        torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x)
    )
    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2 current position
    grid.requires_grad = False

    vgrid = (
        grid + flow
    )  # for all n batches do the same  current position --> warped position
    # scale grid to [-1,1]
    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
    output = F.grid_sample(
        x,
        vgrid_scaled,
        mode=interp_mode,
        padding_mode=padding_mode,
        align_corners=align_corners,
    )

    # TODO, what if align_corners=False
    return output


def resize_flow(flow, size_type, sizes, interp_mode="bilinear", align_corners=False):
    """Resize a flow according to ratio or shape.

    Args:
        flow (Tensor): Precomputed flow. shape [N, 2, H, W].
        size_type (str): 'ratio' or 'shape'.
        sizes (list[int | float]): the ratio for resizing or the final output
            shape.
            1) The order of ratio should be [ratio_h, ratio_w]. For
            downsampling, the ratio should be smaller than 1.0 (i.e., ratio
            < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
            ratio > 1.0).
            2) The order of output_size should be [out_h, out_w].
        interp_mode (str): The mode of interpolation for resizing.
            Default: 'bilinear'.
        align_corners (bool): Whether align corners. Default: False.

    Returns:
        Tensor: Resized flow.
    """
    _, _, flow_h, flow_w = flow.size()
    if size_type == "ratio":
        output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
    elif size_type == "shape":
        output_h, output_w = sizes[0], sizes[1]
    else:
        raise ValueError(
            f"Size type should be ratio or shape, but got type {size_type}."
        )

    input_flow = flow.clone()
    ratio_h = output_h / flow_h
    ratio_w = output_w / flow_w
    input_flow[:, 0, :, :] *= ratio_w
    input_flow[:, 1, :, :] *= ratio_h
    resized_flow = F.interpolate(
        input=input_flow, size=(output_h, output_w), mode=interp_mode
    )  # , align_corners=align_corners)
    return resized_flow


# TODO: may write a cpp file
def pixel_unshuffle(x, scale):
    """Pixel unshuffle.

    Args:
        x (Tensor): Input feature with shape (b, c, hh, hw).
        scale (int): Downsample ratio.

    Returns:
        Tensor: the pixel unshuffled feature.
    """
    b, c, hh, hw = x.size()
    out_channel = c * (scale**2)
    assert hh % scale == 0 and hw % scale == 0
    h = hh // scale
    w = hw // scale
    x_view = x.view(b, c, h, scale, w, scale)
    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2,
        )

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        low = norm_cdf((a - mean) / std)
        up = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [low, up], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * low - 1, 2 * up - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.0))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution.

    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py

    The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.

    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


# From PyTorch
def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            return x
        return tuple(repeat(x, n))

    return parse


to_1tuple = _ntuple(1)
to_2tuple = _ntuple(2)
to_3tuple = _ntuple(3)
to_4tuple = _ntuple(4)
to_ntuple = _ntuple


class ConvResidualBlocks(nn.Module):
    """Conv and residual block used in BasicVSR.

    Args:
        num_in_ch (int): Number of input channels. Default: 3.
        num_out_ch (int): Number of output channels. Default: 64.
        num_block (int): Number of residual blocks. Default: 15.
    """

    def __init__(self, num_in_ch=3, num_out_ch=64, num_block=15):
        super().__init__()
        self.main = nn.Sequential(
            nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1, bias=True),
            nn.LeakyReLU(negative_slope=0.1, inplace=True),
            make_layer(ResidualBlockNoBN, num_block, num_feat=num_out_ch),
        )

    def forward(self, fea):
        return self.main(fea)


class ResidualBlocks(nn.Module):
    """Residual block without shallow conv

    Args:
        num_in_ch (int): Number of input channels. Default: 3.
        num_out_ch (int): Number of output channels. Default: 64.
        num_block (int): Number of residual blocks. Default: 15.
    """

    def __init__(self, num_out_ch=64, num_block=15):
        super().__init__()
        self.main = make_layer(ResidualBlockNoBN, num_block, num_feat=num_out_ch)

    def forward(self, fea):
        return self.main(fea)


class Conv2dwithActication(nn.Module):
    """Conv2d layers with activation"""

    def __init__(self, num_in_ch, num_out_ch, num_layer=1):
        super().__init__()
        main = [
            nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1, bias=True),
            nn.LeakyReLU(negative_slope=0.1, inplace=True),
        ]
        for _ in range(num_layer - 1):
            main.append(nn.Conv2d(num_out_ch, num_out_ch, 3, 1, 1, bias=True))
            main.append(nn.LeakyReLU(negative_slope=0.1, inplace=True))
        self.main = nn.Sequential(*main)

    def forward(self, fea):
        return self.main(fea)


##########################################################################
## LayerNorm


def to_3d(x):
    return rearrange(x, "b c h w -> b (h w) c")


def to_4d(x, h, w):
    return rearrange(x, "b (h w) c -> b c h w", h=h, w=w)


class BiasFree_LayerNorm(nn.Module):
    def __init__(self, normalized_shape):
        super(BiasFree_LayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        normalized_shape = torch.Size(normalized_shape)

        assert len(normalized_shape) == 1

        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.normalized_shape = normalized_shape

    def forward(self, x):
        sigma = x.var(-1, keepdim=True, unbiased=False)
        return x / torch.sqrt(sigma + 1e-5) * self.weight


class WithBias_LayerNorm(nn.Module):
    def __init__(self, normalized_shape):
        super(WithBias_LayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        normalized_shape = torch.Size(normalized_shape)

        assert len(normalized_shape) == 1

        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.normalized_shape = normalized_shape

    def forward(self, x):
        mu = x.mean(-1, keepdim=True)
        sigma = x.var(-1, keepdim=True, unbiased=False)
        return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias


class LayerNorm(nn.Module):
    def __init__(self, dim, LayerNorm_type):
        super(LayerNorm, self).__init__()
        if LayerNorm_type == "BiasFree":
            self.body = BiasFree_LayerNorm(dim)
        else:
            self.body = WithBias_LayerNorm(dim)

    def forward(self, x):
        h, w = x.shape[-2:]
        return to_4d(self.body(to_3d(x)), h, w)


##########################################################################
## Multi-DConv Head Transposed Self-Attention (MDTA)
class Attention(nn.Module):
    def __init__(self, dim_q, dim_k, num_heads, bias):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))

        self.kv = nn.Conv2d(dim_k, dim_k * 2, kernel_size=1, bias=bias)
        self.kv_dwconv = nn.Conv2d(
            dim_k * 2,
            dim_k * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=dim_k * 2,
            bias=bias,
        )
        self.q = nn.Conv2d(dim_q, dim_q, kernel_size=3, stride=1, padding=1, bias=bias)
        self.project_out = nn.Conv2d(dim_q, dim_q, kernel_size=1, bias=bias)

    def forward(self, x, y):
        b, c_kv, h, w = x.shape

        kv = self.kv_dwconv(self.kv(x))
        k, v = kv.chunk(2, dim=1)

        # c_q = y.size(1)
        q = self.q(y)

        q = rearrange(q, "b (head c) h w -> b head c (h w)", head=self.num_heads)
        k = rearrange(k, "b (head c) h w -> b head c (h w)", head=self.num_heads)
        v = rearrange(v, "b (head c) h w -> b head c (h w)", head=self.num_heads)

        q = torch.nn.functional.normalize(q, dim=-1)
        k = torch.nn.functional.normalize(k, dim=-1)

        attn = (q @ k.transpose(-2, -1)) * self.temperature  # c_q, c_kv (64, 32)
        attn = attn.softmax(dim=-1)

        out = attn @ v

        out = rearrange(
            out, "b head c (h w) -> b (head c) h w", head=self.num_heads, h=h, w=w
        )

        out = self.project_out(out)  # (b, 64, h, w)
        return out


##########################################################################
# Attention block with layer norm
class AttentionBlock(nn.Module):
    def __init__(
        self, dim_q, dim_kv, num_heads, bias, LayerNorm_type=None, reduction="sum"
    ):
        super(AttentionBlock, self).__init__()
        if LayerNorm_type is not None:
            self.norm_q = LayerNorm(dim_q, LayerNorm_type)
            self.norm_kv = LayerNorm(dim_kv, LayerNorm_type)
        self.attn = Attention(dim_q, dim_kv, num_heads, bias)
        if reduction == "cat":
            self.fusion = nn.Conv2d(dim_q * 2, dim_q, kernel_size=1, bias=bias)

    def forward(self, x, y):
        if hasattr(self, "norm_kv"):
            x = self.norm_kv(x)
        if hasattr(self, "norm_q"):
            y = self.norm_q(y)

        v = self.attn(x, y)  ## ()

        if hasattr(self, "fusion"):
            return self.fusion(torch.cat((y, v), dim=1))
        else:
            return y + v


##########################################################################
## Simple attention function (parameter-free)


def cross_attention(x, ref):
    b, c0, h, w = x.shape
    _, c1, _, _ = ref.shape
    q, k, v = x.clone(), ref.clone(), ref.clone()

    q = rearrange(q, "b (head c) h w -> b head c (h w)", head=1)
    k = rearrange(k, "b (head c) h w -> b head c (h w)", head=1)
    v = rearrange(v, "b (head c) h w -> b head c (h w)", head=1)

    q = torch.nn.functional.normalize(q, dim=-1)
    k = torch.nn.functional.normalize(k, dim=-1)

    attn = (q @ k.transpose(-2, -1)) * 1  # self.temperature [c0, c1]
    max_, _ = torch.max(attn, dim=-1, keepdim=True)
    attn_submax = attn - max_
    attn = attn.softmax(dim=-1)
    attn = attn * torch.lt(attn_submax, 0).int()
    attn = attn.softmax(dim=-1)

    out = attn @ v  # [c0, hw]

    out = rearrange(out, "b head c (h w) -> b (head c) h w", head=1, h=h, w=w)

    # out = self.project_out(out)
    return out


##########################################################################
## Multi-DConv Head Cross-Attention (MDCA)
class CrossChannelAttention(nn.Module):
    def __init__(self, dim_q, dim_k, num_heads, bias):
        super(CrossChannelAttention, self).__init__()
        self.num_heads = num_heads
        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))

        self.kv = nn.Conv2d(dim_k, dim_k * 2, kernel_size=1, bias=bias)
        self.kv_dwconv = nn.Conv2d(
            dim_k * 2,
            dim_k * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=dim_k * 2,
            bias=bias,
        )
        self.q = nn.Conv2d(dim_q, dim_q, kernel_size=3, stride=1, padding=1, bias=bias)
        self.project_out = nn.Conv2d(dim_q, dim_q, kernel_size=1, bias=bias)

    def forward(self, x, y, return_attn=False):
        b, c_kv, h, w = x.shape

        kv = self.kv_dwconv(self.kv(y))
        k, v = kv.chunk(2, dim=1)

        # c_q = y.size(1)
        q = self.q(x)

        q = rearrange(q, "b (head c) h w -> b head c (h w)", head=self.num_heads)
        k = rearrange(k, "b (head c) h w -> b head c (h w)", head=self.num_heads)
        v = rearrange(v, "b (head c) h w -> b head c (h w)", head=self.num_heads)

        q = torch.nn.functional.normalize(q, dim=-1)
        k = torch.nn.functional.normalize(k, dim=-1)

        attn = (q @ k.transpose(-2, -1)) * self.temperature

        attn = attn.softmax(dim=-1)
        out = attn @ v

        out = rearrange(
            out, "b head c (h w) -> b (head c) h w", head=self.num_heads, h=h, w=w
        )

        out = self.project_out(out)  # (b, 64, h, w)
        if return_attn:
            return out, attn
        else:
            return out


##########################################################################
# Channel-wise Attention block with layer norm
class ChannelAttentionBlock(nn.Module):
    def __init__(
        self,
        dim_q,
        dim_kv,
        dim_out,
        num_heads,
        bias,
        LayerNorm_type=None,
        reduction=True,
        ch_compress=False,
        squeeze_factor=1,
    ):
        super(ChannelAttentionBlock, self).__init__()
        self.ch_compress = ch_compress and (squeeze_factor > 1)
        if self.ch_compress:
            if LayerNorm_type is not None:
                self.norm_q = LayerNorm(dim_q // squeeze_factor, LayerNorm_type)
                self.norm_kv = LayerNorm(dim_kv // squeeze_factor, LayerNorm_type)
            self.attn = CrossChannelAttention(
                dim_q // squeeze_factor, dim_kv // squeeze_factor, num_heads, bias
            )
        else:
            if LayerNorm_type is not None:
                self.norm_q = LayerNorm(dim_q, LayerNorm_type)
                self.norm_kv = LayerNorm(dim_kv, LayerNorm_type)
            self.attn = CrossChannelAttention(dim_q, dim_kv, num_heads, bias)
        self.reduction = reduction
        if reduction:
            self.norm_out = LayerNorm(dim_kv, LayerNorm_type)
            self.ffn = nn.Sequential(
                nn.Conv2d(dim_kv + dim_q, dim_out, kernel_size=1, bias=bias),
                nn.Conv2d(
                    dim_out,
                    dim_out,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    groups=dim_out,
                    bias=bias,
                ),
                nn.GELU(),
                nn.Conv2d(dim_out, dim_out, kernel_size=1, bias=bias),
            )

        if self.ch_compress:
            self.compress_q = nn.Conv2d(
                dim_q,
                dim_q // squeeze_factor,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=bias,
            )
            self.compress_kv = nn.Conv2d(
                dim_kv,
                dim_kv // squeeze_factor,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=bias,
            )
            self.expand = nn.Conv2d(
                dim_q // squeeze_factor,
                dim_q,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=bias,
            )

    def forward(self, x, y, return_attn=False):
        if self.ch_compress:
            x_compressed = self.compress_q(x)
            y_compressed = self.compress_kv(y)
            if hasattr(self, "norm_kv"):
                y_compressed = self.norm_kv(y_compressed)
            if hasattr(self, "norm_q"):
                x_compressed = self.norm_q(x_compressed)
            v = self.expand(self.attn(x_compressed, y_compressed, return_attn))
        else:
            if hasattr(self, "norm_kv"):
                y = self.norm_kv(y)
            if hasattr(self, "norm_q"):
                x = self.norm_q(x)
            v = self.attn(x, y, return_attn)  ## ()
        if self.reduction:
            out = self.ffn(torch.cat([x, self.norm_out(v)], dim=1))
        return out