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)

[((257, 285), 'mmedit.models.registry.COMPONENTS.register_module', 'COMPONENTS.register_module', ([], {}), '()\n', (283, 285), False, 'from mmedit.models.registry import COMPONENTS\n'), ((1389, 1451), 'mmedit.models.build_model', 'build_model', (['model_cfg'], {'train_cfg': 'train_cfg', 'test_cfg': 'test_cfg'}), '(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)\n', (1400, 1451), False, 'from mmedit.models import build_model\n'), ((1650, 1674), 'torch.rand', 'torch.rand', (['(1)', '(3)', '(22)', '(11)'], {}), '(1, 3, 22, 11)\n', (1660, 1674), False, 'import torch\n'), ((1689, 1715), 'torch.rand', 'torch.rand', (['(1)', '(128 * 64)', '(3)'], {}), '(1, 128 * 64, 3)\n', (1699, 1715), False, 'import torch\n'), ((1728, 1754), 'torch.rand', 'torch.rand', (['(1)', '(128 * 64)', '(2)'], {}), '(1, 128 * 64, 2)\n', (1738, 1754), False, 'import torch\n'), ((1766, 1792), 'torch.rand', 'torch.rand', (['(1)', '(128 * 64)', '(2)'], {}), '(1, 128 * 64, 2)\n', (1776, 1792), False, 'import torch\n'), ((2515, 2560), 'torch.is_tensor', 'torch.is_tensor', (["outputs['results']['output']"], {}), "(outputs['results']['output'])\n", (2530, 2560), False, 'import torch\n'), ((2681, 2706), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (2704, 2706), False, 'import torch\n'), ((541, 567), 'torch.nn.Linear', 'nn.Linear', (['in_dim', 'out_dim'], {}), '(in_dim, out_dim)\n', (550, 567), True, 'import torch.nn as nn\n'), ((3483, 3528), 'torch.is_tensor', 'torch.is_tensor', (["outputs['results']['output']"], {}), "(outputs['results']['output'])\n", (3498, 3528), False, 'import torch\n')]