# Copyright (c) OpenMMLab. All rights reserved.
import torch
from torch.nn.modules.batchnorm import _BatchNorm

from mmpose.models.backbones import HRNet
from mmpose.models.backbones.hrnet import HRModule
from mmpose.models.backbones.resnet import BasicBlock, Bottleneck


def is_block(modules):
    """Check if is HRModule building block."""
    if isinstance(modules, (HRModule, )):
        return True
    return False


def is_norm(modules):
    """Check if is one of the norms."""
    if isinstance(modules, (_BatchNorm, )):
        return True
    return False


def all_zeros(modules):
    """Check if the weight(and bias) is all zero."""
    weight_zero = torch.equal(modules.weight.data,
                              torch.zeros_like(modules.weight.data))
    if hasattr(modules, 'bias'):
        bias_zero = torch.equal(modules.bias.data,
                                torch.zeros_like(modules.bias.data))
    else:
        bias_zero = True

    return weight_zero and bias_zero


def test_hrmodule():
    # Test HRModule forward
    block = HRModule(
        num_branches=1,
        blocks=BasicBlock,
        num_blocks=(4, ),
        in_channels=[
            64,
        ],
        num_channels=(64, ))

    x = torch.randn(2, 64, 56, 56)
    x_out = block([x])
    assert x_out[0].shape == torch.Size([2, 64, 56, 56])


def test_hrnet_backbone():
    extra = dict(
        stage1=dict(
            num_modules=1,
            num_branches=1,
            block='BOTTLENECK',
            num_blocks=(4, ),
            num_channels=(64, )),
        stage2=dict(
            num_modules=1,
            num_branches=2,
            block='BASIC',
            num_blocks=(4, 4),
            num_channels=(32, 64)),
        stage3=dict(
            num_modules=4,
            num_branches=3,
            block='BASIC',
            num_blocks=(4, 4, 4),
            num_channels=(32, 64, 128)),
        stage4=dict(
            num_modules=3,
            num_branches=4,
            block='BASIC',
            num_blocks=(4, 4, 4, 4),
            num_channels=(32, 64, 128, 256)))

    model = HRNet(extra, in_channels=3)

    imgs = torch.randn(2, 3, 224, 224)
    feat = model(imgs)
    assert len(feat) == 1
    assert feat[0].shape == torch.Size([2, 32, 56, 56])

    # Test HRNet zero initialization of residual
    model = HRNet(extra, in_channels=3, zero_init_residual=True)
    model.init_weights()
    for m in model.modules():
        if isinstance(m, Bottleneck):
            assert all_zeros(m.norm3)
    model.train()

    imgs = torch.randn(2, 3, 224, 224)
    feat = model(imgs)
    assert len(feat) == 1
    assert feat[0].shape == torch.Size([2, 32, 56, 56])

    # Test HRNet with the first three stages frozen
    frozen_stages = 3
    model = HRNet(extra, in_channels=3, frozen_stages=frozen_stages)
    model.init_weights()
    model.train()
    if frozen_stages >= 0:
        assert model.norm1.training is False
        assert model.norm2.training is False
        for layer in [model.conv1, model.norm1, model.conv2, model.norm2]:
            for param in layer.parameters():
                assert param.requires_grad is False

    for i in range(1, frozen_stages + 1):
        if i == 1:
            layer = getattr(model, 'layer1')
        else:
            layer = getattr(model, f'stage{i}')
        for mod in layer.modules():
            if isinstance(mod, _BatchNorm):
                assert mod.training is False
        for param in layer.parameters():
            assert param.requires_grad is False

        if i < 4:
            layer = getattr(model, f'transition{i}')
        for mod in layer.modules():
            if isinstance(mod, _BatchNorm):
                assert mod.training is False
        for param in layer.parameters():
            assert param.requires_grad is False