# 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/relax/amber_minimize.py
# ENERGY = unit.kilocalories_per_mole
# LENGTH = unit.angstroms
# def will_restrain(atom: openmm_app.Atom, rset: str) -> bool:
# def _add_restraints(
#     system: openmm.System,
#     reference_pdb: openmm_app.PDBFile,
#     stiffness: unit.Unit,
#     rset: str,
#     exclude_residues: Sequence[int]):
# def _openmm_minimize(
#     pdb_str: str,
#     max_iterations: int,
#     tolerance: unit.Unit,
#     stiffness: unit.Unit,
#     restraint_set: str,
#     exclude_residues: Sequence[int],
#     use_gpu: bool):
# def _get_pdb_string(topology: openmm_app.Topology, positions: unit.Quantity):
# def _check_cleaned_atoms(pdb_cleaned_string: str, pdb_ref_string: str):
# def _check_residues_are_well_defined(prot: protein.Protein):
# def _check_atom_mask_is_ideal(prot):
# def clean_protein(
#     prot: protein.Protein,
#     checks: bool = True):
# def make_atom14_positions(prot):
# def find_violations(prot_np: protein.Protein):
# def get_violation_metrics(prot: protein.Protein):
# def _run_one_iteration(
#     *,
#     pdb_string: str,
#     max_iterations: int,
#     tolerance: float,
#     stiffness: float,
#     restraint_set: str,
#     max_attempts: int,
#     use_gpu: bool,
#     exclude_residues: Optional[Collection[int]] = None):
# def run_pipeline(
#     prot: protein.Protein,
#     stiffness: float,
#     use_gpu: bool,
#     max_outer_iterations: int = 1,
#     place_hydrogens_every_iteration: bool = True,
#     max_iterations: int = 0,
#     tolerance: float = 2.39,
#     restraint_set: str = "non_hydrogen",
#     max_attempts: int = 100,
#     checks: bool = True,
#     exclude_residues: Optional[Sequence[int]] = None):
# def get_initial_energies(pdb_strs: Sequence[str],
#                          stiffness: float = 0.0,
#                          restraint_set: str = "non_hydrogen",
#                          exclude_residues: Optional[Sequence[int]] = None):
# 
# Path: alphafold/relax/utils.py
# def overwrite_pdb_coordinates(pdb_str: str, pos) -> str:
# def overwrite_b_factors(pdb_str: str, bfactors: np.ndarray) -> str:
# def assert_equal_nonterminal_atom_types(
#     atom_mask: np.ndarray, ref_atom_mask: np.ndarray):


from typing import Any, Dict, Sequence, Tuple
from alphafold.common import protein
from alphafold.relax import amber_minimize
from alphafold.relax import utils
import numpy as np
       slowdown. In the worst case we do 20 outer iterations.
      use_gpu: Whether to run on GPU.
    """

    self._max_iterations = max_iterations
    self._tolerance = tolerance
    self._stiffness = stiffness
    self._exclude_residues = exclude_residues
    self._max_outer_iterations = max_outer_iterations
    self._use_gpu = use_gpu

  def process(self, *,
              prot: protein.Protein) -> Tuple[str, Dict[str, Any], np.ndarray]:
    """Runs Amber relax on a prediction, adds hydrogens, returns PDB string."""
    out = amber_minimize.run_pipeline(
        prot=prot, max_iterations=self._max_iterations,
        tolerance=self._tolerance, stiffness=self._stiffness,
        exclude_residues=self._exclude_residues,
        max_outer_iterations=self._max_outer_iterations,
        use_gpu=self._use_gpu)
    min_pos = out['pos']
    start_pos = out['posinit']
    rmsd = np.sqrt(np.sum((start_pos - min_pos)**2) / start_pos.shape[0])
    debug_data = {
        'initial_energy': out['einit'],
        'final_energy': out['efinal'],
        'attempts': out['min_attempts'],
        'rmsd': rmsd
    }
    pdb_str = amber_minimize.clean_protein(prot)

# 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/relax/amber_minimize.py
# ENERGY = unit.kilocalories_per_mole
# LENGTH = unit.angstroms
# def will_restrain(atom: openmm_app.Atom, rset: str) -> bool:
# def _add_restraints(
#     system: openmm.System,
#     reference_pdb: openmm_app.PDBFile,
#     stiffness: unit.Unit,
#     rset: str,
#     exclude_residues: Sequence[int]):
# def _openmm_minimize(
#     pdb_str: str,
#     max_iterations: int,
#     tolerance: unit.Unit,
#     stiffness: unit.Unit,
#     restraint_set: str,
#     exclude_residues: Sequence[int],
#     use_gpu: bool):
# def _get_pdb_string(topology: openmm_app.Topology, positions: unit.Quantity):
# def _check_cleaned_atoms(pdb_cleaned_string: str, pdb_ref_string: str):
# def _check_residues_are_well_defined(prot: protein.Protein):
# def _check_atom_mask_is_ideal(prot):
# def clean_protein(
#     prot: protein.Protein,
#     checks: bool = True):
# def make_atom14_positions(prot):
# def find_violations(prot_np: protein.Protein):
# def get_violation_metrics(prot: protein.Protein):
# def _run_one_iteration(
#     *,
#     pdb_string: str,
#     max_iterations: int,
#     tolerance: float,
#     stiffness: float,
#     restraint_set: str,
#     max_attempts: int,
#     use_gpu: bool,
#     exclude_residues: Optional[Collection[int]] = None):
# def run_pipeline(
#     prot: protein.Protein,
#     stiffness: float,
#     use_gpu: bool,
#     max_outer_iterations: int = 1,
#     place_hydrogens_every_iteration: bool = True,
#     max_iterations: int = 0,
#     tolerance: float = 2.39,
#     restraint_set: str = "non_hydrogen",
#     max_attempts: int = 100,
#     checks: bool = True,
#     exclude_residues: Optional[Sequence[int]] = None):
# def get_initial_energies(pdb_strs: Sequence[str],
#                          stiffness: float = 0.0,
#                          restraint_set: str = "non_hydrogen",
#                          exclude_residues: Optional[Sequence[int]] = None):
# 
# Path: alphafold/relax/utils.py
# def overwrite_pdb_coordinates(pdb_str: str, pos) -> str:
# def overwrite_b_factors(pdb_str: str, bfactors: np.ndarray) -> str:
# def assert_equal_nonterminal_atom_types(
#     atom_mask: np.ndarray, ref_atom_mask: np.ndarray):


# Path: alphafold/relax/relax.py
from typing import Any, Dict, Sequence, Tuple
from alphafold.common import protein
from alphafold.relax import amber_minimize
from alphafold.relax import utils
import numpy as np

       slowdown. In the worst case we do 20 outer iterations.
      use_gpu: Whether to run on GPU.
    """

    self._max_iterations = max_iterations
    self._tolerance = tolerance
    self._stiffness = stiffness
    self._exclude_residues = exclude_residues
    self._max_outer_iterations = max_outer_iterations
    self._use_gpu = use_gpu

  def process(self, *,
              prot: protein.Protein) -> Tuple[str, Dict[str, Any], np.ndarray]:
    """Runs Amber relax on a prediction, adds hydrogens, returns PDB string."""
    out = amber_minimize.run_pipeline(
        prot=prot, max_iterations=self._max_iterations,
        tolerance=self._tolerance, stiffness=self._stiffness,
        exclude_residues=self._exclude_residues,
        max_outer_iterations=self._max_outer_iterations,
        use_gpu=self._use_gpu)
    min_pos = out['pos']
    start_pos = out['posinit']
    rmsd = np.sqrt(np.sum((start_pos - min_pos)**2) / start_pos.shape[0])
    debug_data = {
        'initial_energy': out['einit'],
        'final_energy': out['efinal'],
        'attempts': out['min_attempts'],
        'rmsd': rmsd
    }
    pdb_str = amber_minimize.clean_protein(prot)

    min_pdb = utils.overwrite_pdb_coordinates(pdb_str, min_pos)