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, 4, 42, 54), 'numpy.testing.assert_almost_equal', 'np.testing.assert_almost_equal', ({(42, 35, 42, 49): "acc['p_mpjpe']", (42, 51, 42, 53): '(0.0)'}, {}), "(acc['p_mpjpe'], 0.0)", True, 'import numpy as np\n'), ((54, 11, 54, 40), 'numpy.random.random', 'np.random.random', ({(54, 28, 54, 39): 'input_shape'}, {}), '(input_shape)', True, 'import numpy as np\n'), ((55, 11, 55, 34), 'torch.FloatTensor', 'torch.FloatTensor', ({(55, 29, 55, 33): 'inps'}, {}), '(inps)', False, 'import torch\n'), ((17, 9, 17, 38), 'pytest.raises', 'pytest.raises', ({(17, 23, 17, 37): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((23, 9, 23, 38), 'pytest.raises', 'pytest.raises', ({(23, 23, 23, 37): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((32, 24, 32, 46), 'torch.Size', 'torch.Size', ({(32, 35, 32, 45): '[1, 17, 3]'}, {}), '([1, 17, 3])', False, 'import torch\n'), ((34, 35, 34, 55), 'torch.ones_like', 'torch.ones_like', ({(34, 51, 34, 54): 'out'}, {}), '(out)', False, 'import torch\n'), ((35, 44, 35, 60), 'torch.tensor', 'torch.tensor', ({(35, 57, 35, 59): '(0.0)'}, {}), '(0.0)', False, 'import torch\n'), ((40, 38, 40, 58), 'torch.ones_like', 'torch.ones_like', ({(40, 54, 40, 57): 'out'}, {}), '(out)', False, 'import torch\n')]