#!/bin/env python
# -*- coding: utf-8 -*-
#
#--------------------------------------------------------------------------------
#cattraj.py v1.1, Copyright Bjoern Olausson
#--------------------------------------------------------------------------------
#This program is free software; you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation; either version 2 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#GNU General Public License for more details.
#
#To view the license visit
#http://www.gnu.org/licenses/old-licenses/gpl-2.0.html
#or write to
#Free Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
#--------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
#
# cattraj.py is intended to be used on NAMD 2.8 colvars trajectory files. It's usage is similar to catdcd and the default behavior
# is to concatenate colvar trajectory files from consecutive runs. As with catdcd it is able to skip frames and thus reduce the time resolution.
# The first file read will always contain one more frame then the following files because it either includes frame zero or the last step from the previous file

import sys, glob, re, os
from optparse import OptionParser, OptionGroup

usage = '''cattraj.py is intended to be used on NAMD 2.8 colvars trajectory files. It's usage is similar to catdcd and the default behavior
is to concatenate colvar trajectory files from consecutive runs. As with catdcd it is able to skip frames and thus reduce the time resolution.
The first file read will always contain one more frame then the following files because it either includes frame zero or the last step from the previous file'''

opa = OptionParser(usage=usage)

opa.add_option("-i", "--in", action="append", dest="infiles",  type="string", metavar="STR*",
	help="Files to process. If you use wild-cards, quote the sting. The resulting list is sorted to comply with the human sort order expectation. This option can be used multiple times.")
opa.add_option("-o", "--out", action="store", dest="outfile", default="cattraj.out.traj", type="string", metavar="STR*",
	help="Filename for the new reduced/concatenated trajectory")
opa.add_option("-f", "--first", action="store", dest="first", default=1, type="int", metavar="INT",
	help="Start with frame. The frames are listed consecutively starting from 1")
opa.add_option("-l", "--last", action="store", dest="last", default=-1, type="int", metavar="INT",
	help="Stop at frame. The frames are listed consecutively starting from 1")
opa.add_option("-s", "--stride", action="store", dest="stride", default=1, type="int", metavar="INT",
	help="Number of frames to skip")

# Print the help if no args are supplied
if len(sys.argv) == 1:
	opa.print_help()
	sys.exit()

# Instruct optparse to parse the program's command line
(options, args) = opa.parse_args()

infiles = options.infiles
outfile = options.outfile
first =  options.first
last = options.last
stride = options.stride

if last < first and last != -1:
	print " Last must be >= first."
	sys.exit(2)

def sortnicely(l):
	""" Sort the given list in the way that humans expect.
	http://www.codinghorror.com/blog/2007/12/sorting-for-humans-natural-sort-order.html"""
	convert = lambda text: int(text) if text.isdigit() else text
	alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key)]
	l.sort( key=alphanum_key )
	return l

filelist_unsorted = []
for args in infiles:
	expanded = glob.glob(args)
	filelist_unsorted.extend(expanded)

filelist = sortnicely(filelist_unsorted)

# Check if the output file exists
try:
	open(outfile, "r")
except IOError:
	pass
else:
	print "File %(saveas)s exists!" %{"saveas": outfile}
	sys.exit(1)

fo = open(outfile, "a")

counter = 1
global_file_counter_w = 0
write_next_line = first
last_step = -1
done = False
# Act on all files found by glob
for f in filelist:
	per_file_counter_r =   0
	per_file_counter_w =   0
	# Open the trajectory file read only
	fname = os.path.realpath(f)
	#fname = os.path.basename(f)
	print "\nReading %(file)s" %{"file": fname}
	trj = open(f, "r")
	# Iterate trough all lines in the trajectory file
	for line in trj:
	# Write the comment only once
		if line.startswith("#") and counter == 1:
			fo.write(line)

		# Skip commented lines
		if not "#" in line and line != "\n":
			# Skip a line if the timestep is identical - This happens when you concatenate files from consecutive runs.
			# The last timestep in the trajectory file N is repeated in the trajectory file N+1 from the following run
			current_step =  int(line.split()[0])
			if current_step != last_step:
				# Only write lines matching the stride counter
				if counter == write_next_line:
					fo.write(line)
					per_file_counter_w += 1
					write_next_line += stride
			counter += 1
			per_file_counter_r += 1
		# Exit if we reached the last frame specified by --last option.
		if counter == last + 1:
			done = True
			break
	# Save the last timestep from the current file to compare it with the first timestep of the next file
	last_step = int(line.split()[0])
	global_file_counter_w += per_file_counter_w

	print "Frames - r: %(frms)10i" %{"frms": per_file_counter_r}
	print "Frames - w: %(written)10i" %{"written": per_file_counter_w}

	if done:
		break

print ""
print "-----------------------------------"
print "-----------------------------------"
print "Overall frames - r: %(frms)10i" %{"frms": counter - 1}
print "Overall frames - w: %(written)10i" %{"written": global_file_counter_w}