import numpy as np

from constants import CA_C_DIST, N_CA_DIST, N_CA_C_ANGLE


def rotation_matrix(angle, axis):
    """
    Args:
        angle: (n,)
        axis: 0=x, 1=y, 2=z
    Returns:
        (n, 3, 3)
    """
    n = len(angle)
    R = np.eye(3)[None, :, :].repeat(n, axis=0)

    axis = 2 - axis
    start = axis // 2
    step = axis % 2 + 1
    s = slice(start, start + step + 1, step)

    R[:, s, s] = np.array(
        [[np.cos(angle), (-1) ** (axis + 1) * np.sin(angle)],
         [(-1) ** axis * np.sin(angle), np.cos(angle)]]
    ).transpose(2, 0, 1)
    return R


def get_bb_transform(n_xyz, ca_xyz, c_xyz):
    """
    Compute translation and rotation of the canoncical backbone frame (triangle N-Ca-C) from a position with
    Ca at the origin, N on the x-axis and C in the xy-plane to the global position of the backbone frame

    Args:
        n_xyz: (n, 3)
        ca_xyz: (n, 3)
        c_xyz: (n, 3)

    Returns:
        quaternion represented as array of shape (n, 4)
        translation vector which is an array of shape (n, 3)
    """

    translation = ca_xyz
    n_xyz = n_xyz - translation
    c_xyz = c_xyz - translation

    # Find rotation matrix that aligns the coordinate systems
    #    rotate around y-axis to move N into the xy-plane
    theta_y = np.arctan2(n_xyz[:, 2], -n_xyz[:, 0])
    Ry = rotation_matrix(theta_y, 1)
    n_xyz = np.einsum('noi,ni->no', Ry.transpose(0, 2, 1), n_xyz)

    #    rotate around z-axis to move N onto the x-axis
    theta_z = np.arctan2(n_xyz[:, 1], n_xyz[:, 0])
    Rz = rotation_matrix(theta_z, 2)
    # n_xyz = np.einsum('noi,ni->no', Rz.transpose(0, 2, 1), n_xyz)

    #    rotate around x-axis to move C into the xy-plane
    c_xyz = np.einsum('noj,nji,ni->no', Rz.transpose(0, 2, 1),
                      Ry.transpose(0, 2, 1), c_xyz)
    theta_x = np.arctan2(c_xyz[:, 2], c_xyz[:, 1])
    Rx = rotation_matrix(theta_x, 0)

    # Final rotation matrix
    R = np.einsum('nok,nkj,nji->noi', Ry, Rz, Rx)

    # Convert to quaternion
    # q = w + i*u_x + j*u_y + k * u_z
    quaternion = rotation_matrix_to_quaternion(R)

    return quaternion, translation


def get_bb_coords_from_transform(ca_coords, quaternion):
    """
    Args:
        ca_coords: (n, 3)
        quaternion: (n, 4)
    Returns:
        backbone coords (n*3, 3), order is [N, CA, C]
        backbone atom types as a list of length n*3
    """
    R = quaternion_to_rotation_matrix(quaternion)
    bb_coords = np.tile(np.array(
        [[N_CA_DIST, 0, 0],
         [0, 0, 0],
         [CA_C_DIST * np.cos(N_CA_C_ANGLE), CA_C_DIST * np.sin(N_CA_C_ANGLE), 0]]),
        [len(ca_coords), 1])
    bb_coords = np.einsum('noi,ni->no', R.repeat(3, axis=0), bb_coords) + ca_coords.repeat(3, axis=0)
    bb_atom_types = [t for _ in range(len(ca_coords)) for t in ['N', 'C', 'C']]

    return bb_coords, bb_atom_types


def quaternion_to_rotation_matrix(q):
    """
    x_rot = R x

    Args:
        q: (n, 4)
    Returns:
        R: (n, 3, 3)
    """
    # Normalize
    q = q / (q ** 2).sum(1, keepdims=True) ** 0.5

    # https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion
    w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
    R = np.stack([
        np.stack([1 - 2 * y ** 2 - 2 * z ** 2, 2 * x * y - 2 * z * w,
                  2 * x * z + 2 * y * w], axis=1),
        np.stack([2 * x * y + 2 * z * w, 1 - 2 * x ** 2 - 2 * z ** 2,
                  2 * y * z - 2 * x * w], axis=1),
        np.stack([2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w,
                  1 - 2 * x ** 2 - 2 * y ** 2], axis=1)
    ], axis=1)

    return R


def rotation_matrix_to_quaternion(R):
    """
    https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion
    Args:
        R: (n, 3, 3)
    Returns:
        q: (n, 4)
    """

    t = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]
    r = np.sqrt(1 + t)
    w = 0.5 * r
    x = np.sign(R[:, 2, 1] - R[:, 1, 2]) * np.abs(
        0.5 * np.sqrt(1 + R[:, 0, 0] - R[:, 1, 1] - R[:, 2, 2]))
    y = np.sign(R[:, 0, 2] - R[:, 2, 0]) * np.abs(
        0.5 * np.sqrt(1 - R[:, 0, 0] + R[:, 1, 1] - R[:, 2, 2]))
    z = np.sign(R[:, 1, 0] - R[:, 0, 1]) * np.abs(
        0.5 * np.sqrt(1 - R[:, 0, 0] - R[:, 1, 1] + R[:, 2, 2]))

    return np.stack((w, x, y, z), axis=1)