# Copyright (C) 2002, Thomas Hamelryck (thamelry@vub.ac.be)
# 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.  
from __future__ import generators

from Numeric import Float0
from copy import copy

from PDBExceptions import PDBConstructionException

__doc__="""
Base class for Residue, Chain, Model and Structure classes. 
It is a simple container class, with list and dictiuonary like properties.
"""


class Entity:
    """
    Basic container object. Structure, Model, Chain and Residue
    are subclasses of Entity. It deals with storage and lookup.
    """
    def __init__(self, id):
        self.id=id
        self.full_id=None
        self.parent=None
        self.child_list=[]
        self.child_dict={}
    
    # Special methods   

    def __len__(self):
        "Return the number of children."
        return len(self.child_list)

    def __getitem__(self, id):
        "Return the child with given id."
        return self.child_dict[id]

    def __delitem__(self, id):
        "Delete (and destroy) the child with given id."
        child=self[id]
        child.destroy()
        del self.child_dict[id]
        self.child_list=self.child_dict.values()
        self.child_list.sort(self._sort)

    # Private methods

    def _sort(self, e1, e2):
        """Sort the children in the Entity object.

        This method is implemented by the Entity subclasses. A Chain object
        e.g. sorts Residue objects according to hetero field, sequence
        identifier and insertion code.
        """
        raise NotImplementedError

    # Public methods    

    def get_level(self):
        """Return level in hierarchy.

        A - atom
        R - residue
        C - chain
        M - model
        S - structure
        """
        return self.level

    def sort(self, sort_function=None):
        """Sort the children in the Entity object.

        Arguments:
        o sort_function - Optional. A function that compares two Entities. If this
        argument is present, the default sort method is replaced with this function.
        """
        if sort_function is None:
            self.child_list.sort(self._sort)
        else:
            self.child_list.sort(sort_function)
        for child in self.child_list:
            child.sort()

    def set_parent(self, entity):
        "Set the parent Entity object."
        self.parent=entity

    def detach_parent(self):
        "Detach the parent."
        self.parent=None

    def detach_child(self, id):
        "Remove (but do not destroy) a child."
        child=self.child_dict[id] 
        child.detach_parent()
        del self.child_dict[id]
        self.child_list=self.child_dict.values()
        self.child_list.sort(self._sort)

    def destroy(self):
        """Destroy the Entity.

        This trashes the Entity object and breaks all circular references
        (eh, yes, I'm still using 1.5).
        """
        for child in self.child_list:
            child.destroy()
        del self.parent
        del self.child_list
        del self.child_dict
        del self.id

    def add(self, entity):
        "Add a child to the Entity."
        entity_id=entity.get_id()
        if self.has_id(entity_id):
            raise PDBConstructionException, "%s defined twice" % entity.get_full_id()
        entity.set_parent(self)
        self.child_list.append(entity)
        #self.child_list.sort(self._sort)
        self.child_dict[entity_id]=entity
    
    def get_iterator(self):
        "Return iterator over children."
        for child in self.child_list:
            yield child

    def get_list(self):
        "Return a copy of the list of children."
        return copy(self.child_list)

    def has_id(self, id):
        "Return 1 if a child with given id exists, otherwise 0."
        return self.child_dict.has_key(id)

    def get_parent(self):
        "Return the parent Entity object."
        return self.parent

    def get_id(self):
        "Return the id."
        return self.id

    def get_full_id(self):
        """Return the full id.

        The full id is a tuple containing all id's starting from
        the top object (Structure) down to the current object. A full id for
        a Residue object e.g. is something like:

        ("1abc", 0, "A", (" ", 10, "A"))

        This corresponds to:

        Structure with id "1abc"
        Model with id 0
        Chain with id "A"
        Residue with id (" ", 10, "A")

        The Residue id indicates that the residue is not a hetero-residue 
        (or a water) beacuse it has a blanc hetero field, that its sequence 
        identifier is 10 and its insertion code "A".
        """
        if self.full_id==None:
            entity_id=self.get_id()
            l=[entity_id]   
            parent=self.get_parent()
            while parent!=None:
                entity_id=parent.get_id()
                l.append(entity_id)
                parent=parent.get_parent()
            l.reverse()
            self.full_id=tuple(l)
        return self.full_id



class DisorderedEntityWrapper:
    """
    This class is a simple wrapper class that groups a number of equivalent
    Entities and forwards all method calls to one of them (the currently selected 
    object). DisorderedResidue and DisorderedAtom are subclasses of this class.
    
    E.g.: A DisorderedAtom object contains a number of Atom objects,
    where each Atom object represents a specific position of a disordered
    atom in the structure.
    """
    def __init__(self, id):
        self.id=id
        self.child_dict={}
        self.selected_child=None
        self.parent=None    

    # Special methods

    def __getattr__(self, method):
        "Forward the method call to the selected child."
        return getattr(self.selected_child, method)

    def __setitem__(self, id, child):
        "Add a child, associated with a certain id."
        self.child_dict[id]=child

    # Public methods    

    def get_id(self):
        "Return the id."
        return self.id

    def disordered_has_id(self, id):
        "Return 1 if there is an object present associated with this id."
        return self.child_dict.has_key(id)

    def destroy(self):
        "Destroy the object (and its children)."
        for child in self.disordered_get_list():
            child.destroy()
        del self.parent
        del self.selected_child
        del self.id
        del self.child_dict

    def detach_parent(self):
        "Detach the parent"
        self.parent=None
        for child in self.disordered_get_list():
            child.detach_parent()

    def get_parent(self):
        "Return parent."
        return self.parent

    def set_parent(self, parent):
        "Set the parent for the object and its children."
        self.parent=parent
        for child in self.disordered_get_list():
            child.set_parent(parent)

    def disordered_select(self, id):
        """Select the object with given id as the currently active object.

        Uncaught method calls are forwarded to the selected child object.
        """
        self.selected_child=self.child_dict[id]
    
    def disordered_add(self, child):
        "This is implemented by DisorderedAtom and DisorderedResidue."
        raise NotImplementedError

    def is_disordered(self):
        """
        Return 2, indicating that this Entity is a collection of Entities.
        """
        return 2

    def disordered_get_id_list(self):
        "Return a list of id's."
        l=self.child_dict.keys()
        # sort id list alphabetically
        l.sort()
        return l
        
    def disordered_get(self, id=None):
        """Get the child object associated with id.

        If id is None, the currently selected child is returned.
        """
        if id==None:
            return self.selected_child
        return self.child_dict[id]

    def disordered_get_list(self):
        "Return list of children."
        return self.child_dict.values()

        
