import math
import numbers

import torch
import torch.nn as nn
from einops import rearrange
from torch.nn import functional as F


# ---------------------------------------------------------------------------------------------------------------------
# Layer Norm
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)


# ---------------------------------------------------------------------------------------------------------------------


# FFN
class FeedForward(nn.Module):
    def __init__(self, dim, ffn_expansion_factor, bias, input_resolution=None):
        super(FeedForward, self).__init__()

        self.input_resolution = input_resolution
        self.dim = dim
        self.ffn_expansion_factor = ffn_expansion_factor

        hidden_features = int(dim * ffn_expansion_factor)
        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
        self.dwconv = nn.Conv2d(
            hidden_features * 2,
            hidden_features * 2,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=hidden_features * 2,
            bias=bias,
        )
        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.project_in(x)
        x1, x2 = self.dwconv(x).chunk(2, dim=1)
        x = F.gelu(x1) * x2
        x = self.project_out(x)
        return x


class NormLayer(nn.Module):
    """Normalization Layers.
    ------------
    # Arguments
        - channels: input channels, for batch norm and instance norm.
        - input_size: input shape without batch size, for layer norm.
    """

    def __init__(self, channels, normalize_shape=None, norm_type="bn"):
        super(NormLayer, self).__init__()
        norm_type = norm_type.lower()
        if norm_type == "bn":
            self.norm = nn.BatchNorm2d(channels)
        elif norm_type == "in":
            self.norm = nn.InstanceNorm2d(channels, affine=True)
        elif norm_type == "gn":
            self.norm = nn.GroupNorm(32, channels, affine=True)
        elif norm_type == "pixel":
            self.norm = lambda x: F.normalize(x, p=2, dim=1)
        elif norm_type == "layer":
            self.norm = nn.LayerNorm(normalize_shape)
        elif norm_type == "none":
            self.norm = lambda x: x
        else:
            assert 1 == 0, "Norm type {} not support.".format(norm_type)

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


class ReluLayer(nn.Module):
    """Relu Layer.
    ------------
    # Arguments
        - relu type: type of relu layer, candidates are
            - ReLU
            - LeakyReLU: default relu slope 0.2
            - PRelu
            - SELU
            - none: direct pass
    """

    def __init__(self, channels, relu_type="relu"):
        super(ReluLayer, self).__init__()
        relu_type = relu_type.lower()
        if relu_type == "relu":
            self.func = nn.ReLU(True)
        elif relu_type == "leakyrelu":
            self.func = nn.LeakyReLU(0.2, inplace=True)
        elif relu_type == "prelu":
            self.func = nn.PReLU(channels)
        elif relu_type == "selu":
            self.func = nn.SELU(True)
        elif relu_type == "none":
            self.func = lambda x: x
        else:
            assert 1 == 0, "Relu type {} not support.".format(relu_type)

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


class ConvLayer(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        scale="none",
        norm_type="none",
        relu_type="none",
        use_pad=True,
    ):
        super(ConvLayer, self).__init__()
        self.use_pad = use_pad

        bias = True if norm_type in ["pixel", "none"] else False
        stride = 2 if scale == "down" else 1

        self.scale_func = lambda x: x
        if scale == "up":
            self.scale_func = lambda x: nn.functional.interpolate(
                x, scale_factor=2, mode="nearest"
            )

        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
        self.conv2d = nn.Conv2d(
            in_channels, out_channels, kernel_size, stride, bias=bias
        )

        self.relu = ReluLayer(out_channels, relu_type)
        self.norm = NormLayer(out_channels, norm_type=norm_type)

    def forward(self, x):
        out = self.scale_func(x)
        if self.use_pad:
            out = self.reflection_pad(out)
        out = self.conv2d(out)
        out = self.norm(out)
        out = self.relu(out)
        return out


class ResidualBlock(nn.Module):
    """
    Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
    ------------------
    # Args
        - hg_depth: depth of HourGlassBlock. 0: don't use attention map.
        - use_pmask: whether use previous mask as HourGlassBlock input.
    """

    def __init__(
        self,
        c_in,
        c_out,
        relu_type="prelu",
        norm_type="bn",
        scale="none",
        hg_depth=2,
        att_name="spar",
    ):
        super(ResidualBlock, self).__init__()
        self.c_in = c_in
        self.c_out = c_out
        self.norm_type = norm_type
        self.relu_type = relu_type
        self.hg_depth = hg_depth

        kwargs = {"norm_type": norm_type, "relu_type": relu_type}

        if scale == "none" and c_in == c_out:
            self.shortcut_func = lambda x: x
        else:
            self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)

        self.preact_func = nn.Sequential(
            NormLayer(c_in, norm_type=self.norm_type),
            ReluLayer(c_in, self.relu_type),
        )

        if scale == "down":
            scales = ["none", "down"]
        elif scale == "up":
            scales = ["up", "none"]
        elif scale == "none":
            scales = ["none", "none"]

        self.conv1 = ConvLayer(c_in, c_out, 3, scales[0], **kwargs)
        self.conv2 = ConvLayer(
            c_out, c_out, 3, scales[1], norm_type=norm_type, relu_type="none"
        )

        if att_name.lower() == "spar":
            c_attn = 1
        elif att_name.lower() == "spar3d":
            c_attn = c_out
        else:
            raise Exception("Attention type {} not implemented".format(att_name))

        self.att_func = HourGlassBlock(self.hg_depth, c_out, c_attn, **kwargs)

    def forward(self, x):
        identity = self.shortcut_func(x)
        out = self.preact_func(x)
        out = self.conv1(out)
        out = self.conv2(out)
        out = identity + self.att_func(out)
        return out


class HourGlassBlock(nn.Module):
    """Simplified HourGlass block.
    Reference: https://github.com/1adrianb/face-alignment
    --------------------------
    """

    def __init__(
        self,
        depth,
        c_in,
        c_out,
        c_mid=64,
        norm_type="bn",
        relu_type="prelu",
    ):
        super(HourGlassBlock, self).__init__()
        self.depth = depth
        self.c_in = c_in
        self.c_mid = c_mid
        self.c_out = c_out
        self.kwargs = {"norm_type": norm_type, "relu_type": relu_type}

        if self.depth:
            self._generate_network(self.depth)
            self.out_block = nn.Sequential(
                ConvLayer(self.c_mid, self.c_out, norm_type="none", relu_type="none"),
                nn.Sigmoid(),
            )

    def _generate_network(self, level):
        if level == self.depth:
            c1, c2 = self.c_in, self.c_mid
        else:
            c1, c2 = self.c_mid, self.c_mid

        self.add_module("b1_" + str(level), ConvLayer(c1, c2, **self.kwargs))
        self.add_module(
            "b2_" + str(level), ConvLayer(c1, c2, scale="down", **self.kwargs)
        )
        if level > 1:
            self._generate_network(level - 1)
        else:
            self.add_module(
                "b2_plus_" + str(level),
                ConvLayer(self.c_mid, self.c_mid, **self.kwargs),
            )

        self.add_module(
            "b3_" + str(level),
            ConvLayer(self.c_mid, self.c_mid, scale="up", **self.kwargs),
        )

    def _forward(self, level, in_x):
        up1 = self._modules["b1_" + str(level)](in_x)
        low1 = self._modules["b2_" + str(level)](in_x)
        if level > 1:
            low2 = self._forward(level - 1, low1)
        else:
            low2 = self._modules["b2_plus_" + str(level)](low1)

        up2 = self._modules["b3_" + str(level)](low2)
        if up1.shape[2:] != up2.shape[2:]:
            up2 = nn.functional.interpolate(up2, up1.shape[2:])

        return up1 + up2

    def forward(self, x, pmask=None):
        if self.depth == 0:
            return x
        input_x = x
        x = self._forward(self.depth, x)
        self.att_map = self.out_block(x)
        x = input_x * self.att_map
        return x


class UpSampler(nn.Sequential):
    def __init__(self, scale, n_feats, out_feats):
        m = []
        if scale == 8:
            kernel_size = 3
        elif scale == 16:
            kernel_size = 5
        else:
            kernel_size = 1

        if (scale & (scale - 1)) == 0:  # Is scale = 2^n?
            for _ in range(int(math.log(scale, 2))):
                m.append(
                    nn.Conv2d(
                        in_channels=n_feats,
                        out_channels=4 * out_feats,
                        kernel_size=kernel_size,
                        stride=1,
                        padding=kernel_size // 2,
                    )
                )
                m.append(nn.PixelShuffle(upscale_factor=2))
                m.append(nn.PReLU())
        super(UpSampler, self).__init__(*m)