import torch

from torch import nn
from einops.layers.torch import Rearrange

def init_layer(layer):
    nn.init.xavier_uniform_(layer.weight)

    if hasattr(layer, "bias") and layer.bias is not None: 
        layer.bias.data.fill_(0.0)

def init_bn(bn):
    bn.bias.data.fill_(0.0)
    bn.weight.data.fill_(1.0)
    bn.running_mean.data.fill_(0.0)
    bn.running_var.data.fill_(1.0)

class BiGRU(nn.Module):
    def __init__(
        self, 
        patch_size, 
        channels, 
        depth
    ):
        super(BiGRU, self).__init__()
        patch_width, patch_height = patch_size
        patch_dim = channels * patch_height * patch_width

        self.to_patch_embedding = nn.Sequential(
            Rearrange(
                'b c (w p1) (h p2) -> b (w h) (p1 p2 c)', 
                p1=patch_width, 
                p2=patch_height
            )
        )

        self.gru = nn.GRU(
            patch_dim, 
            patch_dim // 2, 
            num_layers=depth, 
            batch_first=True, 
            bidirectional=True
        )

    def forward(self, x):
        x = self.to_patch_embedding(x)

        try:
            return self.gru(x)[0]
        except:
            torch.backends.cudnn.enabled = False
            return self.gru(x)[0]

class ResConvBlock(nn.Module):
    def __init__(
        self, 
        in_planes, 
        out_planes
    ):
        super(ResConvBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(
            in_planes, 
            momentum=0.01
        )
        self.bn2 = nn.BatchNorm2d(
            out_planes, 
            momentum=0.01
        )
        self.act1 = nn.PReLU()
        self.act2 = nn.PReLU()
        self.conv1 = nn.Conv2d(
            in_planes, 
            out_planes, 
            (3, 3), 
            padding=(1, 1), 
            bias=False
        )
        self.conv2 = nn.Conv2d(
            out_planes, 
            out_planes, 
            (3, 3), 
            padding=(1, 1), 
            bias=False
        )
        self.is_shortcut = False

        if in_planes != out_planes:
            self.shortcut = nn.Conv2d(
                in_planes, 
                out_planes, 
                (1, 1)
            )
            self.is_shortcut = True

        self.init_weights()

    def init_weights(self):
        init_bn(self.bn1)
        init_bn(self.bn2)

        init_layer(self.conv1)
        init_layer(self.conv2)

        if self.is_shortcut: init_layer(self.shortcut)

    def forward(self, x):
        out = self.conv1(
            self.act1(self.bn1(x))
        )
        out = self.conv2(
            self.act2(self.bn2(out))
        )

        if self.is_shortcut: return self.shortcut(x) + out
        else: return out + x