# Copyright 2012 by Eric Talevich.  All rights reserved.
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

import unittest
import warnings

try:
    import numpy
    from numpy import dot  # Missing on PyPy's micronumpy
    del dot
    # We don't need this (?) but Bio.PDB imports it automatically :(
    from numpy.linalg import svd, det  # Missing in PyPy 2.0 numpypy
except ImportError:
    from Bio import MissingPythonDependencyError
    raise MissingPythonDependencyError(
        "Install NumPy if you want to use PDB formats with SeqIO.")

from Bio import SeqIO
from Bio import BiopythonWarning
from Bio.PDB.PDBExceptions import PDBConstructionWarning


class TestPdbSeqres(unittest.TestCase):
    def test_seqres_parse(self):
        """Parse a multi-chain PDB by SEQRES entries.

        Reference:
        http://www.rcsb.org/pdb/files/fasta.txt?structureIdList=2BEG
        """
        chains = list(SeqIO.parse('PDB/2BEG.pdb', 'pdb-seqres'))
        self.assertEqual(len(chains), 5)
        actual_seq = 'DAEFRHDSGYEVHHQKLVFFAEDVGSNKGAIIGLMVGGVVIA'
        for chain, chn_id in zip(chains, 'ABCDE'):
            self.assertEqual(chain.id, '2BEG:' + chn_id)
            self.assertEqual(chain.annotations['chain'], chn_id)
            self.assertEqual(str(chain.seq), actual_seq)

    def test_seqres_read(self):
        """Read a single-chain PDB by SEQRES entries.

        Reference:
        http://www.rcsb.org/pdb/files/fasta.txt?structureIdList=1A8O
        """
        chain = SeqIO.read('PDB/1A8O.pdb', 'pdb-seqres')
        self.assertEqual(chain.id, '1A8O:A')
        self.assertEqual(chain.annotations['chain'], 'A')
        self.assertEqual(str(chain.seq),
                         'MDIRQGPKEPFRDYVDRFYKTLRAEQASQEVKNWMTETLLVQNANPDCKTIL'
                         'KALGPGATLEEMMTACQG')

    def test_seqres_missing(self):
        """Parse a PDB with no SEQRES entries."""
        chains = list(SeqIO.parse('PDB/a_structure.pdb', 'pdb-seqres'))
        self.assertEqual(len(chains), 0)


class TestPdbAtom(unittest.TestCase):
    def test_atom_parse(self):
        """Parse a multi-chain PDB by ATOM entries.

        Reference:
        http://www.rcsb.org/pdb/files/fasta.txt?structureIdList=2BEG
        """
        chains = list(SeqIO.parse('PDB/2BEG.pdb', 'pdb-atom'))
        self.assertEqual(len(chains), 5)
        actual_seq = 'LVFFAEDVGSNKGAIIGLMVGGVVIA'
        for chain, chn_id in zip(chains, 'ABCDE'):
            self.assertEqual(chain.id, '2BEG:' + chn_id)
            self.assertEqual(chain.annotations['chain'], chn_id)
            self.assertEqual(str(chain.seq), actual_seq)

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", PDBConstructionWarning)
            chains = list(SeqIO.parse('PDB/2XHE.pdb', 'pdb-atom'))
        actual_seq = 'DRLSRLRQMAAENQXXXXXXXXXXXXXXXXXXXXXXXPEPFMADFFNRVK'\
                     'RIRDNIEDIEQAIEQVAQLHTESLVAVSKEDRDRLNEKLQDTMARISALG'\
                     'NKIRADLKQIEKENKRAQQEGTFEDGTVSTDLRIRQSQHSSLSRKFVKVM'\
                     'TRYNDVQAENKRRYGENVARQCRVVEPSLSDDAIQKVIEHGXXXXXXXXX'\
                     'XXXXXXXXNEIRDRHKDIQQLERSLLELHEMFTDMSTLVASQGEMIDRIE'\
                     'FSVEQSHNYV'
        self.assertEqual(str(chains[1].seq), actual_seq)

    def test_atom_read(self):
        """Read a single-chain PDB by ATOM entries.

        Reference:
        http://www.rcsb.org/pdb/files/fasta.txt?structureIdList=1A8O
        """
        chain = SeqIO.read('PDB/1A8O.pdb', 'pdb-atom')
        self.assertEqual(chain.id, '1A8O:A')
        self.assertEqual(chain.annotations['chain'], 'A')
        self.assertEqual(str(chain.seq),
                         'MDIRQGPKEPFRDYVDRFYKTLRAEQASQEVKNWMTETLLVQNANPDCKTIL'
                         'KALGPGATLEEMMTACQG')
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", PDBConstructionWarning)
            warnings.simplefilter("ignore", BiopythonWarning)
            chain = SeqIO.read('PDB/a_structure.pdb', 'pdb-atom')
        self.assertEqual(chain.id, '????:A')
        self.assertEqual(chain.annotations['chain'], 'A')
        self.assertEqual(str(chain.seq), 'E')

    def test_atom_noheader(self):
        """Parse a PDB with no HEADER line."""
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', PDBConstructionWarning)
            warnings.simplefilter("ignore", BiopythonWarning)
            chains = list(SeqIO.parse('PDB/1LCD.pdb', 'pdb-atom'))

        self.assertEqual(len(chains), 1)
        self.assertEqual(str(chains[0].seq), 'MKPVTLYDVAEYAGVSYQTVSRVVNQASHVSAKTREKVEAAMAELNYIPNR')


if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)
