import numpy as np
import torch
import torch.nn as nn
from mmcv.runner import obj_from_dict
from mmcv.utils.config import Config

from mmedit.models import build_model
from mmedit.models.losses import L1Loss
from mmedit.models.registry import COMPONENTS


@COMPONENTS.register_module()
class BP(nn.Module):
    """A simple BP network for testing LIIF.

    Args:
        in_dim (int): Input dimension.
        out_dim (int): Output dimension.
    """

    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.layer = nn.Linear(in_dim, out_dim)

    def forward(self, x):
        shape = x.shape[:-1]
        x = self.layer(x.view(-1, x.shape[-1]))
        return x.view(*shape, -1)


def test_liif():

    model_cfg = dict(
        type='LIIF',
        generator=dict(
            type='EDSR',
            in_channels=3,
            out_channels=3,
            mid_channels=8,
            num_blocks=1),
        imnet=dict(type='BP', in_dim=8, out_dim=3),
        local_ensemble=True,
        feat_unfold=True,
        cell_decode=True,
        rgb_mean=(0.4488, 0.4371, 0.4040),
        rgb_std=(1., 1., 1.),
        eval_bsize=30000,
        pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'))

    scale_max = 4
    train_cfg = None
    test_cfg = Config(dict(metrics=['PSNR', 'SSIM'], crop_border=scale_max))

    # build restorer
    restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)

    # test attributes
    assert restorer.__class__.__name__ == 'LIIF'
    assert isinstance(restorer.imnet, BP)
    assert isinstance(restorer.pixel_loss, L1Loss)

    # prepare data
    inputs = torch.rand(1, 3, 22, 11)
    targets = torch.rand(1, 128 * 64, 3)
    coord = torch.rand(1, 128 * 64, 2)
    cell = torch.rand(1, 128 * 64, 2)
    data_batch = {'lq': inputs, 'gt': targets, 'coord': coord, 'cell': cell}

    # prepare optimizer
    optim_cfg = dict(type='Adam', lr=1e-4, betas=(0.9, 0.999))
    optimizer = obj_from_dict(optim_cfg, torch.optim,
                              dict(params=restorer.parameters()))

    # test train_step and forward_test (cpu)
    outputs = restorer.train_step(data_batch, optimizer)
    assert isinstance(outputs, dict)
    assert isinstance(outputs['log_vars'], dict)
    assert isinstance(outputs['log_vars']['loss_pix'], float)
    assert outputs['num_samples'] == 1
    assert outputs['results']['lq'].shape == data_batch['lq'].shape
    assert outputs['results']['gt'].shape == data_batch['gt'].shape
    assert torch.is_tensor(outputs['results']['output'])
    assert outputs['results']['output'].size() == (1, 128 * 64, 3)

    # test train_step and forward_test (gpu)
    if torch.cuda.is_available():
        restorer = restorer.cuda()
        data_batch = {
            'lq': inputs.cuda(),
            'gt': targets.cuda(),
            'coord': coord.cuda(),
            'cell': cell.cuda()
        }

        # train_step
        optimizer = obj_from_dict(optim_cfg, torch.optim,
                                  dict(params=restorer.parameters()))
        outputs = restorer.train_step(data_batch, optimizer)
        assert isinstance(outputs, dict)
        assert isinstance(outputs['log_vars'], dict)
        assert isinstance(outputs['log_vars']['loss_pix'], float)
        assert outputs['num_samples'] == 1
        assert outputs['results']['lq'].shape == data_batch['lq'].shape
        assert outputs['results']['gt'].shape == data_batch['gt'].shape
        assert torch.is_tensor(outputs['results']['output'])
        assert outputs['results']['output'].size() == (1, 128 * 64, 3)

        # val_step
        result = restorer.val_step(data_batch, meta=[{'gt_path': ''}])
        assert isinstance(result, dict)
        assert isinstance(result['eval_result'], dict)
        assert result['eval_result'].keys() == set({'PSNR', 'SSIM'})
        assert isinstance(result['eval_result']['PSNR'], np.float64)
        assert isinstance(result['eval_result']['SSIM'], np.float64)

[((12, 1, 12, 29), 'mmedit.models.registry.COMPONENTS.register_module', 'COMPONENTS.register_module', ({}, {}), '()', False, 'from mmedit.models.registry import COMPONENTS\n'), ((55, 15, 55, 77), 'mmedit.models.build_model', 'build_model', (), '', False, 'from mmedit.models import build_model\n'), ((63, 13, 63, 37), 'torch.rand', 'torch.rand', ({(63, 24, 63, 25): '1', (63, 27, 63, 28): '3', (63, 30, 63, 32): '22', (63, 34, 63, 36): '11'}, {}), '(1, 3, 22, 11)', False, 'import torch\n'), ((64, 14, 64, 40), 'torch.rand', 'torch.rand', ({(64, 25, 64, 26): '1', (64, 28, 64, 36): '128 * 64', (64, 38, 64, 39): '3'}, {}), '(1, 128 * 64, 3)', False, 'import torch\n'), ((65, 12, 65, 38), 'torch.rand', 'torch.rand', ({(65, 23, 65, 24): '1', (65, 26, 65, 34): '128 * 64', (65, 36, 65, 37): '2'}, {}), '(1, 128 * 64, 2)', False, 'import torch\n'), ((66, 11, 66, 37), 'torch.rand', 'torch.rand', ({(66, 22, 66, 23): '1', (66, 25, 66, 33): '128 * 64', (66, 35, 66, 36): '2'}, {}), '(1, 128 * 64, 2)', False, 'import torch\n'), ((82, 11, 82, 56), 'torch.is_tensor', 'torch.is_tensor', ({(82, 27, 82, 55): "outputs['results']['output']"}, {}), "(outputs['results']['output'])", False, 'import torch\n'), ((86, 7, 86, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((23, 21, 23, 47), 'torch.nn.Linear', 'nn.Linear', ({(23, 31, 23, 37): 'in_dim', (23, 39, 23, 46): 'out_dim'}, {}), '(in_dim, out_dim)', True, 'import torch.nn as nn\n'), ((105, 15, 105, 60), 'torch.is_tensor', 'torch.is_tensor', ({(105, 31, 105, 59): "outputs['results']['output']"}, {}), "(outputs['results']['output'])", False, 'import torch\n')]