# Path: alphafold/common/protein.py
# PDB_CHAIN_IDS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
# PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)  # := 62.
# class Protein:
#   def __post_init__(self):
# def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
# def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
# def to_pdb(prot: Protein) -> str:
# def ideal_atom_mask(prot: Protein) -> np.ndarray:
# def from_prediction(
#     features: FeatureDict,
#     result: ModelOutput,
#     b_factors: Optional[np.ndarray] = None,
#     remove_leading_feature_dimension: bool = True) -> Protein:
#   def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
# 
# 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()


import os
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from alphafold.common import protein
from alphafold.common import residue_constants

# Internal import (7716).

TEST_DATA_DIR = 'alphafold/common/testdata/'


class ProteinTest(parameterized.TestCase):

  def _check_shapes(self, prot, num_res):
    """Check that the processed shapes are correct."""
    num_atoms = residue_constants.atom_type_num
    self.assertEqual((num_res, num_atoms, 3), prot.atom_positions.shape)
    self.assertEqual((num_res,), prot.aatype.shape)
    self.assertEqual((num_res, num_atoms), prot.atom_mask.shape)
    self.assertEqual((num_res,), prot.residue_index.shape)
    self.assertEqual((num_res,), prot.chain_index.shape)
    self.assertEqual((num_res, num_atoms), prot.b_factors.shape)

  @parameterized.named_parameters(
      dict(testcase_name='chain_A',
           pdb_file='2rbg.pdb', chain_id='A', num_res=282, num_chains=1),
      dict(testcase_name='chain_B',
           pdb_file='2rbg.pdb', chain_id='B', num_res=282, num_chains=1),
      dict(testcase_name='multichain',
           pdb_file='2rbg.pdb', chain_id=None, num_res=564, num_chains=2))
  def test_from_pdb_str(self, pdb_file, chain_id, num_res, num_chains):
    pdb_file = os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
                            pdb_file)
    with open(pdb_file) as f:
      pdb_string = f.read()

# Path: alphafold/common/protein.py
# PDB_CHAIN_IDS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
# PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)  # := 62.
# class Protein:
#   def __post_init__(self):
# def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
# def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
# def to_pdb(prot: Protein) -> str:
# def ideal_atom_mask(prot: Protein) -> np.ndarray:
# def from_prediction(
#     features: FeatureDict,
#     result: ModelOutput,
#     b_factors: Optional[np.ndarray] = None,
#     remove_leading_feature_dimension: bool = True) -> Protein:
#   def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
# 
# 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/common/protein_test.py
import os
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from alphafold.common import protein
from alphafold.common import residue_constants


# Internal import (7716).

TEST_DATA_DIR = 'alphafold/common/testdata/'


class ProteinTest(parameterized.TestCase):

  def _check_shapes(self, prot, num_res):
    """Check that the processed shapes are correct."""
    num_atoms = residue_constants.atom_type_num
    self.assertEqual((num_res, num_atoms, 3), prot.atom_positions.shape)
    self.assertEqual((num_res,), prot.aatype.shape)
    self.assertEqual((num_res, num_atoms), prot.atom_mask.shape)
    self.assertEqual((num_res,), prot.residue_index.shape)
    self.assertEqual((num_res,), prot.chain_index.shape)
    self.assertEqual((num_res, num_atoms), prot.b_factors.shape)

  @parameterized.named_parameters(
      dict(testcase_name='chain_A',
           pdb_file='2rbg.pdb', chain_id='A', num_res=282, num_chains=1),
      dict(testcase_name='chain_B',
           pdb_file='2rbg.pdb', chain_id='B', num_res=282, num_chains=1),
      dict(testcase_name='multichain',
           pdb_file='2rbg.pdb', chain_id=None, num_res=564, num_chains=2))
  def test_from_pdb_str(self, pdb_file, chain_id, num_res, num_chains):
    pdb_file = os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
                            pdb_file)
    with open(pdb_file) as f:
      pdb_string = f.read()

    prot = protein.from_pdb_string(pdb_string, chain_id)