#!/usr/bin/python

# NOTE: one would better use random.sample to subsample a db result ...

# =================================================================================
# Python builtins
import sys
from os import mkdir

# FAHTools
sys.path.append( ".." )
from FAHdb import FAHdb
from FAHMisc import UniformlyDistributedInteger 

# mmtools - for xtc2gro and GromacsStructures
sys.path.append( "/home/server_c/modules/mmtools/gromacstools" )
from GromacsData import xtc2gro

# =================================================================================
# Settings ... have fun with this!
database = "project4010"
datadir = "/home/server_c/server2/data/PROJ4010"
workDirName = "work"
def buildqlist( db ):
	"""
	qbase = "select run, clone, gen, snapshot from frames where %s=%d"
	defs = ("def1", "def2" )
	nums = (0, 1)
	for d in defs :
       		for num in nums :
                	db.addquery( qbase % ( d, num ) )
	"""
	qbase = "select run, clone, gen, snapshot from frames where rmsd<=%f and rmsd>%f"
	for j in range( 12 ):
		db.addquery( qbase % ( j, j+1 ) )
		if j > 3 and j < 7 :
			db.addquery( qbase % ( j, j + 1 ) )
	return db

# =================================================================================
# working directory -- for intermediate files
# things will be overwritten
try:
	mkdir( workDirName )
except:
	pass

# =================================================================================
# Build the database connection and query list
db = FAHdb( database )
db = buildqlist( db )

# =================================================================================
# The query loop
rcgsList = []
while True:
	data = db.query()
	if data:
                maxval = len( data ) - 1
                randomFrame = UniformlyDistributedInteger( min = 0, max = maxval )
                rcgs = data[ randomFrame ]
                rcgsList.append( rcgs )
	else:
		break

# =================================================================================
# convert everyone to a gro file, store in workDirName directory
for rcgs in rcgsList :
	( run, clone, gen, snapshot ) = rcgs

	rcdir = "RUN%d/CLONE%d" % ( run, clone )
	dir = "%s/%s" % ( datadir , rcdir )

	xtcFile = "%s/frame%d.xtc" % ( dir, gen )
	tprFile = "%s/frame0.tpr" %  dir

	gro = xtc2gro( xtcFile, tprFile, snapshot, "echo 0 0" )
	fname = "%d-%d-%d-%d.gro" % rcgs
	name = "%s/%s" % ( workDirName, fname )
	print "Writing %s" % name
	gro.write( name )