from glob import glob
from os import chdir, mkdir, popen, popen2
from os.path import basename
import subprocess
from sys import exit
from time import time

# base class for points determiner
# to do:
# - write a class derived from PointsDeterminer, replace 'modify' so that the mdp file is modded instead.
class PointsDeterminer :
	"""
	ARGUMENTS:
	
	(required) server - server name, and maybe user name, eg, "server@vspmf93", "vspmf93"
	(required) projectdir - the place on the remote server where your project *data* resides
	(required) projectconf - the location on the remote server of projectX.conf file
	(required) core - the name and local location of the core appropriate for your project

	(recommended) workdir - a bit of a fudge, I admit, but your projectX.conf file probably 
				lists files with their paths relative to server2 -- this is where you
				provide that path *up to* the actual directory where your files are.
				For example:
					1. project2170.conf is in server2/dan/2170
					2. project2170.conf refers to dan/2170/restrt1.inc, etc.
				so:
					server="vspmf93"
					projectdir="server2/dan/2170"
					projectconf="server2/dan/2170/project.conf"
					workdir="dan"
				... I hope that makes sense.

	makeaunitbin - the server program that can build .dat files that your core can read. Default 
				"./MakeAUnit" will probably suffice 
	fraction - 1/fraction is the part of the WU you will run, ie, fraction=20 means run 5%. Default
				20
	verbose - you'll get a lot of info anyway, but why not?	
	"""
	def __init__( self, server, projectdir, projectconf, core, workdir="project", makeaunitbin="./MakeAUnit", fraction=20, verbose=False ):
		self.server=server
		self.projectdir=projectdir
		self.projectconf=projectconf

		self.core=core
		self.workdir=workdir
		self.makeaunitbin=makeaunitbin
		self.fraction=float( fraction )
		self.verbose=verbose
		
		try: mkdir( self.workdir )
		except: pass

		try: mkdir( self.refdir )
		except: pass

	def _vmesg( self, message ):
		if self.verbose: print "PointsDeterminer: %s" % message

	def _do_cp( self, source, target ):
		self._do_popen( "cp %s %s" % ( source, target ) )

	def _do_scp( self, source, target ):
		self._do_popen( "scp -r %s %s" % ( source, target ) )

	def _do_mv( self, source, target ):
		self._do_popen( "mv %s %s" % ( source, target ) )

	def _do_popen( self, command, special=False ):
		if special :
			d = popen( "pwd" )
			print d.read()
			print command
			subprocess.call( command.split() )
		else :
			result=popen( command )
			result=result.read()
			if result: self._vmesg( "returned %s when running '%s'" % ( result, command ) )

	def setup( self ) : 
		# scp the project configuration
		# note this will overwrite a local copy
		self._vmesg( "copying '%s' from server" % self.projectdir )
		remotepath ="%s:%s" % ( self.server, self.projectdir )
		self._do_scp( remotepath, self.workdir )

		self._vmesg( "copying '%s' from server" % self.projectconf )
		remotepath="%s:%s" % ( self.server, self.projectconf )
		self._do_scp( remotepath, self.workdir )

		self.projectconf="%s/%s" % ( self.workdir, basename( self.projectconf ) )
		
		# write a project.conf that refers to self.project
		FILE=open( "project.conf","w" )
		FILE.write( "SERVER_CONF %s ASSIGN_AND_ACCEPT 1\n" % self.projectconf )

	def modify( self ):
		# change MAX_ITER in projectX.conf
		FILE=open( self.projectconf )
		newconf = ""
		lines = FILE.readlines()
		FILE.close()

		for line in lines :
			if "MAX_ITER" in line :
				niter = float( line.split()[1] )
				niter = int( niter/self.fraction )
				newconf += "MAX_ITER %d # changed by script\n" % niter
			else :
				newconf += line
		self._do_mv( self.projectconf, self.projectconf+".old" )
		FILE=open( self.projectconf, "w" )
		FILE.write( newconf )
		FILE.close()

	def makeaunit( self ):
		self._vmesg( "running %s" % self.makeaunitbin )
		self._do_popen( self.makeaunitbin )
		self._do_cp( self.core, self.workdir )
		try: mkdir( "%s/work" % self.workdir )
		except: pass
		fahwork = "%s/work/wudata_01.dat" % self.workdir 
		self._do_mv( "test.dat", fahwork )

	def benchmark( self ):
		# here's the fun part ...
		chdir( self.workdir )
		corename = self.core.split("/")[-1]
		t0 = time()
		self._do_popen( "./%s -dir work -suffix 01 -echo" % corename, special=True )
		return time()-t0	

class AmberPointsDeterminer( PointsDeterminer ):
	def modify( self ):
		# change nstlim in the md input file
		# 1. get the md.inp file
		dir = self.projectdir.replace("server2/","")
		inp = glob( "%s/*inp" % dir )[0]
		self._vmesg( "using and modifying '%s'" % inp )
		FILE=open( inp )
		newinp=""
		lines = FILE.readlines()
		FILE.close()
		for line in lines :
			if "nstlim" in line :
				parts = line.strip().split( "," )
				for part in parts :
					if "nstlim" in part :
						niter = float( part.split("=")[1] )
						niter = int( niter/self.fraction )
						newinp += "nstlim=%d," % niter
					elif part :
						newinp += "%s, " % part

				newinp += "\n"
			else :
				newinp += "%s\n" % line.strip() 

		print "check for errors:"
		print newinp
		self._do_mv( inp, inp+".old" )
		FILE=open( inp, "w" )
		FILE.write( newinp )
		FILE.close()