# Path: alphafold/relax/cleanup.py
# def fix_pdb(pdbfile, alterations_info):
# def clean_structure(pdb_structure, alterations_info):
# def _remove_heterogens(fixer, alterations_info, keep_water):
# def _replace_met_se(pdb_structure, alterations_info):
# def _remove_chains_of_length_one(pdb_structure, alterations_info):


import io
from absl.testing import absltest
from alphafold.relax import cleanup
from simtk.openmm.app.internal import pdbstructure
  handle = io.StringIO(pdb_str)
  return pdbstructure.PdbStructure(handle)


def _lines_to_structure(pdb_lines):
  return _pdb_to_structure('\n'.join(pdb_lines))


class CleanupTest(absltest.TestCase):

  def test_missing_residues(self):
    pdb_lines = ['SEQRES   1 C    3  CYS GLY LEU',
                 'ATOM      1  N   CYS C   1     -12.262  20.115  60.959  1.00 '
                 '19.08           N',
                 'ATOM      2  CA  CYS C   1     -11.065  20.934  60.773  1.00 '
                 '17.23           C',
                 'ATOM      3  C   CYS C   1     -10.002  20.742  61.844  1.00 '
                 '15.38           C',
                 'ATOM      4  O   CYS C   1     -10.284  20.225  62.929  1.00 '
                 '16.04           O',
                 'ATOM      5  N   LEU C   3      -7.688  18.700  62.045  1.00 '
                 '14.75           N',
                 'ATOM      6  CA  LEU C   3      -7.256  17.320  62.234  1.00 '
                 '16.81           C',
                 'ATOM      7  C   LEU C   3      -6.380  16.864  61.070  1.00 '
                 '16.95           C',
                 'ATOM      8  O   LEU C   3      -6.551  17.332  59.947  1.00 '
                 '16.97           O']
    input_handle = io.StringIO('\n'.join(pdb_lines))
    alterations = {}

# Path: alphafold/relax/cleanup.py
# def fix_pdb(pdbfile, alterations_info):
# def clean_structure(pdb_structure, alterations_info):
# def _remove_heterogens(fixer, alterations_info, keep_water):
# def _replace_met_se(pdb_structure, alterations_info):
# def _remove_chains_of_length_one(pdb_structure, alterations_info):


# Path: alphafold/relax/cleanup_test.py
import io
from absl.testing import absltest
from alphafold.relax import cleanup
from simtk.openmm.app.internal import pdbstructure

  handle = io.StringIO(pdb_str)
  return pdbstructure.PdbStructure(handle)


def _lines_to_structure(pdb_lines):
  return _pdb_to_structure('\n'.join(pdb_lines))


class CleanupTest(absltest.TestCase):

  def test_missing_residues(self):
    pdb_lines = ['SEQRES   1 C    3  CYS GLY LEU',
                 'ATOM      1  N   CYS C   1     -12.262  20.115  60.959  1.00 '
                 '19.08           N',
                 'ATOM      2  CA  CYS C   1     -11.065  20.934  60.773  1.00 '
                 '17.23           C',
                 'ATOM      3  C   CYS C   1     -10.002  20.742  61.844  1.00 '
                 '15.38           C',
                 'ATOM      4  O   CYS C   1     -10.284  20.225  62.929  1.00 '
                 '16.04           O',
                 'ATOM      5  N   LEU C   3      -7.688  18.700  62.045  1.00 '
                 '14.75           N',
                 'ATOM      6  CA  LEU C   3      -7.256  17.320  62.234  1.00 '
                 '16.81           C',
                 'ATOM      7  C   LEU C   3      -6.380  16.864  61.070  1.00 '
                 '16.95           C',
                 'ATOM      8  O   LEU C   3      -6.551  17.332  59.947  1.00 '
                 '16.97           O']
    input_handle = io.StringIO('\n'.join(pdb_lines))
    alterations = {}

    result = cleanup.fix_pdb(input_handle, alterations)