from dipy.align.transforms import regtransforms, Transform
import numpy as np
from numpy.testing import (assert_array_equal,
                           assert_array_almost_equal,
                           assert_equal,
                           assert_raises)


def test_number_of_parameters():
    expected_params = {('TRANSLATION', 2): 2,
                       ('TRANSLATION', 3): 3,
                       ('ROTATION', 2): 1,
                       ('ROTATION', 3): 3,
                       ('RIGID', 2): 3,
                       ('RIGID', 3): 6,
                       ('SCALING', 2): 1,
                       ('SCALING', 3): 1,
                       ('AFFINE', 2): 6,
                       ('AFFINE', 3): 12,
                       ('RIGIDSCALING', 2): 5,
                       ('RIGIDSCALING', 3): 9,
                       ('RIGIDISOSCALING', 2): 4,
                       ('RIGIDISOSCALING', 3): 7}

    for ttype, transform in regtransforms.items():
        assert_equal(
            transform.get_number_of_parameters(),
            expected_params[ttype])


def test_param_to_matrix_2d():
    rng = np.random.RandomState()
    # Test translation matrix 2D
    transform = regtransforms[('TRANSLATION', 2)]
    dx, dy = rng.uniform(size=(2,))
    theta = np.array([dx, dy])
    expected = np.array([[1, 0, dx], [0, 1, dy], [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_equal(actual, expected)

    # Test rotation matrix 2D
    transform = regtransforms[('ROTATION', 2)]
    angle = rng.uniform()
    theta = np.array([angle])
    ct = np.cos(angle)
    st = np.sin(angle)
    expected = np.array([[ct, -st, 0], [st, ct, 0], [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid matrix 2D
    transform = regtransforms[('RIGID', 2)]
    angle, dx, dy = rng.uniform(size=(3,))
    theta = np.array([angle, dx, dy])
    ct = np.cos(angle)
    st = np.sin(angle)
    expected = np.array([[ct, -st, dx], [st, ct, dy], [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test scaling matrix 2D
    transform = regtransforms[('SCALING', 2)]
    factor = rng.uniform()
    theta = np.array([factor])
    expected = np.array([[factor, 0, 0], [0, factor, 0], [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid isoscaling matrix 2D
    transform = regtransforms[('RIGIDISOSCALING', 2)]
    angle, dx, dy, factor = rng.uniform(size=(4,))
    theta = np.array([angle, dx, dy, factor])
    ct = np.cos(angle)
    st = np.sin(angle)
    expected = np.array([[ct*factor, -st*factor, dx],
                         [st*factor, ct*factor, dy],
                         [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid scaling matrix 2D
    transform = regtransforms[('RIGIDSCALING', 2)]
    angle, dx, dy, sx, sy = rng.uniform(size=(5,))
    theta = np.array([angle, dx, dy, sx, sy])
    ct = np.cos(angle)
    st = np.sin(angle)
    expected = np.array([[ct*sx, -st*sx, dx],
                         [st*sy, ct*sy, dy],
                         [0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test affine 2D
    transform = regtransforms[('AFFINE', 2)]
    theta = rng.uniform(size=(6,))
    expected = np.eye(3)
    expected[0, :] = theta[:3]
    expected[1, :] = theta[3:6]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Verify that ValueError is raised if incorrect number of parameters
    for transform in regtransforms.values():
        n = transform.get_number_of_parameters()
        # Set incorrect number of parameters
        theta = np.zeros(n + 1, dtype=np.float64)
        assert_raises(ValueError, transform.param_to_matrix, theta)


def test_param_to_matrix_3d():
    rng = np.random.RandomState()
    # Test translation matrix 3D
    transform = regtransforms[('TRANSLATION', 3)]
    dx, dy, dz = rng.uniform(size=(3,))
    theta = np.array([dx, dy, dz])
    expected = np.array([[1, 0, 0, dx],
                         [0, 1, 0, dy],
                         [0, 0, 1, dz],
                         [0, 0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_equal(actual, expected)

    # Create helpful rotation matrix function
    def get_rotation_matrix(alpha, beta, gamma):
        ca = np.cos(alpha)
        sa = np.sin(alpha)
        cb = np.cos(beta)
        sb = np.sin(beta)
        cc = np.cos(gamma)
        sc = np.sin(gamma)

        X = np.array([[1, 0, 0],
                      [0, ca, -sa],
                      [0, sa, ca]])
        Y = np.array([[cb, 0, sb],
                      [0, 1, 0],
                      [-sb, 0, cb]])
        Z = np.array([[cc, -sc, 0],
                      [sc, cc, 0],
                      [0, 0, 1]])

        return Z.dot(X.dot(Y))  # Apply in order: Y, X, Z (Y goes to the right)

    # Test rotation matrix 3D
    transform = regtransforms[('ROTATION', 3)]
    theta = rng.uniform(size=(3,))
    R = get_rotation_matrix(theta[0], theta[1], theta[2])
    expected = np.eye(4)
    expected[:3, :3] = R[:3, :3]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid matrix 3D
    transform = regtransforms[('RIGID', 3)]
    theta = rng.uniform(size=(6,))
    R = get_rotation_matrix(theta[0], theta[1], theta[2])
    expected = np.eye(4)
    expected[:3, :3] = R[:3, :3]
    expected[:3, 3] = theta[3:6]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid isoscaling matrix 3D
    transform = regtransforms[('RIGIDISOSCALING', 3)]
    theta = rng.uniform(size=(7,))
    R = get_rotation_matrix(theta[0], theta[1], theta[2])
    expected = np.eye(4)
    expected[:3, :3] = R[:3, :3] * theta[6]
    expected[:3, 3] = theta[3:6]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test rigid scaling matrix 3D
    transform = regtransforms[('RIGIDSCALING', 3)]
    theta = rng.uniform(size=(9,))
    R = get_rotation_matrix(theta[0], theta[1], theta[2])
    expected = np.eye(4)
    R[0, :3] *= theta[6]
    R[1, :3] *= theta[7]
    R[2, :3] *= theta[8]
    expected[:3, :3] = R[:3, :3]
    expected[:3, 3] = theta[3:6]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test scaling matrix 3D
    transform = regtransforms[('SCALING', 3)]
    factor = rng.uniform()
    theta = np.array([factor])
    expected = np.array([[factor, 0, 0, 0],
                         [0, factor, 0, 0],
                         [0, 0, factor, 0],
                         [0, 0, 0, 1]])
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Test affine 3D
    transform = regtransforms[('AFFINE', 3)]
    theta = rng.uniform(size=(12,))
    expected = np.eye(4)
    expected[0, :] = theta[:4]
    expected[1, :] = theta[4:8]
    expected[2, :] = theta[8:12]
    actual = transform.param_to_matrix(theta)
    assert_array_almost_equal(actual, expected)

    # Verify that ValueError is raised if incorrect number of parameters
    for transform in regtransforms.values():
        n = transform.get_number_of_parameters()
        # Set incorrect number of parameters
        theta = np.zeros(n + 1, dtype=np.float64)
        assert_raises(ValueError, transform.param_to_matrix, theta)


def test_identity_parameters():
    for transform in regtransforms.values():
        dim = transform.get_dim()
        theta = transform.get_identity_parameters()

        expected = np.eye(dim + 1)
        actual = transform.param_to_matrix(theta)
        assert_array_almost_equal(actual, expected)


def test_jacobian_functions():
    rng = np.random.RandomState()
    # Compare the analytical Jacobians with their numerical approximations
    h = 1e-8
    nsamples = 50

    for transform in regtransforms.values():
        n = transform.get_number_of_parameters()
        dim = transform.get_dim()

        expected = np.empty((dim, n))
        theta = rng.uniform(size=(n,))
        T = transform.param_to_matrix(theta)

        for j in range(nsamples):
            x = 255 * (rng.uniform(size=(dim,)) - 0.5)
            actual = transform.jacobian(theta, x)

            # Approximate with finite differences
            x_hom = np.ones(dim + 1)
            x_hom[:dim] = x[:]
            for i in range(n):
                dtheta = theta.copy()
                dtheta[i] += h
                dT = np.array(transform.param_to_matrix(dtheta))
                g = (dT - T).dot(x_hom) / h
                expected[:, i] = g[:dim]

            assert_array_almost_equal(actual, expected, decimal=5)

    # Test ValueError is raised when theta parameter doesn't have the right
    # length
    for transform in regtransforms.values():
        n = transform.get_number_of_parameters()
        # Wrong number of parameters
        theta = np.zeros(n + 1)
        x = np.zeros(dim)
        assert_raises(ValueError, transform.jacobian, theta, x)


def test_invalid_transform():
    # Note: users should not attempt to use the base class Transform:
    # they should get an instance of one of its derived classes from the
    # regtransforms dictionary (the base class is not contained there)
    # If for some reason the user instantiates it and attempts to use it,
    # however, it will raise exceptions when attempting to retrieve its
    # Jacobian, identity parameters or its matrix representation. It will
    # return -1 if queried about its dimension or number of parameters
    transform = Transform()
    theta = np.ndarray(3)
    x = np.ndarray(3)
    assert_raises(ValueError, transform.jacobian, theta, x)
    assert_raises(ValueError, transform.get_identity_parameters)
    assert_raises(ValueError, transform.param_to_matrix, theta)

    expected = -1
    actual = transform.get_number_of_parameters()
    assert_equal(actual, expected)

    actual = transform.get_dim()
    assert_equal(actual, expected)
