# 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, 14, 9, 31), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ({}, {}), '()', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((10, 8, 10, 37), 'torch.randn', 'torch.randn', ({(10, 20, 10, 36): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((19, 14, 19, 57), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((20, 8, 20, 37), 'torch.randn', 'torch.randn', ({(20, 20, 20, 36): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((30, 14, 30, 62), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((31, 8, 31, 37), 'torch.randn', 'torch.randn', ({(31, 20, 31, 36): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((40, 7, 40, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((88, 11, 88, 51), 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((89, 8, 89, 36), 'torch.rand', 'torch.rand', ({(89, 19, 89, 35): '(2, 128, 64, 64)'}, {}), '((2, 128, 64, 64))', False, 'import torch\n'), ((90, 11, 90, 38), 'torch.zeros', 'torch.zeros', ({(90, 23, 90, 37): '(2, 1, 64, 64)'}, {}), '((2, 1, 64, 64))', False, 'import torch\n'), ((98, 7, 98, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((41, 18, 41, 35), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ({}, {}), '()', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((42, 12, 42, 41), 'torch.randn', 'torch.randn', ({(42, 24, 42, 40): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((51, 18, 51, 61), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((52, 12, 52, 41), 'torch.randn', 'torch.randn', ({(52, 24, 52, 40): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((62, 18, 62, 66), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((63, 12, 63, 41), 'torch.randn', 'torch.randn', ({(63, 24, 63, 40): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((73, 18, 74, 56), 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), ((75, 12, 75, 41), 'torch.randn', 'torch.randn', ({(75, 24, 75, 40): '(2, 5, 256, 256)'}, {}), '((2, 5, 256, 256))', False, 'import torch\n'), ((105, 15, 106, 52), 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', (), '', False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n')]