import torch
import torch.nn as nn

__all__ = ['PreActOctResNet', 'pre_act_oct_resnet26', 'pre_act_oct_resnet50',
            'pre_act_oct_resnet101', 'pre_act_oct_resnet152', 'pre_act_oct_resnet200']


class OctaveConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, alpha_in=0.5, alpha_out=0.5, stride=1, padding=0, dilation=1,
                 groups=1, bias=False):
        super(OctaveConv, self).__init__()
        self.downsample = nn.AvgPool2d(kernel_size=(2, 2), stride=2)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
        assert stride == 1 or stride == 2, "Stride should be 1 or 2."
        self.stride = stride
        assert 0 <= alpha_in <= 1 and 0 <= alpha_out <= 1, "Alphas should be in the interval from 0 to 1."
        self.alpha_in, self.alpha_out = alpha_in, alpha_out
        self.conv_l2l = None if alpha_in == 0 or alpha_out == 0 else \
                        nn.Conv2d(int(alpha_in * in_channels), int(alpha_out * out_channels),
                                  kernel_size, 1, padding, dilation, groups, bias)
        self.conv_l2h = None if alpha_in == 0 or alpha_out == 1 else \
                        nn.Conv2d(int(alpha_in * in_channels), out_channels - int(alpha_out * out_channels),
                                  kernel_size, 1, padding, dilation, groups, bias)
        self.conv_h2l = None if alpha_in == 1 or alpha_out == 0 else \
                        nn.Conv2d(in_channels - int(alpha_in * in_channels), int(alpha_out * out_channels),
                                  kernel_size, 1, padding, dilation, groups, bias)
        self.conv_h2h = None if alpha_in == 1 or alpha_out == 1 else \
                        nn.Conv2d(in_channels - int(alpha_in * in_channels), out_channels - int(alpha_out * out_channels),
                                  kernel_size, 1, padding, dilation, groups, bias)

    def forward(self, x):
        x_h, x_l = x if type(x) is tuple else (x, None)

        if x_h is not None:
            x_h = self.downsample(x_h) if self.stride == 2 else x_h
            x_h2h = self.conv_h2h(x_h)
            x_h2l = self.conv_h2l(self.downsample(x_h)) if self.alpha_out > 0 else None
        if x_l is not None:
            x_l2h = self.conv_l2h(x_l)
            x_l2h = self.upsample(x_l2h) if self.stride == 1 else x_l2h
            x_l2l = self.downsample(x_l) if self.stride == 2 else x_l
            x_l2l = self.conv_l2l(x_l2l) if self.alpha_out > 0 else None
            x_h = x_l2h + x_h2h
            x_l = x_h2l + x_l2l if x_h2l is not None and x_l2l is not None else None
            return x_h, x_l
        else:
            return x_h2h, x_h2l

class OctaveBatchNorm2d(nn.Module):
    def __init__(self, in_channels, alpha=0.5, norm_layer=nn.BatchNorm2d):
        super(OctaveBatchNorm2d, self).__init__()
        self.bn_h = None if alpha == 1 else norm_layer(int(in_channels * (1 - alpha)))
        self.bn_l = None if alpha == 0 else norm_layer(int(in_channels * alpha))
    
    def forward(self, x):
        x_h, x_l = x if type(x) is tuple else (x, None)
        if x_h is not None:
            x_h = self.bn_h(x_h)
        if x_l is not None:
            x_l = self.bn_l(x_l)
        return x_h, x_l


class OctaveReLU(nn.Module):
    def __init__(self, in_channels, alpha=0.5, norm_layer=nn.BatchNorm2d):
        super(OctaveReLU, self).__init__()
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x_h, x_l = x if type(x) is tuple else (x, None)
        x_h  = self.relu(x_h)
        x_l = self.relu(x_l) if x_l is not None else None
        return x_h, x_l

class OctaveDropout2d(nn.Module):
    def __init__(self, p=0.5):
        super(OctaveDropout2d, self).__init__(self)
        self.dropout = nn.Dropout2d(p=p)

    def forward(self, x):
        x_h, x_l = x if type(x) is tuple else (x, None)
        x_h = self.dropout(x_h)
        x_l = self.dropout(x_l) if x_l is not None else None
        return x_h, x_l


#######################

class Conv_BN(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, alpha_in=0.5, alpha_out=0.5, stride=1, padding=0, dilation=1,
                 groups=1, bias=False, norm_layer=nn.BatchNorm2d):
        super(Conv_BN, self).__init__()
        self.conv = OctaveConv(in_channels, out_channels, kernel_size, alpha_in, alpha_out, stride, padding, dilation,
                               groups, bias)
        self.bn_h = None if alpha_out == 1 else norm_layer(int(out_channels * (1 - alpha_out)))
        self.bn_l = None if alpha_out == 0 else norm_layer(int(out_channels * alpha_out))

    def forward(self, x):
        x_h, x_l = self.conv(x)
        x_h = self.bn_h(x_h)
        x_l = self.bn_l(x_l) if x_l is not None else None
        return x_h, x_l


class Conv_BN_ACT(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, alpha_in=0.5, alpha_out=0.5, stride=1, padding=0, dilation=1,
                 groups=1, bias=False, norm_layer=nn.BatchNorm2d, activation_layer=nn.ReLU):
        super(Conv_BN_ACT, self).__init__()
        self.conv = OctaveConv(in_channels, out_channels, kernel_size, alpha_in, alpha_out, stride, padding, dilation,
                               groups, bias)
        self.bn_h = None if alpha_out == 1 else norm_layer(int(out_channels * (1 - alpha_out)))
        self.bn_l = None if alpha_out == 0 else norm_layer(int(out_channels * alpha_out))
        self.act = activation_layer(inplace=True)

    def forward(self, x):
        x_h, x_l = self.conv(x)
        x_h = self.act(self.bn_h(x_h))
        x_l = self.act(self.bn_l(x_l)) if x_l is not None else None
        return x_h, x_l

###################


class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, alpha_in=0.5, alpha_out=0.5, norm_layer=None, output=False):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        self.conv1 = Conv_BN_ACT(inplanes, width, kernel_size=3, alpha_in=alpha_in, alpha_out=alpha_out,
                                 norm_layer=norm_layer)
        self.conv2 = Conv_BN(width, planes * self.expansion, kernel_size=3, norm_layer=norm_layer,
                             alpha_in=0 if output else 0.5, alpha_out=0 if output else 0.5)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity_h = x[0] if type(x) is tuple else x
        identity_l = x[1] if type(x) is tuple else None

        x_h, x_l = self.conv1(x)
        x_h, x_l = self.conv2((x_h, x_l))
        if self.downsample is not None:
            identity_h, identity_l = self.downsample(x)

        x_h += identity_h
        x_l = x_l + identity_l if identity_l is not None else None

        x_h = self.relu(x_h)
        x_l = self.relu(x_l) if x_l is not None else None

        return x_h, x_l


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, alpha_in=0.5, alpha_out=0.5, norm_layer=None, output=False, dropout=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups

        self.downsample = downsample
        self.stride = stride


        self.bn1 = OctaveBatchNorm2d(inplanes, alpha_in)
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = OctaveConv(inplanes, width, kernel_size=1,
                                bias=False, alpha_in=alpha_in, alpha_out=alpha_out)
        self.bn2 = OctaveBatchNorm2d(width, alpha_out)
        self.conv2 = OctaveConv(width, width, kernel_size=3, stride=stride, padding=1, bias=False, groups=groups,
                                alpha_in=0 if output else 0.5, alpha_out=0 if output else 0.5)
        self.bn3 = OctaveBatchNorm2d(width, alpha=0 if output else 0.5)
        self.conv3 = OctaveConv(width, planes * self.expansion, kernel_size=1, bias=False,
                                alpha_in=0 if output else 0.5, alpha_out=0 if output else 0.5)
        
        self.dropout = nn.Dropout2d(p=dropout) if dropout is not None else None

    def forward(self, x):
        x_h, x_l = self.bn1(x)
        x_h = self.relu(x_h)
        x_l = self.relu(x_l) if x_l is not None else None

        identity_h = x_h
        identity_l = x_l

        x_h, x_l = self.conv1((x_h, x_l))

        x_h, x_l = self.bn2((x_h, x_l))
        x_h  = self.relu(x_h)
        x_l = self.relu(x_l) if x_l is not None else None
        x_h, x_l = self.conv2((x_h, x_l))

        x_h, x_l = self.bn3((x_h, x_l))
        x_h  = self.relu(x_h)
        x_l = self.relu(x_l) if x_l is not None else None
        x_h = self.dropout(x_h) if self.dropout is not None else x_h
        x_l = self.dropout(x_l) if self.dropout is not None else x_l
        x_h, x_l = self.conv3((x_h, x_l))

        if self.downsample is not None:
            identity_h, identity_l = self.downsample(x)

        x_h += identity_h
        x_l = x_l + identity_l if identity_l is not None else None

        x_l = self.relu(x_l) if x_l is not None else None

        return x_h, x_l


class PreActOctResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 groups=1, width_per_group=64, norm_layer=None):
        super(PreActOctResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d

        self.inplanes = 64
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer, alpha_in=0)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer, alpha_out=0, output=True)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    if m.bn3.bn_h is not None:
                        nn.init.constant_(m.bn3.bn_h.weight, 0)
                    if m.bn3.bn_l is not None:
                        nn.init.constant_(m.bn3.bn_l.weight, 0)
                elif isinstance(m, BasicBlock):
                    if m.conv2.bn_h is not None:
                        nn.init.constant_(m.conv2.bn_h.weight, 0)
                    if m.conv2.bn_l is not None:
                        nn.init.constant_(m.conv2.bn_l.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, alpha_in=0.5, alpha_out=0.5, norm_layer=None, output=False, dropout=None):
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = OctaveConv(self.inplanes, planes * block.expansion, kernel_size=1,
                                        stride=stride, alpha_in=alpha_in, alpha_out=alpha_out)

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, alpha_in, alpha_out, norm_layer, output))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, norm_layer=norm_layer,
                                alpha_in=0 if output else 0.5, alpha_out=0 if output else 0.5, output=output))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool(x)
        x_h, x_l = self.layer1(x)
        x_h, x_l = self.layer2((x_h,x_l))
        x_h, x_l = self.layer3((x_h,x_l))
        x_h, x_l = self.layer4((x_h,x_l))
        x = self.avgpool(x_h)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x


def pre_act_oct_resnet26(pretrained=False, **kwargs):
    """Constructs a Octave ResNet-26 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PreActOctResNet(Bottleneck, [2, 2, 2, 2], **kwargs)
    return model


def pre_act_oct_resnet50(pretrained=False, **kwargs):
    """Constructs a Octave ResNet-50 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PreActOctResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    return model


def pre_act_oct_resnet101(pretrained=False, **kwargs):
    """Constructs a Octave ResNet-101 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PreActOctResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    return model


def pre_act_oct_resnet152(pretrained=False, **kwargs):
    """Constructs a Octave ResNet-152 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PreActOctResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
    return model


def pre_act_oct_resnet200(pretrained=False, **kwargs):
    """Constructs a Octave ResNet-200 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = PreActOctResNet(Bottleneck, [3, 24, 36, 3], **kwargs)
    return model


def test():
    inputs = torch.FloatTensor(4, 3, 256, 256)
    model = pre_act_oct_resnet50(zero_init_residual=True)
    print(model(inputs).shape)


if __name__ == '__main__':
    test()