# --- UCSF Chimera Copyright ---
# Copyright (c) 2000 Regents of the University of California.
# All rights reserved.  This software provided pursuant to a
# license agreement containing restrictions on its disclosure,
# duplication and use.  This notice must be embedded in or
# attached to all copies, including partial copies, of the
# software or any revisions or derivations thereof.
# --- UCSF Chimera Copyright ---

import string
import chimera
from chimera import elements
import os
import operator

FORMAT_MOL2 = '.mol2'
FORMAT_PDB = '.pdb'

class Atom:
	"""Class for holding information about a single atom"""

	def __init__(self, aname, rtype, rseq, x, y, z):
		self.aname = aname
		self.rtype = rtype
		self.rseq = rseq
		self.x = x
		self.y = y
		self.z = z

	def makeChimera(self, mol):
		nameList = [ 'LP' ]	# default to lone pair
		nameList.insert(0, self.aname[0])
		if len(self.aname) > 1 and self.aname[1] in string.uppercase:
			s = self.aname[0] + string.lower(self.aname[1])
			nameList.insert(0, s)
			nameList.insert(0, self.aname[:2])
		for name in nameList:
			try:
				num = elements.name.index(name)
				break
			except ValueError:
				pass
		e = chimera.Element(num)
		self.chimeraAtom = mol.newAtom(self.aname, e)
		coord = chimera.Coord()
		coord.x = self.x
		coord.y = self.y
		coord.z = self.z
		self.chimeraAtom.setCoord(coord)

class Compound:
	"""Class for holding information about a single compound"""

	Undefined = -1
	Viable = 0
	Deleted = 1
	Purged = 2

	def __init__(self):
		self.label = 'Unlabeled'
		self.atoms = {}
		self.conect = []
		self.header = ''
		self.compnd = ''
		self.author = ''
		self.rawText = []
		self.fields = {}
		self.state = Compound.Viable
		self.listed = 0
		self.data = []
		self.prevData = []
		self.textArray = []

	def readPDB(self, f):
		while 1:
			line = f.readline()
			if not line:
				break
			if line[:3] == 'TER':
				self.data.append(line)
				break
			elif line[:6] == 'ATOM  ':
				self._pdbParseAtom(line)
				self.data.append(line)
			elif line[:6] == 'CONECT':
				self._pdbParseConect(line)
				self.data.append(line)
			elif line[:6] == 'REMARK':
				if self._pdbParseRemark(line):
					self.rawText.append(line)
			elif line[:6] == 'HEADER':
				self._pdbParseHeader(line)
				self.rawText.append(line)
			elif line[:6] == 'COMPND':
				self._pdbParseCompnd(line)
				self.rawText.append(line)
			elif line[:6] == 'AUTHOR':
				self._pdbParseAuthor(line)
				self.rawText.append(line)

	def _pdbParseAtom(self, line):
		field = string.split(line)
		if len(field) != 8:
			if len(field) == 7:
				# Mostly this is caused by 4-letter
				# atom names
				aname = field[2][:4]
				rtype = field[2][4:]
				del field[2]
				field.insert(2, rtype)
				field.insert(2, aname)
			else:
				print 'ATOM record with %d fields!' % len(field)
				print line
				return
		serial = string.atoi(field[1])
		aname = field[2]
		rtype = field[3]
		rseq = string.atoi(field[4])
		x = string.atof(field[5])
		y = string.atof(field[6])
		z = string.atof(field[7])
		self.atoms[serial] = Atom(aname, rtype, rseq, x, y, z)

	def _pdbParseConect(self, line):
		field = string.split(line)
		if len(field) < 3:
			print 'CONECT record with %d fields!' % len(field)
			print line
			return
		n = string.atoi(field[1])
		for d in field[2:]:
			m = string.atoi(d)
			if m > n:
				self.conect.append((n, m))

	def _pdbParseRemark(self, line):
		return self._parseDock(line, len('REMARK'))

	def _pdbParseHeader(self, line):
		self.header = string.strip(line[6:])
		self.textArray.append(line)

	def _pdbParseCompnd(self, line):
		self.compnd = string.strip(line[6:])
		self.textArray.append(line)

	def _pdbParseAuthor(self, line):
		self.author = string.strip(line[6:])
		self.textArray.append(line)

	def _parseDock(self, line, prefixLen):
		keepRemark = 1
		field = string.split(line[prefixLen:], ':')
		if len(field) == 2:
			key = string.strip(field[0])
			value = string.strip(field[1])
			try:
				value = string.atoi(value)
			except ValueError:
				try:
					value = string.atof(value)
				except ValueError:
					pass
			if key == 'viewdock state' and type(value) == type(0):
				self.state = value
				keepRemark = 0
			else:
				self.fields[key] = value
		else:
			field = string.split(line[prefixLen:])
			key = None
			if len(field) > 1:
				last = field[-1]
				value = None
				try:
					value = string.atoi(last)
				except ValueError:
					try:
						value = string.atof(last)
					except ValueError:
						pass
				if value:
					key = line[prefixLen:-len(last)-1]
					self.fields[string.strip(key)] = value
		if keepRemark:
			self.textArray.append(line)
		return keepRemark

	def readMol2(self, f):
		if self._mol2ReadDockComments(f) < 0:
			return
		if self._mol2ReadMoleculeSection(f) < 0:
			return
		if self._mol2ReadAtomSection(f) < 0:
			return
		if self._mol2ReadBondSection(f) < 0:
			return

	def _mol2ReadDockComments(self, f):
		state = 0
		while 1:
			line = f.readline()
			if not line:
				return -1
			if line[0] == '#':
				state = 1
				if self._parseDock(line, len('##########')):
					self.rawText.append(line)
			elif state == 1:
				self.data.append(line)
				line = string.lower(string.strip(line))
				if line == '@<tripos>molecule':
					break
			else:
				self.prevData.append(line)
		return 0

	def _mol2ReadMoleculeSection(self, f):
		state = 1
		while 1:
			line = f.readline()
			if not line:
				return -1
			# Note that state 0 of checking for the
			# correct head is done in _mol2ReadDockComments
			self.data.append(line)
			if state == 1:
				# mol_name
				self.name = line
				state = 2
			elif state == 2:
				# atoms bonds subst feat sets
				data = string.split(line)
				if len(data) >= 3:
					self.numAtoms = string.atoi(data[0])
					self.numBonds = string.atoi(data[1])
					self.numSubst = string.atoi(data[2])
				state = 3
			elif state == 3:
				# mol_type (e.g., SMALL)
				state = 4
			elif state == 4:
				# charge_type (e.g., NO_CHARGES)
				state = 5
			elif state == 5:
				# status_bits (e.g., analyzed)
				state = 6
			elif state == 6:
				# mol_comment
				break
		return 0

	def _mol2ReadAtomSection(self, f):
		while 1:
			line = f.readline()
			if not line:
				return -1
			self.data.append(line)
			line = string.lower(string.strip(line))
			if line == '@<tripos>atom':
				break
		for i in range(self.numAtoms):
			line = f.readline()
			if not line:
				self.atoms = {}
				return -1
			data = string.split(line)
			if len(data) < 6:
				self.atoms = {}
				return -1
			self.data.append(line)
			serial = string.atoi(data[0])
			aname = data[1]
			x = string.atof(data[2])
			y = string.atof(data[3])
			z = string.atof(data[4])
			try:
				rseq = string.atoi(data[6])
				rtype = data[7]
			except (ValueError, IndexError):
				rseq = 0
				rtype = 'UNK'
			self.atoms[serial] = Atom(aname, rtype, rseq, x, y, z)
		return 0

	def _mol2ReadBondSection(self, f):
		while 1:
			line = f.readline()
			if not line:
				return -1
			self.data.append(line)
			line = string.lower(string.strip(line))
			if line == '@<tripos>bond':
				break
		for i in range(self.numBonds):
			line = f.readline()
			if not line:
				self.atoms = {}
				return -1
			data = string.split(line)
			if len(data) < 4:
				self.atoms = {}
				return -1
			self.data.append(line)
			n = string.atoi(data[1])
			m = string.atoi(data[2])
			self.conect.append((n, m))
		return 0

	def buildChimera(self):
		self.text = string.join(self.textArray, '')
		del self.textArray
		mol = chimera.Molecule()
		mol.name = self.compnd
		for a in self.atoms.values():
			a.makeChimera(mol)
		res = mol.newResidue(a.rtype, ' ', a.rseq, ' ')
		for a in self.atoms.values():
			res.addAtom(a.chimeraAtom)
		for f, t in self.conect:
			try:
				a1 = self.atoms[f]
				a2 = self.atoms[t]
				mol.newBond(a1.chimeraAtom, a2.chimeraAtom)
			except KeyError:
				pass
		self.chimeraModel = mol

	def buildLabel(self):
		self.label = 'Chimera Model %s' % self.chimeraModel.oslIdent()

	def hide(self):
		mList = chimera.openModels.list(id=self.chimeraModel.id,
						subid=self.chimeraModel.subid)
		for m in mList:
			m.display = 0

	def show(self):
		mList = chimera.openModels.list(id=self.chimeraModel.id,
						subid=self.chimeraModel.subid)
		for m in mList:
			m.display = 1

	def displayed(self):
		return self.chimeraModel.display

	def save(self, format, f):
		f.write(string.join(self.rawText, ''))
		if format is FORMAT_MOL2:
			f.write('########## viewdock state: %d\n' % self.state)
		else:
			f.write('REMARK viewdock state: %d\n' % self.state)
		f.write(string.join(self.data, ''))

class Results:
	"""Class for holding information about a set of DOCK results"""

	def __init__(self, filename, notifySelectionChange):
		self.notifySelectionChange = notifySelectionChange
		self.filename = filename
		n = string.rindex(filename, '.')
		format = FORMAT_PDB
		if n > 0:
			if filename[n:] == FORMAT_MOL2:
				format = FORMAT_MOL2
		f = open(filename)
		self.compoundList = []
		prevC = None
		while 1:
			c = Compound()
			if format is FORMAT_MOL2:
				c.readMol2(f)
				if prevC and c.prevData:
					prevC.data = prevC.data + c.prevData
				prevC = c
			else:
				c.readPDB(f)
			if len(c.atoms) == 0:
				break
			self.compoundList.append(c)
		f.close()
		self.format = format
		for i in range(len(self.compoundList)):
			self.compoundList[i].buildChimera()
		chimera.openModels.add(map(lambda c: c.chimeraModel,
						self.compoundList))
		keys = {}
		for c in self.compoundList:
			for k in c.fields.keys():
				keys[k] = 1
			c.buildLabel()
		self.fields = keys.keys()
		self.fields.sort()

		self.titleWidth = {}
		self.dataWidth = {}
		self.alignment = {}
		for k in self.fields:
			self.titleWidth[k] = len(k)
			self.dataWidth[k] = 0
			self.alignment[k] = ''
		for c in self.compoundList:
			for k, v in c.fields.items():
				if type(v) == type(''):
					self.alignment[k] = '-'
				w = len(str(v))
				if w > self.dataWidth[k]:
					self.dataWidth[k] = w
		self.listBoxMap = {}
		self.selected = self.compoundList[:]

		listBoxKeys = []
		if keys.has_key('Number'):
			listBoxKeys.append('Number')
		if keys.has_key('Name'):
			listBoxKeys.append('Name')
		if len(listBoxKeys) == 0:
			listBoxKeys.append(self.fields[0])
		self.setListBoxKeys(listBoxKeys)
		self.sortedCompoundList = self.compoundList[:]

	def setListBoxKeys(self, keys):
		self.listBoxKeys = tuple(['S'] + keys)
		titleFmts = ['%1s']
		dataFmts = ['%1s']
		for k in keys:
			t, d = self._makeFormats(self.titleWidth[k],
							self.dataWidth[k],
							self.alignment[k])
			titleFmts.append(t)
			dataFmts.append(d)
		self.titleFmt = string.join(titleFmts, '  ')
		self.dataFmt = string.join(dataFmts, '  ')

	def _makeFormats(self, tw, dw, a):
		if tw > dw:
			tpad = 0
			twidth = tw
			dpad = (tw - dw) / 2
			dwidth = twidth - dpad
		else:
			dpad = 0
			dwidth = dw
			tpad = (dw - tw) / 2
			twidth = dwidth - tpad
		return '%%%ds%s' % (twidth, tpad * ' '), \
			'%%%s%ds%s' % (a, dwidth, dpad * ' ')

	def addColumn(self, col):
		if col not in self.fields or col in self.listBoxKeys:
			return 0
		self.setListBoxKeys(list(self.listBoxKeys)[1:] + [col])
		return 1

	def deleteColumn(self, col):
		if col not in self.fields or col not in self.listBoxKeys:
			return 0
		l = list(self.listBoxKeys)[1:]
		l.remove(col)
		self.setListBoxKeys(l)
		return 1

	def setSortKey(self, k):
		l = []
		for c in self.compoundList:
			try:
				v = c.fields[k]
			except KeyError:
				v = ''
			l.append((v, c))
		l.sort()
		self.sortedCompoundList = map(lambda t: t[1], l)

	def updateListBox(self, show, label, listbox):
		values = []
		compounds = []
		for c in self.sortedCompoundList:
			if not show[c.state]:
				c.listed = 0
				continue
			c.listed = 1
			if c.state == Compound.Viable:
				v = ['V']
			elif c.state == Compound.Deleted:
				v = ['D']
			else:
				v = ['P']
			for k in self.listBoxKeys[1:]:
				try:
					v.append(str(c.fields[k]))
				except KeyError:
					v.append('')
			values.append(v)
			compounds.append(c)
		label.config(text=self.titleFmt % self.listBoxKeys)
		sList = map(lambda v, fmt=self.dataFmt: fmt % tuple(v), values)
		self.listBoxMap = {}
		map(lambda k, v, m=self.listBoxMap: operator.setitem(m, k, v),
			sList, compounds)
		listbox.setlist(sList)
		selected = []
		n = -1
		for c in self.selected:
			try:
				n = compounds.index(c)
				selected.append(c)
				listbox.selection_set(n)
			except ValueError:
				c.hide()
		if n >= 0:
			listbox.see(n)
		self.selected = selected
		self.notifySelectionChange(self.selected)

	def pruneModels(self, molList):
		compounds = self._getCompounds(molList)
		for c in compounds:
			if c.state == Compound.Viable:
				c.state = Compound.Deleted
			c.hide()

	def selectModels(self, molList):
		self.selected = self._getCompounds(molList)
		self.notifySelectionChange(self.selected)

	def _getCompounds(self, molList):
		compounds = []
		for c in self.compoundList:
			if c.chimeraModel in molList:
				compounds.append(c)
		return compounds

	def updateSelectedCompounds(self, listbox):
		compounds = []
		for s in listbox.getcurselection():
			try:
				compounds.append(self.listBoxMap[s])
			except KeyError:
				pass
		for c in self.selected:
			if c not in compounds:
				c.hide()
		for c in compounds:
			if c not in self.selected:
				c.show()
		self.selected = compounds
		self.notifySelectionChange(self.selected)

	def setSelectedState(self, state):
		for c in self.selected:
			c.state = state

	def hideAll(self):
		for c in self.compoundList:
			c.hide()

	def hideViable(self):
		for c in self.compoundList:
			if c.state == Compound.Viable:
				c.hide()

	def hideDeleted(self):
		for c in self.compoundList:
			if c.state == Compound.Deleted:
				c.hide()

	def hidePurged(self):
		for c in self.compoundList:
			if c.state == Compound.Purged:
				c.hide()

	def displayAll(self):
		for c in self.compoundList:
			c.show()

	def displayViable(self):
		for c in self.compoundList:
			if c.state == Compound.Viable:
				c.show()

	def displayDeleted(self):
		for c in self.compoundList:
			if c.state == Compound.Deleted:
				c.show()

	def displayPurged(self):
		for c in self.compoundList:
			if c.state == Compound.Purged:
				c.show()

	def save(self, filename=None, skipPurged=0):
		if filename == None:
			filename = self.filename
		try:
			os.rename(filename, filename + '.save')
		except os.error:
			pass
		f = open(filename, 'w')
		for c in self.compoundList:
			if not skipPurged or c.state != Compound.Purged:
				c.save(self.format, f)
		f.close()

	def movieSetup(self, listbox):
		self.movieList = []
		first = None
		for c in self.sortedCompoundList:
			if not c.listed:
				continue
			self.movieList.append(c)
			if c.displayed():
				if first:
					c.hide()
				else:
					first = c
		if first is None:
			first = self.movieList[0]
			first.show()
		if len(self.movieList) < 2:
			print 'Not enough compounds displayed in list'
			return 0
		listbox.selection_clear(0, 'end')
		self.movieIndex = self.movieList.index(first)
		listbox.selection_set(self.movieIndex)
		listbox.see(self.movieIndex)
		self.selected = [first]
		self.notifySelectionChange(self.selected)
		return 1

	def movieStep(self, listbox):
		prev = self.movieList[self.movieIndex]
		listbox.selection_clear(self.movieIndex)
		self.movieIndex = self.movieIndex + 1
		if self.movieIndex >= len(self.movieList):
			self.movieIndex = 0
		listbox.selection_set(self.movieIndex)
		listbox.see(self.movieIndex)
		next = self.movieList[self.movieIndex]
		prev.hide()
		next.show()
		self.selected = [next]
		self.notifySelectionChange(self.selected)
		return 1

	def movieStop(self):
		return

if 0:
	r = Results('sample.pdb')
	for m in chimera.openModels.list():
		print m.name, m.id, m.subid
	title, strings = r.listBoxStrings(['Name', 'Number'])
	print title
	for s in strings:
		print s