#!/usr/bin/python

#Workflow based on Trajectory Converter - v1.5 by: Justin Lemkul 
#completely reimplemented and improved by Peter Schmidtke & Jesus Seco

import sys,os,re,fnmatch

if len(sys.argv)>4 :
	f=sys.argv[1]
	if not os.path.exists(f):
		sys.exit(" ERROR : Something you provided does not exist. Breaking up.\n\nUSAGE : python trajconv_peter.py amberCrd amberTop trajDir trajPattern outPutPrefix\n\n \
Example : python amber2xtc.py mdcrd.crd mdcrd.top md *.x.gz md_gromacs\n")
else :
	sys.exit(" \n  USAGE : python amber2xtc.py AMBERCRD AMBERTOP TRAJDIR TRAJPATTERN OUTPUTPREFIX\n\
      Example : python amber2xtc.py mdcrd.crd mdcrd.top md *.x.gz md_gromacs\n\
      Note that the AmberCrd can also be a PDB file.\n")

crd=sys.argv[1]
top=sys.argv[2]
trajdir=sys.argv[3]
pattern=sys.argv[4]
outputPref=sys.argv[5]

traj_files=fnmatch.filter(os.listdir(trajdir),pattern) #get the fpocket output folders


RE_DIGIT = re.compile(r'(\d+)')     #set a pattern to find digits
ALPHANUM_KEY = lambda s: [int(g) if g.isdigit() else g for g in RE_DIGIT.split(s)]      #create on the fly function (lambda) to return numbers in filename strings
traj_files.sort(key=ALPHANUM_KEY)      #sort by these numbers in filenames

print "Will convert the following files : "
print traj_files

csn=1
for file in traj_files :
	ptrajtmp=open("ptraj_tmp.ptr","w")
	print "currently converting "+file
	ptrajtmp.write("trajin "+trajdir+os.sep+file+"\n")
	ptrajtmp.write("reference "+crd+"\n")
	#ptrajtmp.write("center ~:WAT,CIO mass origin\n")
	#ptrajtmp.write("image origin center :* byres familiar\n")
	
	ptrajtmp.write("trajout pdb_tmp/mdcrd.pdb pdb")
	ptrajtmp.close()
	if not os.path.exists("pdb_tmp"):
		os.mkdir("pdb_tmp")

	os.system("ptraj "+top +" ptraj_tmp.ptr >& /dev/null")
	if not os.path.exists("xtc_tmp"):
		os.mkdir("xtc_tmp")
	#move to *.pdb
	os.system("cd pdb_tmp; ls *.pdb.* | cut -f3 -d\".\" | awk '{print \"mv mdcrd.pdb.\"$0\" mdcrd_\"$0\".pdb\" }' | sh ; cd ../")	
	pdb_files=fnmatch.filter(os.listdir("pdb_tmp"),"*.pdb")
	pdb_files.sort(key=ALPHANUM_KEY)      #sort by these numbers in filenames	
	if csn==1:
		os.system("editconf -f pdb_tmp/mdcrd_1.pdb -o "+outputPref+"_t1_top.gro >& /dev/null")
	
	for pdb in pdb_files:
		os.system("echo \"0\" | trjconv -s pdb_tmp/"+pdb+" -f pdb_tmp/"+pdb+" -o xtc_tmp/traj_"+str(csn)+".pdb.xtc -t0 "+str(csn)+" >& /dev/null")
		csn+=1
	if os.path.exists(outputPref+"_traj.xtc"):
		os.system("trjcat -f "+outputPref+"_traj.xtc xtc_tmp/*.pdb.xtc -o "+outputPref+"_traj.xtc >& trajcat.log")
	else :
		os.system("trjcat -f xtc_tmp/*.pdb.xtc -o "+outputPref+"_traj.xtc >& trajcat.log")
	os.system("rm -rf pdb_tmp/*.pdb")
	
	os.system("rm -rf xtc_tmp/*.xtc")


os.remove("ptraj_tmp.ptr")
os.system("rmdir pdb_tmp")
os.system("rmdir xtc_tmp")