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

from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder
from mmedit.models.common import SimpleGatedConvModule


def test_deepfill_enc():
    encoder = DeepFillEncoder()
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.stride == (2, 2)
    assert encoder.enc2.out_channels == 64

    encoder = DeepFillEncoder(encoder_type='stage2_conv')
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.out_channels == 32
    assert encoder.enc3.out_channels == 64
    assert encoder.enc4.out_channels == 64

    encoder = DeepFillEncoder(encoder_type='stage2_attention')
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.out_channels == 32
    assert encoder.enc3.out_channels == 64
    assert encoder.enc4.out_channels == 128
    if torch.cuda.is_available():
        encoder = DeepFillEncoder().cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.stride == (2, 2)
        assert encoder.enc2.out_channels == 64

        encoder = DeepFillEncoder(encoder_type='stage2_conv').cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.out_channels == 32
        assert encoder.enc3.out_channels == 64
        assert encoder.enc4.out_channels == 64

        encoder = DeepFillEncoder(encoder_type='stage2_attention').cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.out_channels == 32
        assert encoder.enc3.out_channels == 64
        assert encoder.enc4.out_channels == 128

        encoder = DeepFillEncoder(
            conv_type='gated_conv', channel_factor=0.75).cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 96, 64, 64)
        assert isinstance(encoder.enc2, SimpleGatedConvModule)
        assert encoder.enc2.conv.stride == (2, 2)
        assert encoder.enc2.conv.out_channels == 48 * 2


def test_deepfill_contextual_attention_neck():
    # TODO: add unittest for contextual attention module
    neck = ContextualAttentionNeck(in_channels=128)
    x = torch.rand((2, 128, 64, 64))
    mask = torch.zeros((2, 1, 64, 64))
    mask[..., 20:100, 23:90] = 1.

    res, offset = neck(x, mask)

    assert res.shape == (2, 128, 64, 64)
    assert offset.shape == (2, 32, 32, 32, 32)

    if torch.cuda.is_available():
        neck.cuda()
        res, offset = neck(x.cuda(), mask.cuda())

        assert res.shape == (2, 128, 64, 64)
        assert offset.shape == (2, 32, 32, 32, 32)

        neck = ContextualAttentionNeck(
            in_channels=128, conv_type='gated_conv').cuda()
        res, offset = neck(x.cuda(), mask.cuda())
        assert res.shape == (2, 128, 64, 64)
        assert offset.shape == (2, 32, 32, 32, 32)
        assert isinstance(neck.conv1, SimpleGatedConvModule)

[(9, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (10, 'torch.randn', 'torch.randn', False, 'import torch\n'), (19, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (20, 'torch.randn', 'torch.randn', False, 'import torch\n'), (30, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (31, 'torch.randn', 'torch.randn', False, 'import torch\n'), (40, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (88, 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (89, 'torch.rand', 'torch.rand', False, 'import torch\n'), (90, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (98, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (41, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (42, 'torch.randn', 'torch.randn', False, 'import torch\n'), (51, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (52, 'torch.randn', 'torch.randn', False, 'import torch\n'), (62, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (63, 'torch.randn', 'torch.randn', False, 'import torch\n'), (73, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (75, 'torch.randn', 'torch.randn', False, 'import torch\n'), (105, 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n')]