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', True, 'import numpy as np\n'), (54, 'numpy.random.random', 'np.random.random', True, 'import numpy as np\n'), (55, 'torch.FloatTensor', 'torch.FloatTensor', False, 'import torch\n'), (17, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (23, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (32, 'torch.Size', 'torch.Size', False, 'import torch\n'), (34, 'torch.ones_like', 'torch.ones_like', False, 'import torch\n'), (35, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (40, 'torch.ones_like', 'torch.ones_like', False, 'import torch\n')]