# 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
               exclude_residues: Sequence[int],
               max_outer_iterations: int,
               use_gpu: bool):
    """Initialize Amber Relaxer.

    Args:
      max_iterations: Maximum number of L-BFGS iterations. 0 means no max.
      tolerance: kcal/mol, the energy tolerance of L-BFGS.
      stiffness: kcal/mol A**2, spring constant of heavy atom restraining
        potential.
      exclude_residues: Residues to exclude from per-atom restraining.
        Zero-indexed.
      max_outer_iterations: Maximum number of violation-informed relax
       iterations. A value of 1 will run the non-iterative procedure used in
       CASP14. Use 20 so that >95% of the bad cases are relaxed. Relax finishes
       as soon as there are no violations, hence in most cases this causes no
       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."""

# 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

               exclude_residues: Sequence[int],
               max_outer_iterations: int,
               use_gpu: bool):
    """Initialize Amber Relaxer.

    Args:
      max_iterations: Maximum number of L-BFGS iterations. 0 means no max.
      tolerance: kcal/mol, the energy tolerance of L-BFGS.
      stiffness: kcal/mol A**2, spring constant of heavy atom restraining
        potential.
      exclude_residues: Residues to exclude from per-atom restraining.
        Zero-indexed.
      max_outer_iterations: Maximum number of violation-informed relax
       iterations. A value of 1 will run the non-iterative procedure used in
       CASP14. Use 20 so that >95% of the bad cases are relaxed. Relax finishes
       as soon as there are no violations, hence in most cases this causes no
       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(