# FAHProject.py
#
# Daniel L. Ensign
# Pande Group
# Dept. of Chemistry
# Stanford University 
#
# 9/10/07
# 
# This file contains a Python object suitable for representing a Folding@home
# project.
#
# 10/1/07 Added addRun method
#
"""
Instantiation:
	project4100 = FAHProject( "/home/server/server2/projects/project4100.conf" )

Attributes:
conffile - .conf file describing the project. Must be specified when instantiating
	the object.
verbose - whether to print extra info. May be set upon instantiation or after.
conf - a ProjectConfigurationFile object representing the project settings
datapath - a copy of conf.datapath, tells where the data is held.

Methods:
_vmesg - hidden, useful internal shorthand for "if self.verbose: print message"
countReturnedWUs - counts the number of frame*.xtc files in self.datapath directories 
addRun - will add a .gro file to the project. Will not run SetupDataDirs or restart
	the server -- that should be done using the a FAHServer object.
	Usage: project4100.addRun( "newconf.gro" )
addRuns - does addRun on a list of gro files: 
	>> grolist = ("conf1.gro" , "conf2.gro", "conf3.gro" )
	>> project4100.addRuns( grolist )
writeConfig - calls self.conf.write() method for saving a new configuration file.
	Note that configurations are always backed up, and this method will err on
	the side of not writing a new file if it can't back up the older file.
"""

import sys
from glob import glob
from os.path import abspath, exists

sys.path.append( abspath("../adaptiveSampling/") ) # this is wrong
from ProjectConfigurationFile import *
from FAHMisc import partition # should be changed if we go to Python 2.5
from FAHBaseType import FAHBaseType

class FAHProject( FAHBaseType ):

	def __init__ ( self, conffile, verbose = False ):
		self.conffile = conffile
		self.verbose = verbose

		self.conf = ProjectConfigurationFile( self.conffile )
		self.datapath = self.conf.datapath
		
	def countReturnedWUs( self ):
		datapath = self.conf.datapath
		nRuns = self.conf.nruns
		nClones = self.conf.nclones

		nret = 0

		run = 0
		while run < nRuns :
			clone = 0
			while clone < nClones :
				dirname = datapath + "/" + "RUN%d/CLONE%d/" % ( run, clone )
				rc = glob( dirname + "frame*.xtc" )
				nret += len( rc )
				clone += 1
	
			run += 1
		
		return nret

	def addRun( self, gro ):
		# add the .gro file 'gro' to this project.
		# - need to increment NUM_RUNS
		# - need to add to MULTI_GRO

		# if no MULTI_GRO, exit without doing anything!
		if not self.conf.settings[ 'MULTI_GRO' ] :
			self._vmesg( "This project has no MULTI_GRO; cannot add runs" )
			return 

		if not exists( gro ):
			self._vmesg( ".gro '%s' not found" % gro )
			return

		# increment NUM_RUNS
		currentConfig = self.conf.fileContents
		newConfig1 = []
		for line in currentConfig:
			if "NUM_RUNS" in line :
				( setting, sep, comments ) = partition( line, "#" )
	
				numRunsLine = setting.split()
				numRuns = int( numRunsLine[1] )

				newNumRuns = numRuns + 1
				newLine = "NUM_RUNS %d # %s" % ( newNumRuns, comments )
				newConfig1.append( newLine )

			else:
				newConfig1.append( line )

		# add to MULTI_GRO
		newConfig2 = []
		lastline = ""
		for line in newConfig1 :
			if "END_MULTI_GRO" in line :
				# get tpr, etc. info from lastline
				parts = lastline.split()[1:]
				
				newLine = "%s " % gro
				
				for elem in parts :
					newLine += "%s " % elem
				newLine = newLine.strip() + "\n"
				newConfig2.append( newLine )
				newConfig2.append( "END_MULTI_GRO\n" ) 
			else :
				newConfig2.append( line )
				lastline = line

		# now newConfig2 is a list containing new configuration file lines --
		# need to rebuild the configuration object using its rebuild method
		self.conf.fileContents = newConfig2
		self.conf.rebuild()

	def addRuns( self, grolist ):
		for gro in grolist :
			self.addRun( gro )

	def writeConfig( self ):
		self.conf.write()