# --- UCSF Chimera Copyright ---
# Copyright (c) 2000-2011 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 ---
#
# $Id: copyright 34705 2011-10-19 23:37:43Z pett $

from chimera.baseDialog import ModelessDialog, ModalDialog
from MAViewer import MOD_ASSOC
from prefs import prefs, MODELLER_KEY, MODELLER_PATH, MODELLER_TEMP_PATH, MODELLER_USE_WEB
from chimera import CancelOperation

class LoopsDialog(ModelessDialog):
	ModellerHomeButtonText = "Modeller Home Page"
	buttons = ('OK', 'Apply', 'Close', ModellerHomeButtonText)
	default = 'OK'
	title = "Model Loops / Refine Structure"
	help = "ContributedSoftware/multalignviewer/modeller.html"
	provideStatus = True
	statusPosition = "above"

	def __init__(self, mav, *args, **kw):
		self.mav = mav
		ModelessDialog.__init__(self, *args, **kw)

	def fillInUI(self, parent):
		import Tkinter, Pmw

		row = 0

		f = Tkinter.Frame(parent)
		f.grid(row=row, column=0, columnspan=2)
		f.rowconfigure(0, weight=4) # get model selection approx. off to right of both
		f.rowconfigure(1, weight=1) # structure options
		self.modelingArea = Pmw.RadioSelect(f, buttontype='radiobutton', labelpos='w',
			label_text="Model/remodel:", orient='vertical', pady=0, command=self._areaChoiceCB)
		self.modelingArea.grid(row=0, column=0, rowspan=2)
		self.modelingArea.add("region", text="active region", anchor="w")
		self.modelingArea.add("selection", text="Chimera selection region", anchor="w")
		self.modelingArea.add("internal", text="non-terminal missing structure", anchor="w")
		self.modelingArea.add("missing", text="all missing structure", anchor="w")
		from chimera.widgets import MoleculeOptionMenu
		self.structureMenu = MoleculeOptionMenu(f, command=self._structChoiceCB,
				filtFunc=lambda m, mav=self.mav: m in mav.associations)
		self.structureMenu.grid(row=1, column=1)
		self._assocHandlerID = self.mav.triggers.addHandler(MOD_ASSOC,
				lambda *args: self.structureMenu.refresh(), None)
		row += 1

		from chimera.tkoptions import EnumOption, SymbolicEnumOption, IntOption, StringOption
		class AdjoiningOption(EnumOption):
			values = range(4)
		self.mobileEdgesOption = AdjoiningOption(parent, row, "Allow this many residues"
			" adjacent\nto missing regions to move",
			1, None, balloon="How many of the existing residues on each side of\n"
			"missing residues are allowed to move during loop placement")
		row += 1

		self.numModelsOption = IntOption(parent, row, "Number of models to generate",
			5, None, min=1, max=1000, balloon="Number of models limited to a maximum of 1000")
		row += 1

		class MethodOption(SymbolicEnumOption):
			labels, values = zip(("standard", ""), ("DOPE", "dope_"), ("DOPE-HR", "dopehr_"))
		self.methodOption = MethodOption(parent, row, "Loop modeling protocol",
			"", None, balloon="Method used to refine loops.\n"
			"Note that while the DOPE protocols are more accurate,\n"
			"they are slower and may fail to produce any models.")
		row += 1

		class ExecutionOption(SymbolicEnumOption):
			labels, values = zip(("web service", "web"), ("local installation", "local"))
		self.executionOption = ExecutionOption(parent, row,
			"Run Modeller using", "web", self._executionChangeCB)
		row += 1

		self.keyOption = StringOption(parent, row, "Modeller license key", prefs[MODELLER_KEY],
			None, show="*", balloon="Use '%s' button to register and get a license key"
			% self.ModellerHomeButtonText)
		# alterate option for local installation...
		# don't have two widgets in same cell at once or else forget/manage will get confused
		self.keyOption.forget()
		from chimera.tkoptions import InputFileOption, OutputFileOption
		self.executableLocationOption = InputFileOption(parent, row, "Location of Modeller"
			" executable", prefs[MODELLER_PATH], None, title="Choose Modeller Executable")
		self.executableLocationOption.forget()
		self.keyOption.manage()
		row += 1

		self.scriptLocationOption = InputFileOption(parent, row, "Custom Modeller script file\n"
			"(optional, overrides dialog)", None, None,
			balloon="Use a custom Modeller script instead of the one"
			" that would be generated by the dialog", title="Choose Modeller Script Location")
		self.scriptLocationOption.forget()
		row += 1

		self.scriptButton = Tkinter.Button(parent, text="Get Current Modeller Script",
				command=self._viewScript)
		self.scriptButton.grid(row=row, column=0, columnspan=2)
		self.scriptButton.grid_remove()
		from chimera import help
		help.register(self.scriptButton, balloon="View the current Modeller script, which\n"
			"has incorporated the above options.")
		row += 1

		self.tempFolderOption = OutputFileOption(parent, row, "Temporary folder location"
			" (optional)", prefs[MODELLER_TEMP_PATH], None,
			title="Choose Folder for Temporary Files", dirsOnly=True,
			balloon= "You can specify a folder for temporary files if you wish to check\n"
			"them afterwards.  Otherwise, a location will be generated automatically\n"
			"and will be cleaned up after execution.")
		row += 1

		from ModellerBase import ModellerCitation
		ModellerCitation(parent).grid(row=row, column=0, columnspan=2)
		row += 1

		self.structureMenu.invoke()
		exeOptVal = ExecutionOption.values[not prefs[MODELLER_USE_WEB]]
		if exeOptVal != self.executionOption.get():
			self.executionOption.set(exeOptVal)
			self._executionChangeCB(self.executionOption)

	def Apply(self):
		self.status("")
		try:
			loopInfo, seq, templateModels = self._gatherLoopInfo()
		except CancelOperation:
			return
		if self.executionOption.get() == "web":
			scriptLocation = None
			executableLocation = None
			kw = {'licenseKey': self.keyOption.get()}
		else:
			scriptLocation = self.scriptLocationOption.get().strip()
			executableLocation = self.executableLocationOption.get().strip()
			kw = {'customScript': scriptLocation, 'executableLocation': executableLocation}
		from ModellerBase import model, verifyModelKw
		verified, statusMsg = verifyModelKw(kw)
		if not verified:
			self.status(statusMsg, color="red")
			self.enter()
			return
		model(self.mav, seq, templateModels, self.numModelsOption.get(), True, True, False,
			tempPath=self.tempFolderOption.get().strip(), loopInfo=loopInfo, **kw)

	def destroy(self):
		self.mav.triggers.deleteHandler(MOD_ASSOC, self._assocHandlerID)
		self.mav = None
		ModelessDialog.destroy(self)

	def ModellerHomePage(self):
		from chimera import help
		help.display("http://www.salilab.org/modeller/")

	def _areaChoiceCB(self, butName):
		if butName in ("selection", "region"):
			self.mobileEdgesOption.disable()
		else:
			self.mobileEdgesOption.enable()

	def _executionChangeCB(self, opt):
		if opt.get() == "web":
			# forget must be before manage
			self.executableLocationOption.forget()
			self.scriptLocationOption.forget()
			self.scriptButton.grid_remove()
			self.keyOption.manage()
		else:
			# forget must be before manage
			self.keyOption.forget()
			self.executableLocationOption.manage()
			self.scriptLocationOption.manage()
			self.scriptButton.grid()

	def _gatherLoopInfo(self):
		from chimera import UserError

		modelingArea = self.modelingArea.getvalue()
		if not modelingArea:
			raise AssertionError("Nothing chosen to model")

		if modelingArea in ("selection", "region"):
			if modelingArea == "selection":
				region = self.mav.getRegion(self.mav.SEL_REGION_NAME)
				if not region.blocks:
					self.enter()
					raise UserError("%s region empty\n"
						"Select on the structure to designate part(s) to model/remodel"
						% self.mav.SEL_REGION_NAME)
				regName = self.mav.SEL_REGION_NAME
			else:
				region = self.mav.getRegion()
				if not region:
					self.enter()
					raise UserError("No active region\n"
						"Drag out a region on the sequence to designate part(s) of structure to"
						" model/remodel")
				regName = "Active"

			seq = None
			for s1, s2, p1, p2 in region.blocks:
				if seq is None:
					seq = s1
				elif s1 != seq:
					self.enter()
					raise UserError("%s region must contain only one sequence" % regName)
				if s2 != seq:
					self.enter()
					raise UserError("%s region must contain only one sequence" % regName)

			if seq not in self.mav.seqs:
				self.enter()
				raise UserError("%s region must contain a sequence, not a header" % regName)

			matchMaps = getattr(seq, 'matchMaps', None)

			if not matchMaps:
				self.enter()
				raise UserError("%s region's sequence must be associated with a structure"
						% regName)
			elif len(matchMaps) > 1:
				mol = _AskWhichMol(matchMaps.keys()).run(self.uiMaster())
				if not mol:
					raise CancelOperation("multiple mols")
			else:
				mol = matchMaps.values()[0]['mseq'].molecule
		else:
			mol = self.structureMenu.getvalue()
			if not mol:
				self.enter()
				raise UserError("No sequences have associated structures.")
			seq = self.mav.associations[mol]

			matchMaps = seq.matchMaps

		matchMap = matchMaps[mol]

		if modelingArea in ("selection", "region"):
			hasMissing = hasExisting = False
			for s1, s2, p1, p2 in region.blocks:
				for gapped in range(p1, p2+1):
					ungapped = seq.gapped2ungapped(gapped)
					if ungapped in matchMap:
						hasExisting = True
						if hasMissing:
							break
					elif ungapped is not None:
						hasMissing = True
						if hasExisting:
							break
			if hasMissing and hasExisting:
				from chimera.baseDialog import AskYesNoDialog
				if AskYesNoDialog("%s region covers both existing and missing structure.\n"
						"Is that okay?" % regName).run(self.uiMaster()) == "no":
					raise CancelOperation("bad region")

		templateModels = [mol]
		loopData = []
		if modelingArea in ("selection", "region"):
			for s1, s2, p1, p2 in region.blocks:
				for i in range(p1, p2+1):
					start = seq.gapped2ungapped(i)
					if start is not None:
						break
				else:
					continue
				for i in range(p2, p1-1, -1):
					end = seq.gapped2ungapped(i)
					if end is not None:
						break
				loopData.append((start+1, end+1))
			if not loopData:
				raise UserError("%s region covers only gaps" % regName)
		else:
			numAdj = int(self.mobileEdgesOption.get())
			ungappedSeq = seq.ungapped()
			state = "in structure"
			for i in range(len(ungappedSeq)):
				if i in matchMap:
					if state == "in missing":
						end = i
						while end < i+numAdj and end+1 in matchMap:
							end += 1
						if start > 1 or modelingArea != "internal":
							loopData.append((start, end))
						state = "in structure"
				else:
					if state == "in structure":
						start = i+1
						while start > i+1-numAdj and start-2-numAdj in matchMap:
							start -= 1
						state = "in missing"
			if state == "in missing" and modelingArea != "internal":
				loopData.append((start, len(ungappedSeq)))
			if not loopData:
				txt = " internal " if modelingArea == "internal" else " "
				raise UserError("No%sstructure is missing (in parts covered by alignment)"
						% txt)

		return (self.methodOption.get(), loopData), seq, templateModels

	def _structChoiceCB(self, mol):
		if not mol:
			return
		aseq = self.mav.associations[mol]
		matchMap = aseq.matchMaps[mol]
		mseq = matchMap['mseq']
		curRegion = self.mav.currentRegion()
		if None not in mseq.residues:
			selRegion = self.mav.getRegion(self.mav.SEL_REGION_NAME)
			if curRegion == selRegion or (not curRegion and selRegion.blocks):
				self.modelingArea.invoke("selection")
			else:
				self.modelingArea.invoke("region")
			return
		if not curRegion:
			mapping = [matchMap.get(i, None) for i in range(len(aseq.ungapped()))]
			while mapping and mapping[0] == None:
				mapping.pop(0)
			while mapping and mapping[-1] == None:
				mapping.pop()
			if None in mapping:
				self.modelingArea.invoke("internal")
			else:
				self.modelingArea.invoke("missing")
		else:
			self.modelingArea.invoke("region")

	def _viewScript(self):
		# Prepare the Modeller scripts, based on the setting input from the UI.
		loopInfo, seq, templateModels = self._gatherLoopInfo()
		from ModellerBase import writeModellerScripts
		scriptPath, configPath, tmpDir = writeModellerScripts(
			self.keyOption.get().strip(), self.numModelsOption.get(),
			True, True, False, False, loopInfo, self.scriptLocationOption.get().strip(),
			self.tempFolderOption.get().strip(), False, "")
		from Idle import flist
		flist.open_shell().io.open(editFile=scriptPath)

class _AskWhichMol(ModalDialog):
	title = "Choose Molecule to Remodel"
	buttons = ('OK', 'Cancel')

	def __init__(self, possibleMols):
		self.possibleMols = possibleMols
		ModalDialog.__init__(self)

	def fillInUI(self, parent):
		import Pmw
		from chimera.misc import chimeraLabel
		self.molMenu = Pmw.OptionMenu(parent, items=[chimeraLabel(m, modelName=True)
			for m in self.possibleMols], labelpos='n',
			label_text="Active region's sequence is associated with multiple"
				" structures.\nPlease choose which structure you want to remodel:")
		self.molMenu.grid()

	def OK(self):
		import Pmw
		self.Cancel(value=self.possibleMols[self.molMenu.index(Pmw.SELECT)])


class _AskWhichSeq(ModalDialog):
	title = "Choose Sequence to Base Remodeling on"
	buttons = ('OK', 'Cancel')

	def __init__(self, possibleSeqs):
		self.possibleSeqs = possibleSeqs
		ModalDialog.__init__(self)

	def fillInUI(self, parent):
		import Pmw
		self.seqMenu = Pmw.OptionMenu(parent, items=[seq.name for seq in self.possibleSeqs],
			labelpos='n',
			label_text="Multiple sequences in the alignment have associated structures.\n"
				"Please choose which sequence you want to base remodeling on:")
		self.seqMenu.grid()

	def OK(self):
		import Pmw
		self.Cancel(value=self.possibleSeqs[self.seqMenu.index(Pmw.SELECT)])