#format python 
''' PDBAnalyzer script by destine, 2005-05-25
'''
import unittest, os, glob
from cStringIO import StringIO
from Bio.PDB import *

class PDBAnalyzer:
    def __init__(self, name, input_file):
        # first make a PDB parser object
        self.pdb_parser = PDBParser(PERMISSIVE=1)

        # get the structure, call it name
        self.structure = self.pdb_parser.get_structure(name, input_file)

    def getChains(self, model_index):
        return self.structure.get_list()[model_index]

    def getChain(self, model_index, chain_index):
        chains = self.getChains( model_index )
        return chains.get_list()[chain_index]

    def getAtomsCoordInChain(self, model_index, chain_index):
        chain = self.getChain( model_index, chain_index)
        coords = []
        for residue in chain.get_list():
            #print residue.get_list()
            for atom in residue.get_list():
                if atom.name == 'CA':
                    coords.append(atom.coord)
        return coords

class TestPDBAnalyzer(unittest.TestCase):
    def setUp(self):
        input_file = '1C3W.pdb'
        self.pa = PDBAnalyzer("test",input_file)

    def testPolypeptidesUsingC2N(self):
        ppb = PPBuilder()
        ppp = ppb.build_peptides(self.pa.structure, 1)
        self.assertEquals(2,len(ppp))

    def testAtomsCoordInChain(self):
        coords = self.pa.getAtomsCoordInChain(0,0)
        self.assertEquals(222, len(coords))

if __name__ == '__main__':
    unittest.main()