# ProjectConfigurationFile.py
#
# Daniel L. Ensign
# Pande Group
# Department of Chemistry
# Stanford University 
# Stanford, CA
#
# last modified: Tue, 07 Apr 2008
#
# VV:  I added a write() method....
#
# =================================================================================
# TO DO:
# 1. Continue adding keywords to confkeys
# 2. Deal intelligently with paths. There are two options:
#	- specify the path to the constructor ( server2dir ) -- DONE
#	- get the absolute path of the projectX.conf and infer the server2 dir
#	- add AutoPaths setting ( default = True ). Ignore all path info if set
#		to False. Might break TPR_SCHEME RUN_GROMPP, however -- won't be
#		able to compute the number of frames in a WU if we can't find
#		the mdp file automatically 
# 3. Make printing prettier -- right now we just loop through the keywords in 
#    the settings dictionary, and print key, value except for MULTI_GRO. Booleans
#    force having to replace True with "". 
# ==================================================================================

# a list of 'known' configuration keywords
from confkeys import BooleanKeywords, FileKeywords, ValueKeywords

# objects for understanding keywords
from confkeys import KeyWithValue, KeyBoolean, KeyWithFile, MultiGroFiles

from os.path import abspath
class ProjectConfigurationFile( object ):
	"""
	this class is for representing Folding@home project configurations
	internal to a Python script.

	Use:
		conffilename = "/home/server_b/server2/densign/3036/project3036.conf"
		conffile = ProjectConfigurationFile( conffilename )

	and you can get stuff out of it:

		>> print conffile.nruns
		10
		>> print conffile.datapath
		'/home/server_b/server2/data/PROJ3036'
	
	(both are useful for automatic detection by, eg, analysis scripts). Most elements will
	be stored in the conffile.settings dictionary:

		>>> print conffile.settings["GRO_NAME_DEFAULT"]
		densign/3036/equil0.gro
		>> print conffile.settings["MULTI_GRO"]
		True
		>> print conffile.settings["USE_XTC"]
		None

	for MULTI_GRO, list of gro/top/... files are in conffile.grolist.

	Only those keywords listed in confkeys will be stored as part of the dictionary.
		
	Note about directories: unless you specify the server2dir argument, the path will
	be inferred from the absolute path of the configuration file.

	"""

	# =====================================================================================
	# EXCEPTIONS

	class CannotFindConfiguration( Exception ):
		pass	

	# =====================================================================================
	def __init__( self, filename, server2dir = "", verbose = False, debug = False ):
		self.filename = abspath( filename.strip() )
		self.verbose = verbose
		self.debug = debug

		#try:
		self.fileContents = self.readConfig()
		"""except:
			print "whoops! %s" % self.filename
			raise self.CannotFindConfiguration
		"""
		if server2dir :
			self.path = abspath( server2dir )
		else:
			self.path = abspath( self.filename )[:self.filename.find("server2") ] + "server2"

		# self.settings is a dictionary using FAH keywords as keys
		self.settings = self.parseConfig()
		
		# for booleans, set ones not in the file to False
		self.setBools()

		self.datapath = "%s/data/PROJ%s/" % ( self.path, self.settings[ "PROJECTID" ] )

		# setShortNames() builds several useful attributes so they
		# may be more easily accessed than as parts of self.settings
		self.setShortNames()
		
		self.name = "project configuration '%s'" % self.filename
	
	def __str__( self ):
		return self.name

	def __repr__( self ):
		return self.name

	def setShortNames( self ):
		self.nruns = int( self.settings["NUM_RUNS"] )
		self.nclones = int( self.settings["NUM_CLONES"] )
		self.numdumps = int( self.settings["NUM_DUMPS"] )
		self.projectnumber = int( self.settings[ "PROJECTID" ] )
		self.credit = float( self.settings[ "STATSCREDIT" ] )

	def trimComments( self, confline ):
		# remove '#' and everything after
	 	cutindex = confline.find( "#" )
		return confline[0 : cutindex ].strip()	

	def getNumDumps( self, mdpfile ):
		# method for grabbing NUM_DUMPS from an mdp file if TPR_SCHEME RUN_GROMPP
		MDP = open( mdpfile )
		mdp = MDP.readlines()
		MDP.close()

		nsteps = 0
		nstxtcout = 0	

		for line in mdp :
			line = line[ 0 : line.find(";") ]
			if "nsteps" in line :
				nsteps = int( line.split("=")[1] )
			if "nstxtcout" in line :
				nstxtcout = int( line.split("=")[1] )

		try:
			numdumps = nsteps/ nstxtcout
		except:
			numdumps = None
		
		return numdumps			

	def readConfig( self ):
		FILE = open( self.filename )
		filelines = FILE.readlines()
		FILE.close()
		return filelines

	def setBools( self ):
		keys = self.settings.keys()
		for boolean in BooleanKeywords:
			if boolean not in keys:
				self.settings[ boolean ] = False

	def parseConfig( self ):
		# builds the settings list and other attributes
		# uses stuff from confkeys. Keywords fall into these categories:
		# 	- key with value
		# 	- key boolean
		# 	- key with file (slightly different from key with value)
		# 	- multi gro (extremely special case of key boolean)
		#	- tpr scheme (extremely special case of key with value)

		settings = {}
		multigro = []
		countMultiGro = False

		# this gets set appropriately later, but needs a default
		settings["MULTI_GRO"] = False

		for line in self.fileContents :
		
			trimmedLine = self.trimComments( line )
			linelist = trimmedLine.split()

			try: keyword = linelist[0]
			except: keyword = None


			# keyword tests
			if keyword in BooleanKeywords :
				obj = KeyBoolean( trimmedLine )
				settings[ obj.keyword ] = obj.value
				
			if keyword in FileKeywords :
				obj = KeyWithFile( trimmedLine, self.path )
				settings[ obj.keyword ] = obj.value
		
			if keyword in ValueKeywords :
				obj = KeyWithValue( trimmedLine )
				settings[ obj.keyword ] = obj.value

			if keyword == "END_MULTI_GRO" :
				countMultiGro = False
			if countMultiGro :
				multigro.append( trimmedLine )
			if keyword == "MULTI_GRO" : # already have settings[ "MULTI_GRO" ] == True
				countMultiGro = True

		# build list of gro (and auxiliary) files
		if settings[ "MULTI_GRO" ]:
			self.grolist = MultiGroFiles( multigro, self.path ).filelist
		else:
			self.grolist = () # don't need one if no MULTI_GRO 

		# special case: TPR_SCHEME
		# this is special because we easily determine NUM_DUMPS
		# if RUN_EXTERNAL but not quite so easily if we're using
		# RUN_GROMPP
		# strategy: if RUN_EXTERNAL, we're all set
		# if RUN_GROMPP, use MDP_NAME_DEFAULT and assume it lives
		# in the same directory as the project configuration
                if settings.has_key("TPR_SCHEME"):
		   if settings["TPR_SCHEME"]=="RUN_GROMPP" :
			#print "--%s--" % settings["MDP_NAME_DEFAULT"]
			settings["NUM_DUMPS"] = self.getNumDumps( settings["MDP_NAME_DEFAULT"] )

		return settings

	def configFileString( self ):
		# ugly, needs to be prettier
		file = ""
		keys = self.settings.keys()
		for key in keys:
		
			if self.settings[ key ]:

				if key == "MULTI_GRO" :
					file += self.multiGroString()	
				#elif key == True :
				#	file += "%s\n" % key
				else:
					file +=  "%-20s%-20s\n" % ( key, self.settings[ key ] )

		file = file.replace( "True", "" )

		return file

	def multiGroString( self ):
		mgstring = "MULTI_GRO\n"
		for filedict in self.grolist :
			# print each in order gro, top, xvg
			mgstring += "%s %s %s\n" % ( filedict['gro'], filedict['top'], filedict['xvg'] )
		
		mgstring += "END_MULTI_GRO\n"
		return mgstring 

	# =====================================================================================
	# This function should be last, since it's so long and complicated. How to split it 
	# up into pieces?
	def old_parseConfig( self ):
		multigro = []
		multigroRead = False
		settings = {}
		settings[ "MULTI_GRO" ] = False

		for line in self.fileContents :
			trimmedLine = self.trimComments( line )
			splitLine = trimmedLine.split()

			try: 
				keyword = splitLine[0]
			except:
				keyword = None

			if keyword in confkeys: 
				try :
					argument = splitLine[1]
				except : 
					argument = None

				settings[ keyword ] = argument

			# special case: MULTI_GRO
			# this is not in a very useful format: may be a reason	
			# to read the gro lines more carefully. Perhaps build
			# a dictionary for each line with gro/top/xvg keys.

			if keyword == "MULTI_GRO" : 
				multigroRead = True
				settings[ "MULTI_GRO" ] = True

			if multigroRead :
				if keyword == "END_MULTI_GRO" :
					self.grolist = multigro
					multigroRead = False

				elif keyword != "MULTI_GRO" :
					multigro.append( trimmedLine.split() )
	
		# special case: TPR_SCHEME
		# this is special because we easily determine NUM_DUMPS
		# if RUN_EXTERNAL but not quite so easily if we're using
		# RUN_GROMPP
		# strategy: if RUN_EXTERNAL, we're all set
		# if RUN_GROMPP, use MDP_NAME_DEFAULT and assume it lives
		# in the same directory as the project configuration
		if settings["TPR_SCHEME"]=="RUN_GROMPP" :
			settings["MDP_NAME_DEFAULT"] = self.path + settings[ "MDP_NAME_DEFAULT" ].replace("//", "/")
			settings["NUM_DUMPS"] = self.getNumDumps( settings["MDP_NAME_DEFAULT" ] )
	
		return settings

        def write(self, outfile):
            """Writes the configuration file to filename outfile."""

            fout = open(outfile,'w')
            fout.write(self.configFileString())
            fout.close()