# Path: alphafold/model/quat_affine.py
# QUAT_TO_ROT = np.zeros((4, 4, 3, 3), dtype=np.float32)
# QUAT_TO_ROT[0, 0] = [[ 1, 0, 0], [ 0, 1, 0], [ 0, 0, 1]]  # rr
# QUAT_TO_ROT[1, 1] = [[ 1, 0, 0], [ 0,-1, 0], [ 0, 0,-1]]  # ii
# QUAT_TO_ROT[2, 2] = [[-1, 0, 0], [ 0, 1, 0], [ 0, 0,-1]]  # jj
# QUAT_TO_ROT[3, 3] = [[-1, 0, 0], [ 0,-1, 0], [ 0, 0, 1]]  # kk
# QUAT_TO_ROT[1, 2] = [[ 0, 2, 0], [ 2, 0, 0], [ 0, 0, 0]]  # ij
# QUAT_TO_ROT[1, 3] = [[ 0, 0, 2], [ 0, 0, 0], [ 2, 0, 0]]  # ik
# QUAT_TO_ROT[2, 3] = [[ 0, 0, 0], [ 0, 0, 2], [ 0, 2, 0]]  # jk
# QUAT_TO_ROT[0, 1] = [[ 0, 0, 0], [ 0, 0,-2], [ 0, 2, 0]]  # ir
# QUAT_TO_ROT[0, 2] = [[ 0, 0, 2], [ 0, 0, 0], [-2, 0, 0]]  # jr
# QUAT_TO_ROT[0, 3] = [[ 0,-2, 0], [ 2, 0, 0], [ 0, 0, 0]]  # kr
# QUAT_MULTIPLY = np.zeros((4, 4, 4), dtype=np.float32)
# QUAT_MULTIPLY[:, :, 0] = [[ 1, 0, 0, 0],
#                           [ 0,-1, 0, 0],
#                           [ 0, 0,-1, 0],
#                           [ 0, 0, 0,-1]]
# QUAT_MULTIPLY[:, :, 1] = [[ 0, 1, 0, 0],
#                           [ 1, 0, 0, 0],
#                           [ 0, 0, 0, 1],
#                           [ 0, 0,-1, 0]]
# QUAT_MULTIPLY[:, :, 2] = [[ 0, 0, 1, 0],
#                           [ 0, 0, 0,-1],
#                           [ 1, 0, 0, 0],
#                           [ 0, 1, 0, 0]]
# QUAT_MULTIPLY[:, :, 3] = [[ 0, 0, 0, 1],
#                           [ 0, 0, 1, 0],
#                           [ 0,-1, 0, 0],
#                           [ 1, 0, 0, 0]]
# QUAT_MULTIPLY_BY_VEC = QUAT_MULTIPLY[:, 1:, :]
# def rot_to_quat(rot, unstack_inputs=False):
# def rot_list_to_tensor(rot_list):
# def vec_list_to_tensor(vec_list):
# def quat_to_rot(normalized_quat):
# def quat_multiply_by_vec(quat, vec):
# def quat_multiply(quat1, quat2):
# def apply_rot_to_vec(rot, vec, unstack=False):
# def apply_inverse_rot_to_vec(rot, vec):
#   def __init__(self, quaternion, translation, rotation=None, normalize=True,
#                unstack_inputs=False):
#   def to_tensor(self):
#   def apply_tensor_fn(self, tensor_fn):
#   def apply_rotation_tensor_fn(self, tensor_fn):
#   def scale_translation(self, position_scale):
#   def from_tensor(cls, tensor, normalize=False):
#   def pre_compose(self, update):
#   def apply_to_point(self, point, extra_dims=0):
#   def invert_point(self, transformed_point, extra_dims=0):
#   def __repr__(self):
# def _multiply(a, b):
# def make_canonical_transform(
#     n_xyz: jnp.ndarray,
#     ca_xyz: jnp.ndarray,
#     c_xyz: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def make_transform_from_reference(
#     n_xyz: jnp.ndarray,
#     ca_xyz: jnp.ndarray,
#     c_xyz: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
# class QuatAffine(object):


import collections
import jax.numpy as jnp
import tree
from typing import List
from alphafold.model import quat_affine
    origin: Vecs,  # shape (...)
    point_on_xy_plane: Vecs,  # shape (...)
) -> Rigids:  # shape (...)
  """Create Rigids from 3 points.

  Jumper et al. (2021) Suppl. Alg. 21 "rigidFrom3Points"
  This creates a set of rigid transformations from 3 points by Gram Schmidt
  orthogonalization.

  Args:
    point_on_neg_x_axis: Vecs corresponding to points on the negative x axis
    origin: Origin of resulting rigid transformations
    point_on_xy_plane: Vecs corresponding to points in the xy plane
  Returns:
    Rigid transformations from global frame to local frames derived from
    the input points.
  """
  m = rots_from_two_vecs(
      e0_unnormalized=vecs_sub(origin, point_on_neg_x_axis),
      e1_unnormalized=vecs_sub(point_on_xy_plane, origin))

  return Rigids(rot=m, trans=origin)


def rigids_from_list(l: List[jnp.ndarray]) -> Rigids:
  """Converts flat list of arrays to rigid transformations."""
  assert len(l) == 12
  return Rigids(Rots(*(l[:9])), Vecs(*(l[9:])))



# Path: alphafold/model/quat_affine.py
# QUAT_TO_ROT = np.zeros((4, 4, 3, 3), dtype=np.float32)
# QUAT_TO_ROT[0, 0] = [[ 1, 0, 0], [ 0, 1, 0], [ 0, 0, 1]]  # rr
# QUAT_TO_ROT[1, 1] = [[ 1, 0, 0], [ 0,-1, 0], [ 0, 0,-1]]  # ii
# QUAT_TO_ROT[2, 2] = [[-1, 0, 0], [ 0, 1, 0], [ 0, 0,-1]]  # jj
# QUAT_TO_ROT[3, 3] = [[-1, 0, 0], [ 0,-1, 0], [ 0, 0, 1]]  # kk
# QUAT_TO_ROT[1, 2] = [[ 0, 2, 0], [ 2, 0, 0], [ 0, 0, 0]]  # ij
# QUAT_TO_ROT[1, 3] = [[ 0, 0, 2], [ 0, 0, 0], [ 2, 0, 0]]  # ik
# QUAT_TO_ROT[2, 3] = [[ 0, 0, 0], [ 0, 0, 2], [ 0, 2, 0]]  # jk
# QUAT_TO_ROT[0, 1] = [[ 0, 0, 0], [ 0, 0,-2], [ 0, 2, 0]]  # ir
# QUAT_TO_ROT[0, 2] = [[ 0, 0, 2], [ 0, 0, 0], [-2, 0, 0]]  # jr
# QUAT_TO_ROT[0, 3] = [[ 0,-2, 0], [ 2, 0, 0], [ 0, 0, 0]]  # kr
# QUAT_MULTIPLY = np.zeros((4, 4, 4), dtype=np.float32)
# QUAT_MULTIPLY[:, :, 0] = [[ 1, 0, 0, 0],
#                           [ 0,-1, 0, 0],
#                           [ 0, 0,-1, 0],
#                           [ 0, 0, 0,-1]]
# QUAT_MULTIPLY[:, :, 1] = [[ 0, 1, 0, 0],
#                           [ 1, 0, 0, 0],
#                           [ 0, 0, 0, 1],
#                           [ 0, 0,-1, 0]]
# QUAT_MULTIPLY[:, :, 2] = [[ 0, 0, 1, 0],
#                           [ 0, 0, 0,-1],
#                           [ 1, 0, 0, 0],
#                           [ 0, 1, 0, 0]]
# QUAT_MULTIPLY[:, :, 3] = [[ 0, 0, 0, 1],
#                           [ 0, 0, 1, 0],
#                           [ 0,-1, 0, 0],
#                           [ 1, 0, 0, 0]]
# QUAT_MULTIPLY_BY_VEC = QUAT_MULTIPLY[:, 1:, :]
# def rot_to_quat(rot, unstack_inputs=False):
# def rot_list_to_tensor(rot_list):
# def vec_list_to_tensor(vec_list):
# def quat_to_rot(normalized_quat):
# def quat_multiply_by_vec(quat, vec):
# def quat_multiply(quat1, quat2):
# def apply_rot_to_vec(rot, vec, unstack=False):
# def apply_inverse_rot_to_vec(rot, vec):
#   def __init__(self, quaternion, translation, rotation=None, normalize=True,
#                unstack_inputs=False):
#   def to_tensor(self):
#   def apply_tensor_fn(self, tensor_fn):
#   def apply_rotation_tensor_fn(self, tensor_fn):
#   def scale_translation(self, position_scale):
#   def from_tensor(cls, tensor, normalize=False):
#   def pre_compose(self, update):
#   def apply_to_point(self, point, extra_dims=0):
#   def invert_point(self, transformed_point, extra_dims=0):
#   def __repr__(self):
# def _multiply(a, b):
# def make_canonical_transform(
#     n_xyz: jnp.ndarray,
#     ca_xyz: jnp.ndarray,
#     c_xyz: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def make_transform_from_reference(
#     n_xyz: jnp.ndarray,
#     ca_xyz: jnp.ndarray,
#     c_xyz: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
# class QuatAffine(object):


# Path: alphafold/model/r3.py
import collections
import jax.numpy as jnp
import tree
from typing import List
from alphafold.model import quat_affine

    origin: Vecs,  # shape (...)
    point_on_xy_plane: Vecs,  # shape (...)
) -> Rigids:  # shape (...)
  """Create Rigids from 3 points.

  Jumper et al. (2021) Suppl. Alg. 21 "rigidFrom3Points"
  This creates a set of rigid transformations from 3 points by Gram Schmidt
  orthogonalization.

  Args:
    point_on_neg_x_axis: Vecs corresponding to points on the negative x axis
    origin: Origin of resulting rigid transformations
    point_on_xy_plane: Vecs corresponding to points in the xy plane
  Returns:
    Rigid transformations from global frame to local frames derived from
    the input points.
  """
  m = rots_from_two_vecs(
      e0_unnormalized=vecs_sub(origin, point_on_neg_x_axis),
      e1_unnormalized=vecs_sub(point_on_xy_plane, origin))

  return Rigids(rot=m, trans=origin)


def rigids_from_list(l: List[jnp.ndarray]) -> Rigids:
  """Converts flat list of arrays to rigid transformations."""
  assert len(l) == 12
  return Rigids(Rots(*(l[:9])), Vecs(*(l[9:])))



def rigids_from_quataffine(a: quat_affine.QuatAffine) -> Rigids: