# Copyright (c) OpenMMLab. All rights reserved.
import pytest
import torch

from mmcv.cnn.bricks import ContextBlock


def test_context_block():
    with pytest.raises(AssertionError):
        # pooling_type should be in ['att', 'avg']
        ContextBlock(16, 1. / 4, pooling_type='unsupport_type')

    with pytest.raises(AssertionError):
        # fusion_types should be of type list or tuple
        ContextBlock(16, 1. / 4, fusion_types='unsupport_type')

    with pytest.raises(AssertionError):
        # fusion_types should be in ['channel_add', 'channel_mul']
        ContextBlock(16, 1. / 4, fusion_types=('unsupport_type', ))

    # test pooling_type='att'
    imgs = torch.randn(2, 16, 20, 20)
    context_block = ContextBlock(16, 1. / 4, pooling_type='att')
    out = context_block(imgs)
    assert context_block.conv_mask.in_channels == 16
    assert context_block.conv_mask.out_channels == 1
    assert out.shape == imgs.shape

    # test pooling_type='avg'
    imgs = torch.randn(2, 16, 20, 20)
    context_block = ContextBlock(16, 1. / 4, pooling_type='avg')
    out = context_block(imgs)
    assert hasattr(context_block, 'avg_pool')
    assert out.shape == imgs.shape

    # test fusion_types=('channel_add',)
    imgs = torch.randn(2, 16, 20, 20)
    context_block = ContextBlock(16, 1. / 4, fusion_types=('channel_add', ))
    out = context_block(imgs)
    assert context_block.channel_add_conv is not None
    assert context_block.channel_mul_conv is None
    assert out.shape == imgs.shape

    # test fusion_types=('channel_mul',)
    imgs = torch.randn(2, 16, 20, 20)
    context_block = ContextBlock(16, 1. / 4, fusion_types=('channel_mul', ))
    out = context_block(imgs)
    assert context_block.channel_add_conv is None
    assert context_block.channel_mul_conv is not None
    assert out.shape == imgs.shape

    # test fusion_types=('channel_add', 'channel_mul')
    imgs = torch.randn(2, 16, 20, 20)
    context_block = ContextBlock(
        16, 1. / 4, fusion_types=('channel_add', 'channel_mul'))
    out = context_block(imgs)
    assert context_block.channel_add_conv is not None
    assert context_block.channel_mul_conv is not None
    assert out.shape == imgs.shape