<commit_before>"""
Object representations of 3D and secondary structures. 
"""

import os
import re
import copy
import math
import numpy
import csb.pyutils
import csb.math
import csb.bio.utils

from csb.bio.sequence import SequenceTypes, SequenceAlphabets, AlignmentTypes
from itertools import izip


AngleUnits = csb.pyutils.enum( Degrees='deg', Radians='rad' )
"""
Torsion angle unit types
"""

SecStructures = csb.pyutils.enum( Helix='H', Strand='E', Coil='C', Turn='T',
                                  Bend='S', Helix3='G', PiHelix='I', BetaBridge='B', Gap='-' )
"""
Secondary structure types
"""

ChemElements = csb.pyutils.enum( H=1, He=2, Li=3, Be=4, B=5, C=6, N=7, O=8, F=9, Ne=10, Na=11, Mg=12, Al=13, Si=14, P=15, 
                                 S=16, Cl=17, Ar=18, K=19, Ca=20, Sc=21, Ti=22, V=23, Cr=24, Mn=25, Fe=26, Co=27, Ni=28, 
                                 Cu=29, Zn=30, Ga=31, Ge=32, As=33, Se=34, Br=35, Kr=36, Rb=37, Sr=38, Y=39, Zr=40, Nb=41, 
                                 Mo=42, Tc=43, Ru=44, Rh=45, Pd=46, Ag=47, Cd=48, In=49, Sn=50, Sb=51, Te=52, I=53, Xe=54,
                                 Cs=55, Ba=56, Hf=72, Ta=73, W=74, Re=75, Os=76, Ir=77, Pt=78, Au=79, Hg=80, Tl=81, Pb=82, 
                                 Bi=83, Po=84, At=85, Rn=86, Fr=87, Ra=88, Rf=104, Db=105, Sg=106, Bh=107, Hs=108, Mt=109, 
                                 Ds=110, Rg=111, La=57, Ce=58, Pr=59, Nd=60, Pm=61, Sm=62, Eu=63, Gd=64, Tb=65, Dy=66, 
                                 Ho=67, Er=68, Tm=69, Yb=70, Lu=71, Ac=89, Th=90, Pa=91, U=92, Np=93, Pu=94, Am=95, Cm=96, 
                                 Bk=97, Cf=98, Es=99, Fm=100, Md=101, No=102, Lr=103,
                                 x=-1 ) 
"""
Periodic table elements
"""

class Broken3DStructureError(ValueError):
    pass
class Missing3DStructureError(Broken3DStructureError):
    pass       
class InvalidOperation(Exception):
    pass
class DuplicateChainIDError(csb.pyutils.DuplicateKeyError):
    pass
class DuplicateResidueIDError(csb.pyutils.DuplicateKeyError):
    pass
class DuplicateAtomIDError(csb.pyutils.DuplicateKeyError):
    pass
class AlignmentArgumentLengthError(ValueError):
    pass
class BrokenSecStructureError(ValueError):
    pass

class Ensemble(csb.pyutils.AbstractNIContainer):
    """
    Represents an ensemble of multiple L{Structure} models.
    Provides a list-like access to these models::
    
        >>> ensemble[0]
        <Structure Model 1: accn, x chains>
        >>> ensemble.models[1]
        <Structure Model 1: accn, x chains>
    """
    
    def __init__(self):
        self._models = EnsembleModelsCollection()
        
    def __repr__(self):
        return "<Ensemble: {0} models>".format(self.models.length)        
        
    @property
    def _children(self):
        return self._models
    
    @property
    def models(self):
        return self._models
        
    @property
    def first_model(self):
        if len(self._models) > 0:
            return self[0]
        return None
        
class EnsembleModelsCollection(csb.pyutils.CollectionContainer):
    
    def __init__(self):
        super(EnsembleModelsCollection, self).__init__(type=Structure, start_index=1)

class Structure(csb.pyutils.AbstractNIContainer):
    """
    Represents a single model of a PDB 3-Dimensional molecular structure.
    Provides access to the L{Chain} objects, contained in the model::
    
        >>> structure['A']
        <Chain A: Protein>
        >>> structure.chains['A']
        <Chain A: Protein>
        >>> structure.items[0]
        <Chain A: Protein> #(given a chain order A, B...)    
    
    @param accession: accession number of the structure
    @type accession: str
    """
    def __init__(self, accession):
                        
        self._accession = None
        self._chains = StructureChainsTable(self)
        self._chainslist = StructureChainsCollection(self)
        self.model_id = None
        
        self.accession = accession

    def __repr__(self):
        return "<Structure Model {0.model_id}: {0.accession}, {1} chains>".format(self, self.chains.length)

    @property
    def _children(self):
        return self._chains
    
    @property
    def chains(self):
        return self._chains
    
    @property
    def items(self):
        return self._chainslist
    
    @property
    def first_chain(self):
        if len(self._chains) > 0:
            return self._chainslist[0]
        return None
        
    @property
    def accession(self):
        return self._accession
    @accession.setter
    def accession(self, accession):
        if accession is None:
            raise ValueError(accession)
        self._accession = str(accession).strip().lower()
        for c in self.chains:
            self.chains[c]._accession = self._accession
            
    def define_molecule(self, molecule_id, chains):
        """
        Group C{chains}, defined by their IDs, to a polymer with id = C{molecule_id}.
        
        @param molecule_id: ID of the new polymer
        @param chains: a list of chain IDs
        @type chains: list
        
        @raise ValueError: if an invalid chain ID is provided
        """
        assert molecule_id is not None
        cs = [ ]
        
        for c in chains:
            if c not in self.chains:
                raise ValueError('No such chain {0} in structure'.format(c))
            cs.append(self.chains[c])
        for chain in cs:
            chain.molecule_id = molecule_id

    def apply_transformation(self, rotation, translation):
        """
        Apply in place RotationMatrix and translation Vector to all atoms in the structure.
        
        @type rotation: numpy array
        @type translation: numpy array 
        """
        for chain_id in self.chains:
            self.chains[chain_id].apply_transformation(rotation, translation)
                    
    def to_fasta(self):
        """
        Dump the structure in FASTA format. 
        
        @return: FASTA-formatted string with all chains in the structure
        @rtype: str
        """
        fasta = []
        
        for chain_id in self.chains:

            chain = self.chains[chain_id]
            if chain.length > 0:
                fasta.append(chain.header)
                fasta.append(chain.sequence)
        
        return os.linesep.join(fasta)

    def to_pdb(self, output_file=None):
        """
        Dump the whole structure in PDB format.
        
        @param output_file: output file name
        @type output_file: str 
        """
        from datetime import datetime
        import StringIO
        
        def isnull(this, that, null=None):
            if this is null: return that
            return this
            
        class MyStringIO(StringIO.StringIO):
            def writeline(self, data):
                self.write('{0:80}\n'.format(data))
        
        stream = MyStringIO()
        
        stream.writeline('HEADER    {0:40}{1:%d-%b-%y}   {2:4}'.format('.', datetime.now(), self.accession.upper()))
        
        molecules = { }
        for chain_id in self.chains:
            chain = self.chains[chain_id]
            if chain.molecule_id not in molecules:
                molecules[chain.molecule_id] = [ ]
            molecules[chain.molecule_id].append(chain_id)
        
        k = 0
        for mol_id in sorted(molecules):
            
            chains = molecules[mol_id]
            first_chain = self.chains[ chains[0] ]            
            
            stream.writeline('COMPND {0:3} MOL_ID: {1};'.format(k + 1, isnull(mol_id, '0')))
            stream.writeline('COMPND {0:3} MOLECULE: {1};'.format(k + 2, isnull(first_chain.name, '')))
            stream.writeline('COMPND {0:3} CHAIN: {1};'.format(k + 3, ', '.join(chains)))
            k += 3
            
        for chain_id in self.chains:
            
            chain = self.chains[chain_id]
            res = [ chain.format_residue(r) for r in chain.residues ]

            rn = 0
            for j in range(0, chain.length, 13):
                rn += 1
                residues = [ '{0:>3}'.format(r) for r in res[j : j + 13] ]
                stream.writeline('SEQRES {0:>3} {1} {2:>4}  {3}'.format(
                                            rn, chain.id, chain.length, ' '.join(residues) ))
        
        for chain_id in self.chains:
        
            chain = self.chains[chain_id]
            for residue in chain.residues:
        
                atoms = [ ]
                for an in residue.atoms:
                    atom = residue.atoms[an]
                    if type(atom) is DisorderedAtom:
                        for dis_atom in atom: atoms.append(dis_atom)
                    else:
                        atoms.append(atom)
                atoms.sort()
                
                for atom in atoms:

                    alt = atom.alternate
                    if alt is True:
                        alt = 'A'
                    elif alt is False:
                        alt = ' '
                    
                    if atom.element:
                        element = repr(atom.element)
                    else:
                        element = ' '
                    stream.writeline('ATOM  {0:>5} {1:>4}{2}{3:>3} {4}{5:>4}{6}   {7:>8.3f}{8:>8.3f}{9:>8.3f}{10:>6.2f}{11:>6.2f}{12:>12}{13:2}'.format(
                                        atom.serial_number, atom._full_name, isnull(alt, ' '), 
                                        chain.format_residue(residue), chain.id, 
                                        isnull(residue.sequence_number, residue.rank), isnull(residue.insertion_code, ' '), 
                                        atom.vector[0], atom.vector[1], atom.vector[2], isnull(atom.occupancy, 0.0), isnull(atom.temperature, 0.0), 
                                        element, isnull(atom.charge, ' ') ))
            stream.writeline('TER')
        stream.writeline('END')
        
        data = stream.getvalue()        
        stream.close()
        
        if not output_file:
            return data
        else:
            with open(output_file, 'w') as out:
                out.write(data)                

class StructureChainsCollection(csb.pyutils.AbstractContainer):
    
    def __init__(self, container):
        self.__container = container
        
    @property
    def _children(self):
        return [self.__container[c] for c in self.__container]

class StructureChainsTable(csb.pyutils.DictionaryContainer):
    
    def __init__(self, structure=None, chains=None):
        self.__container = structure
        super(StructureChainsTable, self).__init__()
        
        if chains is not None:
            for chain in chains:
                self.append(chain)
        
    def __repr__(self):
        if len(self) > 0:
            return "<StructureChains: {0}>".format(', '.join(self.keys()))
        else:
            return "<StructureChains: empty>"     
    
    def append(self, chain):
        """
        Add a new Chain to the structure.
        
        @param chain: the new chain to be appended
        @type chain: L{Chain}
        
        @raise DuplicateChainIDError: if a chain with same ID is already defined
        @raise InvalidOperation: if the chain is already associated with a structure
        """
        
        if chain._structure and chain._structure is not self.__container:
            raise InvalidOperation('This chain is already part of another structure')
        if chain.id in self:
            raise DuplicateChainIDError('A chain with ID {0} is already defined'.format(chain.id))
            
        super(StructureChainsTable, self).append(chain.id, chain)
        
        if self.__container:
            chain._accession = self.__container.accession
            chain._structure = self.__container

    def remove(self, id):
        """
        Remove a chain from the structure.

        @param id: ID of the chain to be detached
        @type id: str
        """
        chain = self[id]
        self._remove(id)
        chain._structure = None   
    
    def _update_chain_id(self, chain, new_id):
        
        if chain.id not in self or self[chain.id] is not chain:
            raise InvalidOperation(chain)
        
        self._remove(chain.id)
        
        if new_id in self:
            raise DuplicateChainIDError('Chain ID {0} is already defined'.format(id))
        
        super(StructureChainsTable, self).append(new_id, chain)
        
class Chain(csb.pyutils.AbstractNIContainer):
    """
    Represents a polymeric chain. Provides list-like and rank-based access to
    the residues in the chain::
    
        >>> chain[0]
        <ProteinResidue [1]: SER None>
        >>> chain.residues[1]
        <ProteinResidue [1]: SER None>
    
    You can also access residues by their PDB sequence number:
    
        >>> chain.find(sequence_number=5, insertion_code='A')
        <ProteinResidue [1]: SER 5A>
    
    @param chain_id: ID of the new chain
    @type chain_id: str
    @param type: sequence type (a member of the L{SequenceTypes} enum)
    @type type: L{csb.pyutils.EnumItem}
    @param name: name of the chain
    @type name: str
    @param residues: initialization list of L{Residue}-s
    @type residues: list
    @param accession: accession number of the chain
    @type accession: str
    @param molecule_id: MOL ID of the chain, if part of a polymer
    
    """
    def __init__(self, chain_id, type=SequenceTypes.Protein, name='',           #@UndefinedVariable
                 residues=None, accession=None, molecule_id=None):       

        self._id = str(chain_id).strip()
        self._accession = None
        self._type = None
        self._residues = ChainResiduesCollection(self, residues)
        self._secondary_structure = None
        self.molecule_id = molecule_id
        self._torsion_computed = False
        self.name = str(name).strip()
        
        self._structure = None
        
        self.type = type
        if accession is not None:
            self.accession = accession
            
    @property
    def _children(self):
        return self._residues

    def __repr__(self):
        return "<Chain {0.id}: {0.type!r}>".format(self)        

    def __len__(self):
        return self._residues.length

    @property
    def id(self):
        return self._id
    @id.setter
    def id(self, id):
        if not isinstance(id, basestring):
            raise ValueError(id)
        id = id.strip()
        if self._structure:
            self._structure.chains._update_chain_id(self, id)
        self._id = id
    
    @property
    def accession(self):
        return self._accession
    @accession.setter
    def accession(self, accession):
        if self._structure:
            raise InvalidOperation("Only structure's accession can be altered in this case")
        if accession is None:
            raise ValueError(accession)
        self._accession = str(accession).strip()
        
    @property
    def type(self):
        return self._type
    @type.setter
    def type(self, type):
        if not isinstance(type, csb.pyutils.EnumItem):
            raise TypeError(type)
        self._type = type

    @property
    def residues(self):
        return self._residues
    
    @property
    def torsion(self):
        if not self._torsion_computed:
            raise InvalidOperation('The correctness of the data is not guaranteed until chain.compute_torsion() is invoked.')
            
        torsion = TorsionAnglesCollection()
        
        for r in self.residues:
            if r.torsion is None:
                torsion.append(TorsionAngles(None, None, None))
            else:
                torsion.append(r.torsion)
                
        return torsion

    @property
    def length(self):
        return self._residues.length
    
    @property
    def entry_id(self):
        if self._accession and self._id:
            return self._accession + self._id
        else:
            return None
    
    @property
    def header(self):
        return ">{0._accession}_{0._id} mol:{1} length:{0.length} {0.name}".format(self, str(self.type).lower())

    @property
    def sequence(self):        
        sequence = []
        for residue in self._residues:
            if residue and residue.type:
                sequence.append(str(residue.type))
            else:
                sequence.append('-')
        return ''.join(sequence)
    
    @property
    def alphabet(self):
        if self._type == SequenceTypes.Protein:                             #@UndefinedVariable
            return SequenceAlphabets.Protein
        elif self._type == SequenceTypes.NucleicAcid:                       #@UndefinedVariable
            return SequenceAlphabets.Nucleic            
        else:
            raise NotImplementedError()
        
    @property
    def secondary_structure(self):
        return self._secondary_structure
    @secondary_structure.setter
    def secondary_structure(self, ss):
        if not isinstance(ss, SecondaryStructure):
            raise TypeError(ss)
        if len(ss) > 0:
            if (ss[ss.last_index].end > self._residues.last_index):
                raise ValueError('The secondary structure is out of range')
        self._secondary_structure = ss        
        
    def clone(self):
        """
        Make a deep copy of the chain. If this chain is part of a structure, 
        detach from it.
        
        @return: a deep copy of self
        @rtype: L{Chain}
        """
        
        new_chain = csb.pyutils.deepcopy(self) # self.subregion(self.residues.start_index, self.residues.last_index, clone=True)
        new_chain.type = self.type
        new_chain._structure = None
        
        return new_chain       
        
    def subregion(self, start, end, clone=False):
        """
        Extract a subchain defined by [start, end]. If clone is True, this
        is a deep copy of the chain. Otherwise same as:
        
            >>> chain.residues[start : end + 1]
        
        but coordinates are checked and a Chain instance is returned.
        
        @param start: start position of the sub-region
        @type start: int
        @param end: end position
        @type end: int
        @param clone: if True, a deep copy of the sub-region is returned, 
                      otherwise - a shallow one
        @type clone: bool
        
        
        @return: a new chain, made from the residues of the extracted region
        @rtype: L{Chain}
        
        @raise IndexError: if start/end positions are out of range
        """
        if start < self.residues.start_index or start > self.residues.last_index:
            raise IndexError('The start position is out of range {0.start_index} .. {0.last_index}'.format(self.residues))
        if end < self.residues.start_index or end > self.residues.last_index:
            raise IndexError('The end position is out of range {0.start_index} .. {0.last_index}'.format(self.residues))
                
        residues = self.residues[start : end + 1]
        
        if clone:
            residues = csb.pyutils.deepcopy(residues)
        
        chain = Chain(self.id, accession=self.accession, name=self.name, 
                      type=self.type, residues=residues, molecule_id=self.molecule_id)
        chain.secondary_structure = self.secondary_structure.subregion(start, end)
        chain._torsion_computed = self._torsion_computed
        
        return chain        
    
    def format_residue(self, residue):
        """
        Return a PDB-friendly string representation of the residue. Safer than repr(residue).
        
        @param residue: the residue name to format (a member of any L{SequenceAlphabets})
        @type residue: L{csb.pyutils.EnumItem}
        
        @return: PDB-friendly string representation of the residue type
        @rtype: str
        """
        if self.type == SequenceTypes.NucleicAcid:                                  #@UndefinedVariable
            return str(residue.type)
        else:
            return repr(residue.type)        
        
    def find(self, sequence_number, insertion_code=None):
        """
        Get a residue by its original Residue Sequence Number and Insertion Code.<commit_msg>Fixed a bug in Chain.subregion()<commit_after>"""
Object representations of 3D and secondary structures. 
"""

import os
import re
import copy
import math
import numpy
import csb.pyutils
import csb.math
import csb.bio.utils

from csb.bio.sequence import SequenceTypes, SequenceAlphabets, AlignmentTypes
from itertools import izip


AngleUnits = csb.pyutils.enum( Degrees='deg', Radians='rad' )
"""
Torsion angle unit types
"""

SecStructures = csb.pyutils.enum( Helix='H', Strand='E', Coil='C', Turn='T',
                                  Bend='S', Helix3='G', PiHelix='I', BetaBridge='B', Gap='-' )
"""
Secondary structure types
"""

ChemElements = csb.pyutils.enum( H=1, He=2, Li=3, Be=4, B=5, C=6, N=7, O=8, F=9, Ne=10, Na=11, Mg=12, Al=13, Si=14, P=15, 
                                 S=16, Cl=17, Ar=18, K=19, Ca=20, Sc=21, Ti=22, V=23, Cr=24, Mn=25, Fe=26, Co=27, Ni=28, 
                                 Cu=29, Zn=30, Ga=31, Ge=32, As=33, Se=34, Br=35, Kr=36, Rb=37, Sr=38, Y=39, Zr=40, Nb=41, 
                                 Mo=42, Tc=43, Ru=44, Rh=45, Pd=46, Ag=47, Cd=48, In=49, Sn=50, Sb=51, Te=52, I=53, Xe=54,
                                 Cs=55, Ba=56, Hf=72, Ta=73, W=74, Re=75, Os=76, Ir=77, Pt=78, Au=79, Hg=80, Tl=81, Pb=82, 
                                 Bi=83, Po=84, At=85, Rn=86, Fr=87, Ra=88, Rf=104, Db=105, Sg=106, Bh=107, Hs=108, Mt=109, 
                                 Ds=110, Rg=111, La=57, Ce=58, Pr=59, Nd=60, Pm=61, Sm=62, Eu=63, Gd=64, Tb=65, Dy=66, 
                                 Ho=67, Er=68, Tm=69, Yb=70, Lu=71, Ac=89, Th=90, Pa=91, U=92, Np=93, Pu=94, Am=95, Cm=96, 
                                 Bk=97, Cf=98, Es=99, Fm=100, Md=101, No=102, Lr=103,
                                 x=-1 ) 
"""
Periodic table elements
"""

class Broken3DStructureError(ValueError):
    pass
class Missing3DStructureError(Broken3DStructureError):
    pass       
class InvalidOperation(Exception):
    pass
class DuplicateChainIDError(csb.pyutils.DuplicateKeyError):
    pass
class DuplicateResidueIDError(csb.pyutils.DuplicateKeyError):
    pass
class DuplicateAtomIDError(csb.pyutils.DuplicateKeyError):
    pass
class AlignmentArgumentLengthError(ValueError):
    pass
class BrokenSecStructureError(ValueError):
    pass

class Ensemble(csb.pyutils.AbstractNIContainer):
    """
    Represents an ensemble of multiple L{Structure} models.
    Provides a list-like access to these models::
    
        >>> ensemble[0]
        <Structure Model 1: accn, x chains>
        >>> ensemble.models[1]
        <Structure Model 1: accn, x chains>
    """
    
    def __init__(self):
        self._models = EnsembleModelsCollection()
        
    def __repr__(self):
        return "<Ensemble: {0} models>".format(self.models.length)        
        
    @property
    def _children(self):
        return self._models
    
    @property
    def models(self):
        return self._models
        
    @property
    def first_model(self):
        if len(self._models) > 0:
            return self[0]
        return None
        
class EnsembleModelsCollection(csb.pyutils.CollectionContainer):
    
    def __init__(self):
        super(EnsembleModelsCollection, self).__init__(type=Structure, start_index=1)

class Structure(csb.pyutils.AbstractNIContainer):
    """
    Represents a single model of a PDB 3-Dimensional molecular structure.
    Provides access to the L{Chain} objects, contained in the model::
    
        >>> structure['A']
        <Chain A: Protein>
        >>> structure.chains['A']
        <Chain A: Protein>
        >>> structure.items[0]
        <Chain A: Protein> #(given a chain order A, B...)    
    
    @param accession: accession number of the structure
    @type accession: str
    """
    def __init__(self, accession):
                        
        self._accession = None
        self._chains = StructureChainsTable(self)
        self._chainslist = StructureChainsCollection(self)
        self.model_id = None
        
        self.accession = accession

    def __repr__(self):
        return "<Structure Model {0.model_id}: {0.accession}, {1} chains>".format(self, self.chains.length)

    @property
    def _children(self):
        return self._chains
    
    @property
    def chains(self):
        return self._chains
    
    @property
    def items(self):
        return self._chainslist
    
    @property
    def first_chain(self):
        if len(self._chains) > 0:
            return self._chainslist[0]
        return None
        
    @property
    def accession(self):
        return self._accession
    @accession.setter
    def accession(self, accession):
        if accession is None:
            raise ValueError(accession)
        self._accession = str(accession).strip().lower()
        for c in self.chains:
            self.chains[c]._accession = self._accession
            
    def define_molecule(self, molecule_id, chains):
        """
        Group C{chains}, defined by their IDs, to a polymer with id = C{molecule_id}.
        
        @param molecule_id: ID of the new polymer
        @param chains: a list of chain IDs
        @type chains: list
        
        @raise ValueError: if an invalid chain ID is provided
        """
        assert molecule_id is not None
        cs = [ ]
        
        for c in chains:
            if c not in self.chains:
                raise ValueError('No such chain {0} in structure'.format(c))
            cs.append(self.chains[c])
        for chain in cs:
            chain.molecule_id = molecule_id

    def apply_transformation(self, rotation, translation):
        """
        Apply in place RotationMatrix and translation Vector to all atoms in the structure.
        
        @type rotation: numpy array
        @type translation: numpy array 
        """
        for chain_id in self.chains:
            self.chains[chain_id].apply_transformation(rotation, translation)
                    
    def to_fasta(self):
        """
        Dump the structure in FASTA format. 
        
        @return: FASTA-formatted string with all chains in the structure
        @rtype: str
        """
        fasta = []
        
        for chain_id in self.chains:

            chain = self.chains[chain_id]
            if chain.length > 0:
                fasta.append(chain.header)
                fasta.append(chain.sequence)
        
        return os.linesep.join(fasta)

    def to_pdb(self, output_file=None):
        """
        Dump the whole structure in PDB format.
        
        @param output_file: output file name
        @type output_file: str 
        """
        from datetime import datetime
        import StringIO
        
        def isnull(this, that, null=None):
            if this is null: return that
            return this
            
        class MyStringIO(StringIO.StringIO):
            def writeline(self, data):
                self.write('{0:80}\n'.format(data))
        
        stream = MyStringIO()
        
        stream.writeline('HEADER    {0:40}{1:%d-%b-%y}   {2:4}'.format('.', datetime.now(), self.accession.upper()))
        
        molecules = { }
        for chain_id in self.chains:
            chain = self.chains[chain_id]
            if chain.molecule_id not in molecules:
                molecules[chain.molecule_id] = [ ]
            molecules[chain.molecule_id].append(chain_id)
        
        k = 0
        for mol_id in sorted(molecules):
            
            chains = molecules[mol_id]
            first_chain = self.chains[ chains[0] ]            
            
            stream.writeline('COMPND {0:3} MOL_ID: {1};'.format(k + 1, isnull(mol_id, '0')))
            stream.writeline('COMPND {0:3} MOLECULE: {1};'.format(k + 2, isnull(first_chain.name, '')))
            stream.writeline('COMPND {0:3} CHAIN: {1};'.format(k + 3, ', '.join(chains)))
            k += 3
            
        for chain_id in self.chains:
            
            chain = self.chains[chain_id]
            res = [ chain.format_residue(r) for r in chain.residues ]

            rn = 0
            for j in range(0, chain.length, 13):
                rn += 1
                residues = [ '{0:>3}'.format(r) for r in res[j : j + 13] ]
                stream.writeline('SEQRES {0:>3} {1} {2:>4}  {3}'.format(
                                            rn, chain.id, chain.length, ' '.join(residues) ))
        
        for chain_id in self.chains:
        
            chain = self.chains[chain_id]
            for residue in chain.residues:
        
                atoms = [ ]
                for an in residue.atoms:
                    atom = residue.atoms[an]
                    if type(atom) is DisorderedAtom:
                        for dis_atom in atom: atoms.append(dis_atom)
                    else:
                        atoms.append(atom)
                atoms.sort()
                
                for atom in atoms:

                    alt = atom.alternate
                    if alt is True:
                        alt = 'A'
                    elif alt is False:
                        alt = ' '
                    
                    if atom.element:
                        element = repr(atom.element)
                    else:
                        element = ' '
                    stream.writeline('ATOM  {0:>5} {1:>4}{2}{3:>3} {4}{5:>4}{6}   {7:>8.3f}{8:>8.3f}{9:>8.3f}{10:>6.2f}{11:>6.2f}{12:>12}{13:2}'.format(
                                        atom.serial_number, atom._full_name, isnull(alt, ' '), 
                                        chain.format_residue(residue), chain.id, 
                                        isnull(residue.sequence_number, residue.rank), isnull(residue.insertion_code, ' '), 
                                        atom.vector[0], atom.vector[1], atom.vector[2], isnull(atom.occupancy, 0.0), isnull(atom.temperature, 0.0), 
                                        element, isnull(atom.charge, ' ') ))
            stream.writeline('TER')
        stream.writeline('END')
        
        data = stream.getvalue()        
        stream.close()
        
        if not output_file:
            return data
        else:
            with open(output_file, 'w') as out:
                out.write(data)                

class StructureChainsCollection(csb.pyutils.AbstractContainer):
    
    def __init__(self, container):
        self.__container = container
        
    @property
    def _children(self):
        return [self.__container[c] for c in self.__container]

class StructureChainsTable(csb.pyutils.DictionaryContainer):
    
    def __init__(self, structure=None, chains=None):
        self.__container = structure
        super(StructureChainsTable, self).__init__()
        
        if chains is not None:
            for chain in chains:
                self.append(chain)
        
    def __repr__(self):
        if len(self) > 0:
            return "<StructureChains: {0}>".format(', '.join(self.keys()))
        else:
            return "<StructureChains: empty>"     
    
    def append(self, chain):
        """
        Add a new Chain to the structure.
        
        @param chain: the new chain to be appended
        @type chain: L{Chain}
        
        @raise DuplicateChainIDError: if a chain with same ID is already defined
        @raise InvalidOperation: if the chain is already associated with a structure
        """
        
        if chain._structure and chain._structure is not self.__container:
            raise InvalidOperation('This chain is already part of another structure')
        if chain.id in self:
            raise DuplicateChainIDError('A chain with ID {0} is already defined'.format(chain.id))
            
        super(StructureChainsTable, self).append(chain.id, chain)
        
        if self.__container:
            chain._accession = self.__container.accession
            chain._structure = self.__container

    def remove(self, id):
        """
        Remove a chain from the structure.

        @param id: ID of the chain to be detached
        @type id: str
        """
        chain = self[id]
        self._remove(id)
        chain._structure = None   
    
    def _update_chain_id(self, chain, new_id):
        
        if chain.id not in self or self[chain.id] is not chain:
            raise InvalidOperation(chain)
        
        self._remove(chain.id)
        
        if new_id in self:
            raise DuplicateChainIDError('Chain ID {0} is already defined'.format(id))
        
        super(StructureChainsTable, self).append(new_id, chain)
        
class Chain(csb.pyutils.AbstractNIContainer):
    """
    Represents a polymeric chain. Provides list-like and rank-based access to
    the residues in the chain::
    
        >>> chain[0]
        <ProteinResidue [1]: SER None>
        >>> chain.residues[1]
        <ProteinResidue [1]: SER None>
    
    You can also access residues by their PDB sequence number:
    
        >>> chain.find(sequence_number=5, insertion_code='A')
        <ProteinResidue [1]: SER 5A>
    
    @param chain_id: ID of the new chain
    @type chain_id: str
    @param type: sequence type (a member of the L{SequenceTypes} enum)
    @type type: L{csb.pyutils.EnumItem}
    @param name: name of the chain
    @type name: str
    @param residues: initialization list of L{Residue}-s
    @type residues: list
    @param accession: accession number of the chain
    @type accession: str
    @param molecule_id: MOL ID of the chain, if part of a polymer
    
    """
    def __init__(self, chain_id, type=SequenceTypes.Protein, name='',           #@UndefinedVariable
                 residues=None, accession=None, molecule_id=None):       

        self._id = str(chain_id).strip()
        self._accession = None
        self._type = None
        self._residues = ChainResiduesCollection(self, residues)
        self._secondary_structure = None
        self.molecule_id = molecule_id
        self._torsion_computed = False
        self.name = str(name).strip()
        
        self._structure = None
        
        self.type = type
        if accession is not None:
            self.accession = accession
            
    @property
    def _children(self):
        return self._residues

    def __repr__(self):
        return "<Chain {0.id}: {0.type!r}>".format(self)        

    def __len__(self):
        return self._residues.length

    @property
    def id(self):
        return self._id
    @id.setter
    def id(self, id):
        if not isinstance(id, basestring):
            raise ValueError(id)
        id = id.strip()
        if self._structure:
            self._structure.chains._update_chain_id(self, id)
        self._id = id
    
    @property
    def accession(self):
        return self._accession
    @accession.setter
    def accession(self, accession):
        if self._structure:
            raise InvalidOperation("Only structure's accession can be altered in this case")
        if accession is None:
            raise ValueError(accession)
        self._accession = str(accession).strip()
        
    @property
    def type(self):
        return self._type
    @type.setter
    def type(self, type):
        if not isinstance(type, csb.pyutils.EnumItem):
            raise TypeError(type)
        self._type = type

    @property
    def residues(self):
        return self._residues
    
    @property
    def torsion(self):
        if not self._torsion_computed:
            raise InvalidOperation('The correctness of the data is not guaranteed until chain.compute_torsion() is invoked.')
            
        torsion = TorsionAnglesCollection()
        
        for r in self.residues:
            if r.torsion is None:
                torsion.append(TorsionAngles(None, None, None))
            else:
                torsion.append(r.torsion)
                
        return torsion

    @property
    def length(self):
        return self._residues.length
    
    @property
    def entry_id(self):
        if self._accession and self._id:
            return self._accession + self._id
        else:
            return None
    
    @property
    def header(self):
        return ">{0._accession}_{0._id} mol:{1} length:{0.length} {0.name}".format(self, str(self.type).lower())

    @property
    def sequence(self):        
        sequence = []
        for residue in self._residues:
            if residue and residue.type:
                sequence.append(str(residue.type))
            else:
                sequence.append('-')
        return ''.join(sequence)
    
    @property
    def alphabet(self):
        if self._type == SequenceTypes.Protein:                             #@UndefinedVariable
            return SequenceAlphabets.Protein
        elif self._type == SequenceTypes.NucleicAcid:                       #@UndefinedVariable
            return SequenceAlphabets.Nucleic            
        else:
            raise NotImplementedError()
        
    @property
    def secondary_structure(self):
        return self._secondary_structure
    @secondary_structure.setter
    def secondary_structure(self, ss):
        if not isinstance(ss, SecondaryStructure):
            raise TypeError(ss)
        if len(ss) > 0:
            if (ss[ss.last_index].end > self._residues.last_index):
                raise ValueError('The secondary structure is out of range')
        self._secondary_structure = ss        
        
    def clone(self):
        """
        Make a deep copy of the chain. If this chain is part of a structure, 
        detach from it.
        
        @return: a deep copy of self
        @rtype: L{Chain}
        """
        
        new_chain = csb.pyutils.deepcopy(self) # self.subregion(self.residues.start_index, self.residues.last_index, clone=True)
        new_chain.type = self.type
        new_chain._structure = None
        
        return new_chain       
        
    def subregion(self, start, end, clone=False):
        """
        Extract a subchain defined by [start, end]. If clone is True, this
        is a deep copy of the chain. Otherwise same as:
        
            >>> chain.residues[start : end + 1]
        
        but coordinates are checked and a Chain instance is returned.
        
        @param start: start position of the sub-region
        @type start: int
        @param end: end position
        @type end: int
        @param clone: if True, a deep copy of the sub-region is returned, 
                      otherwise - a shallow one
        @type clone: bool
        
        
        @return: a new chain, made from the residues of the extracted region
        @rtype: L{Chain}
        
        @raise IndexError: if start/end positions are out of range
        """
        if start < self.residues.start_index or start > self.residues.last_index:
            raise IndexError('The start position is out of range {0.start_index} .. {0.last_index}'.format(self.residues))
        if end < self.residues.start_index or end > self.residues.last_index:
            raise IndexError('The end position is out of range {0.start_index} .. {0.last_index}'.format(self.residues))
                
        residues = self.residues[start : end + 1]
        
        if clone:
            residues = csb.pyutils.deepcopy(residues)
        
        chain = Chain(self.id, accession=self.accession, name=self.name, 
                      type=self.type, residues=residues, molecule_id=self.molecule_id)
        if chain.secondary_structure:
            chain.secondary_structure = self.secondary_structure.subregion(start, end)
        chain._torsion_computed = self._torsion_computed
        
        return chain        
    
    def format_residue(self, residue):
        """
        Return a PDB-friendly string representation of the residue. Safer than repr(residue).
        
        @param residue: the residue name to format (a member of any L{SequenceAlphabets})
        @type residue: L{csb.pyutils.EnumItem}
        
        @return: PDB-friendly string representation of the residue type
        @rtype: str
        """
        if self.type == SequenceTypes.NucleicAcid:                                  #@UndefinedVariable
            return str(residue.type)
        else:
            return repr(residue.type)        
        
    def find(self, sequence_number, insertion_code=None):
        """
        Get a residue by its original Residue Sequence Number and Insertion Code.
{'fraction_non_alphanumeric': 0.0995666756432541, 'fraction_numerical': 0.017381937473988297, 'mean_word_length': 1.7865475134729518, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 127, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2462929', 'n_tokens_mistral': 12401, 'n_tokens_neox': 11576, 'n_words': 3103}