import numpy as np
import pytest
import torch

from mmpose.models import TemporalRegressionHead


def test_temporal_regression_head():
    """Test temporal head."""
    head = TemporalRegressionHead(
        in_channels=1024,
        num_joints=17,
        loss_keypoint=dict(type='MPJPELoss', use_target_weight=True))

    head.init_weights()

    with pytest.raises(AssertionError):
        # ndim of the input tensor should be 3
        input_shape = (1, 1024, 1, 1)
        inputs = _demo_inputs(input_shape)
        _ = head(inputs)

    with pytest.raises(AssertionError):
        # size of the last dim should be 1
        input_shape = (1, 1024, 3)
        inputs = _demo_inputs(input_shape)
        _ = head(inputs)

    input_shape = (1, 1024, 1)
    inputs = _demo_inputs(input_shape)
    out = head(inputs)
    assert out.shape == torch.Size([1, 17, 3])

    loss = head.get_loss(out, out, torch.ones_like(out))
    assert torch.allclose(loss['reg_loss'], torch.tensor(0.))

    _ = head.inference_model(inputs)
    _ = head.inference_model(inputs, [(0, 1), (2, 3)])

    acc = head.get_accuracy(out, out, torch.ones_like(out))
    assert acc['mpjpe'] == 0.
    np.testing.assert_almost_equal(acc['p_mpjpe'], 0.)


def _demo_inputs(input_shape=(1, 1024, 1)):
    """Create a superset of inputs needed to run head.

    Args:
        input_shape (tuple): input batch dimensions.
            Default: (1, 1024, 1).
    Returns:
        Random input tensor with the size of input_shape.
    """
    inps = np.random.random(input_shape)
    inps = torch.FloatTensor(inps)
    return inps

[(42, 'numpy.testing.assert_almost_equal', 'np.testing.assert_almost_equal', (["acc['p_mpjpe']", '(0.0)'], {}), True, 'import numpy as np\n'), (54, 'numpy.random.random', 'np.random.random', (['input_shape'], {}), True, 'import numpy as np\n'), (55, 'torch.FloatTensor', 'torch.FloatTensor', (['inps'], {}), False, 'import torch\n'), (17, 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), False, 'import pytest\n'), (23, 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), False, 'import pytest\n'), (32, 'torch.Size', 'torch.Size', (['[1, 17, 3]'], {}), False, 'import torch\n'), (34, 'torch.ones_like', 'torch.ones_like', (['out'], {}), False, 'import torch\n'), (35, 'torch.tensor', 'torch.tensor', (['(0.0)'], {}), False, 'import torch\n'), (40, 'torch.ones_like', 'torch.ones_like', (['out'], {}), False, 'import torch\n')]