# 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.  


# My Stuff
from PDBExceptions import PDBConstructionException
from Entity import Entity, DisorderedEntityWrapper


atom_name_dict={}
atom_name_dict["N"]=1
atom_name_dict["CA"]=2
atom_name_dict["C"]=3
atom_name_dict["O"]=4


class Residue(Entity):
	"""
	Represents a residue. A Residue object stores atoms.
	"""
	def __init__(self, id, resname, segid):
		self.level="R"
		self.disordered=0
		self.resname=resname
		self.segid=segid
		Entity.__init__(self, id)

	# Special methods

	def __repr__(self):
		return "<Residue het=%s resseq=%s icode=%s>" % self.get_id()

	# Private methods

	def _sort(self, a1, a2):
		"""Sort the Atom objects.

		Atoms are sorted alphabetically according to their name, 
		but N, CA, C, O always come first.

		Arguments:
		o a1, a2 - Atom objects
		"""
		name1=a1.get_name()
		name2=a2.get_name()
		if name1==name2:
			return(cmp(a1.get_altloc(), a2.get_altloc()))
		if atom_name_dict.has_key(name1):
			index1=atom_name_dict[name1]
		else:
			index1=None
		if atom_name_dict.has_key(name2):
			index2=atom_name_dict[name2]
		else:
			index2=None
		if index1 and index2:
			return cmp(index1, index2)
		if index1:
			return -1
		if index2:
			return 1
		return cmp(name1, name2)

	# Public methods

	def add(self, atom):
		"""Add an Atom object.

		Checks for adding duplicate atoms, and raises a
		PDBConstructionException if so.
		"""
		atom_id=atom.get_id()
		if self.has_id(atom_id):
			raise PDBConstructionException, "Atom %s defined twice in residue %s" % (atom_id, self)
		Entity.add(self, atom)

	def sort(self):
		self.child_list.sort(self._sort)

	def flag_disordered(self):
		"Set the disordered flag."
		self.disordered=1

	def is_disordered(self):
		"Return 1 if the residue contains disordered atoms."
		return self.disordered

	def get_resname(self):
		return self.resname

	def get_unpacked_list(self):
		"""
		Returns the list of all atoms, unpack DisorderedAtoms."
		"""
		atom_list=self.get_list()
		undisordered_atom_list=[]
		for atom in atom_list:
			if atom.is_disordered():
				undisordered_atom_list=(undisordered_atom_list+ atom.disordered_get_list())
			else:
				undisordered_atom_list.append(atom)
		return undisordered_atom_list		

	def get_segid(self):
		return self.segid


class DisorderedResidue(DisorderedEntityWrapper):
	"""
	DisorderedResidue is a wrapper around two or more Residue objects. It is
	used to represent point mutations (e.g. there is a Ser 60 and a Cys 60 residue,
	each with 50 % occupancy).
	"""
	def __init__(self, id):
		DisorderedEntityWrapper.__init__(self, id)

	def __repr__(self):
		return "<DisorderedResidue het=%s resseq=%i icode=%s>" % self.get_id()

	def add(self, atom):
		residue=self.disordered_get()
		if not atom.is_disordered()==2:
			# Atoms in disordered residues should have non-blanc
			# altlocs, and are thus represented by DisorderedAtom objects.
			resname=residue.get_resname()
			het, resseq, icode=residue.get_id()	
			# add atom anyway, if PDBParser ignores exception the atom will be part of the residue
			residue.add(atom)
			raise PDBConstructionException, "Blank altlocs in duplicate residue %s (%s, %i, %s)" % (resname, het, resseq, icode) 
		residue.add(atom)

	def sort(self):
		"Sort the atoms in the child Residue objects."
		for residue in self.disordered_get_list():
			residue.sort() 

	def disordered_add(self, residue):
		"""Add a residue object and use its resname as key.

		Arguments:
		o residue - Residue object
		"""
		resname=residue.get_resname()
		# add chain parent to residue
		chain=self.get_parent()
		residue.set_parent(chain)
		assert(not self.disordered_has_id(resname))
		self[resname]=residue
		self.disordered_select(resname)


	

	
