# Path: alphafold/common/residue_constants.py
# def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]],
#   def make_bond_key(atom1_name, atom2_name):
# def sequence_to_onehot(
#     sequence: str,
#     mapping: Mapping[str, int],
#     map_unknown_to_x: bool = False) -> np.ndarray:
# def _make_standard_atom_mask() -> np.ndarray:
# def chi_angle_atom(atom_index: int) -> np.ndarray:
# def _make_rigid_transformation_4x4(ex, ey, translation):
# def _make_rigid_group_constants():
# def make_atom14_dists_bounds(overlap_tolerance=1.5,
#                              bond_length_tolerance_factor=15):
# HHBLITS_AA_TO_ID = {
#     'A': 0,
#     'B': 2,
#     'C': 1,
#     'D': 2,
#     'E': 3,
#     'F': 4,
#     'G': 5,
#     'H': 6,
#     'I': 7,
#     'J': 20,
#     'K': 8,
#     'L': 9,
#     'M': 10,
#     'N': 11,
#     'O': 20,
#     'P': 12,
#     'Q': 13,
#     'R': 14,
#     'S': 15,
#     'T': 16,
#     'U': 1,
#     'V': 17,
#     'W': 18,
#     'X': 20,
#     'Y': 19,
#     'Z': 3,
#     '-': 21,
# }
# ID_TO_HHBLITS_AA = {
#     0: 'A',
#     1: 'C',  # Also U.
#     2: 'D',  # Also B.
#     3: 'E',  # Also Z.
#     4: 'F',
#     5: 'G',
#     6: 'H',
#     7: 'I',
#     8: 'K',
#     9: 'L',
#     10: 'M',
#     11: 'N',
#     12: 'P',
#     13: 'Q',
#     14: 'R',
#     15: 'S',
#     16: 'T',
#     17: 'V',
#     18: 'W',
#     19: 'Y',
#     20: 'X',  # Includes J and O.
#     21: '-',
# }
# MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
#     restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
#     for i in range(len(restypes_with_x_and_gap)))
# STANDARD_ATOM_MASK = _make_standard_atom_mask()
# 
# Path: alphafold/model/tf/shape_helpers.py
# def shape_list(x):
# 
# Path: alphafold/model/tf/shape_placeholders.py
# NUM_RES = 'num residues placeholder'
# NUM_MSA_SEQ = 'msa placeholder'
# NUM_EXTRA_SEQ = 'extra msa placeholder'
# NUM_TEMPLATES = 'num templates placeholder'
# 
# Path: alphafold/model/tf/utils.py
# def tf_combine_mask(*masks):
#   def __init__(self, initial_seed=0):
#   def __call__(self):
# def make_random_seed():
# class SeedMaker(object):


from alphafold.common import residue_constants
from alphafold.model.tf import shape_helpers
from alphafold.model.tf import shape_placeholders
from alphafold.model.tf import utils
import numpy as np
import tensorflow.compat.v1 as tf
  for k in protein:
    if 'profile' in k:  # Include both hhblits and psiblast profiles
      num_dim = protein[k].shape.as_list()[-1]
      assert num_dim in [20, 21, 22], (
          'num_dim for %s out of expected range: %s' % (k, num_dim))
      protein[k] = tf.tensordot(protein[k], perm_matrix[:num_dim, :num_dim], 1)
  return protein


def squeeze_features(protein):
  """Remove singleton and repeated dimensions in protein features."""
  protein['aatype'] = tf.argmax(
      protein['aatype'], axis=-1, output_type=tf.int32)
  for k in [
      'domain_name', 'msa', 'num_alignments', 'seq_length', 'sequence',
      'superfamily', 'deletion_matrix', 'resolution',
      'between_segment_residues', 'residue_index', 'template_all_atom_masks']:
    if k in protein:
      final_dim = shape_helpers.shape_list(protein[k])[-1]
      if isinstance(final_dim, int) and final_dim == 1:
        protein[k] = tf.squeeze(protein[k], axis=-1)

  for k in ['seq_length', 'num_alignments']:
    if k in protein:
      protein[k] = protein[k][0]  # Remove fake sequence dimension
  return protein


def make_random_crop_to_size_seed(protein):
  """Random seed for cropping residues and templates."""

# Path: alphafold/common/residue_constants.py
# def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]],
#   def make_bond_key(atom1_name, atom2_name):
# def sequence_to_onehot(
#     sequence: str,
#     mapping: Mapping[str, int],
#     map_unknown_to_x: bool = False) -> np.ndarray:
# def _make_standard_atom_mask() -> np.ndarray:
# def chi_angle_atom(atom_index: int) -> np.ndarray:
# def _make_rigid_transformation_4x4(ex, ey, translation):
# def _make_rigid_group_constants():
# def make_atom14_dists_bounds(overlap_tolerance=1.5,
#                              bond_length_tolerance_factor=15):
# HHBLITS_AA_TO_ID = {
#     'A': 0,
#     'B': 2,
#     'C': 1,
#     'D': 2,
#     'E': 3,
#     'F': 4,
#     'G': 5,
#     'H': 6,
#     'I': 7,
#     'J': 20,
#     'K': 8,
#     'L': 9,
#     'M': 10,
#     'N': 11,
#     'O': 20,
#     'P': 12,
#     'Q': 13,
#     'R': 14,
#     'S': 15,
#     'T': 16,
#     'U': 1,
#     'V': 17,
#     'W': 18,
#     'X': 20,
#     'Y': 19,
#     'Z': 3,
#     '-': 21,
# }
# ID_TO_HHBLITS_AA = {
#     0: 'A',
#     1: 'C',  # Also U.
#     2: 'D',  # Also B.
#     3: 'E',  # Also Z.
#     4: 'F',
#     5: 'G',
#     6: 'H',
#     7: 'I',
#     8: 'K',
#     9: 'L',
#     10: 'M',
#     11: 'N',
#     12: 'P',
#     13: 'Q',
#     14: 'R',
#     15: 'S',
#     16: 'T',
#     17: 'V',
#     18: 'W',
#     19: 'Y',
#     20: 'X',  # Includes J and O.
#     21: '-',
# }
# MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
#     restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
#     for i in range(len(restypes_with_x_and_gap)))
# STANDARD_ATOM_MASK = _make_standard_atom_mask()
# 
# Path: alphafold/model/tf/shape_helpers.py
# def shape_list(x):
# 
# Path: alphafold/model/tf/shape_placeholders.py
# NUM_RES = 'num residues placeholder'
# NUM_MSA_SEQ = 'msa placeholder'
# NUM_EXTRA_SEQ = 'extra msa placeholder'
# NUM_TEMPLATES = 'num templates placeholder'
# 
# Path: alphafold/model/tf/utils.py
# def tf_combine_mask(*masks):
#   def __init__(self, initial_seed=0):
#   def __call__(self):
# def make_random_seed():
# class SeedMaker(object):


# Path: alphafold/model/tf/data_transforms.py
from alphafold.common import residue_constants
from alphafold.model.tf import shape_helpers
from alphafold.model.tf import shape_placeholders
from alphafold.model.tf import utils
import numpy as np
import tensorflow.compat.v1 as tf

  for k in protein:
    if 'profile' in k:  # Include both hhblits and psiblast profiles
      num_dim = protein[k].shape.as_list()[-1]
      assert num_dim in [20, 21, 22], (
          'num_dim for %s out of expected range: %s' % (k, num_dim))
      protein[k] = tf.tensordot(protein[k], perm_matrix[:num_dim, :num_dim], 1)
  return protein


def squeeze_features(protein):
  """Remove singleton and repeated dimensions in protein features."""
  protein['aatype'] = tf.argmax(
      protein['aatype'], axis=-1, output_type=tf.int32)
  for k in [
      'domain_name', 'msa', 'num_alignments', 'seq_length', 'sequence',
      'superfamily', 'deletion_matrix', 'resolution',
      'between_segment_residues', 'residue_index', 'template_all_atom_masks']:
    if k in protein:
      final_dim = shape_helpers.shape_list(protein[k])[-1]
      if isinstance(final_dim, int) and final_dim == 1:
        protein[k] = tf.squeeze(protein[k], axis=-1)

  for k in ['seq_length', 'num_alignments']:
    if k in protein:
      protein[k] = protein[k][0]  # Remove fake sequence dimension
  return protein


def make_random_crop_to_size_seed(protein):
  """Random seed for cropping residues and templates."""

  protein['random_crop_to_size_seed'] = utils.make_random_seed()