#!/usr/bin/env python

# Author: Greg Alushin, galushin@rockefeller.edu

# This script is meant to align movies being generated in a directory during data acquisition
# It assumes that movies are being saved as compressed tif files in superresolution mode on the K2
# Dependencies: IMOD, FREALIGN (for resample.exe), UNBLUR, SUMMOVIE, EMAN2 (for generating filtered jpgs),
# proc2D for flipping gain reference, and the corresponding shell script alignOnTheFly_Unblur.csh in your path
#
#
#import modules
#
import os
import subprocess
import sys
import optparse
import time

#
#========================
def setupParserOptions():
	parser = optparse.OptionParser()
	parser.set_usage("%prog --apix <angstroms per pixel superres> --rate <dose rate> --kV <voltage>")
	parser.add_option("--apix",dest="apix",type="float",metavar="float",
		help="Angstroms per pixel (superresolution!)")
	parser.add_option("--rate",dest="rate",type="float",metavar="float",
		help="Dose rate per frame")
	parser.add_option("--kV",dest="kV",type="int",metavar="int",
		help="Accelerating voltage") 
	parser.add_option("--patch",dest="patch",type="int",metavar="int",default="0",
		help="Number of patches (single integer, default 0, full-frame alignment)")
	parser.add_option("--bfact",dest="bfact",type="float",metavar="float",default="100.0",
		help="Bfactor (default 100.0)")
	parser.add_option("--gpu",dest="gpu",type="string",metavar="string",default="0",
		help="Gpu to use. Default = 0. For multiple GPUs, enclose in single quotes and separate by space. E.G. '0 1'")
	parser.add_option("--group",dest="group",type="int",metavar="int",default="1",
		help="Number of frames to average before alignment. Default 1, no averaging.")
	parser.add_option("--throw",dest="throw",type="int",metavar="int",default="0",
		help="Number of frames from the beginning of movie to discard from average. Default 0, none discarded.")
	parser.add_option("--trunc",dest="trunc",type="int",metavar="int",default="0",
		help="Number of frames from the end of movie to discard from average. Default 0, none discarded.")
	options,args = parser.parse_args()

	if len(args) > 0:
		parser.error("Unknown commandline options: " +str(args))

	if len(sys.argv) < 7:
		parser.print_help()
		sys.exit()
	params={}
	for i in parser.option_list:
		if isinstance(i.dest,str):
			params[i.dest] = getattr(options,i.dest)
	return params
#=========================
def checkExistence(f_name):
	if not os.path.exists(f_name):
		print "Error: file %s does not exist.  Exiting..."%f_name
		sys.exit()
#=========================
def transferDone(f):
	#check last time file was modified verse system clock
	#if modified more than 10 seconds ago, return True
	#else, return False	
	return ((time.time()-os.stat(f).st_mtime) > 10)
#==========================	
def findNewFiles(donelist):
	#find TIF files in the directory
	TIFs = [s for s in os.listdir('.') if s[-4:]==".tif"]
	#compare to list of done files and return new files
	newTIFs = list(set(TIFs)-set(donelist))
	#simple sort command, which just works, thanks Python!	
	newTIFs.sort()	
	return newTIFs
#=========================
def setUpRefs():
	#prepare gain reference from dm4
	cmd1 = "ln -s *.dm4 reference.dm4"
	subprocess.Popen(cmd1,shell=True).wait()
	cmd2 = "dm2mrc reference.dm4 reference.mrc"
	subprocess.Popen(cmd2,shell=True).wait()
	cmd3 = "proc2d reference.mrc reference_flip.mrc flip"
	subprocess.Popen(cmd3,shell=True).wait()
	cmd4 = "newstack -rot -90 reference_flip.mrc reference_fliprot90.mrc"
	subprocess.Popen(cmd4,shell=True).wait()

	#link defects file
	#cmd5 = "ln -s defects*.txt defects.txt"
	#subprocess.Popen(cmd5,shell=True).wait()

	#defects file is not used by MotionCor2

#===========================
if __name__ == "__main__":
	params=setupParserOptions()
	
	#only make refs if they haven't been generated already
	if not os.path.isfile('reference_fliprot90.mrc'):	
		setUpRefs()
	#make sure defects file and gain reference are present
	checkExistence('reference_fliprot90.mrc')
	
	#checkExistence('defects.txt')
	#defects currently not used by motioncor2

	#check to see if donelist exists: if not initiate blank done list
	if os.path.isfile('done_list.txt'):
		with open('done_list.txt') as f1:
			doneList = f1.read().splitlines()
	else:
		doneList=[]
	
	#create directories for receiving files to be created.  Only do this if they don't already exist.
	if not os.path.exists('aligned_doseweight'):
		os.mkdir('aligned_doseweight')
		os.mkdir('aligned_nodoseweight')
		os.mkdir('jpgs')

	
	#see if new files have arrived
	filesToCheck = findNewFiles(doneList)
	#make sure those files are done
	processList=[]		
	for f in filesToCheck:
		#if transferDone(f):
		#This does not work with current transfer method due to time stamp issues
		processList.append(f)
	#make sure there are files to process, if not reinitiate loop
	if not processList:
		print "All tif files in this directory have already been aligned. Exiting..."
		sys.exit()
		#sequentially process every file in process list with shell script, then add to done list
	for f in processList:
		cmd="alignOnTheFly_MotionCor2.csh %s reference_fliprot90.mrc %s %s %s %s %s '%s' %s %s %s"%(str(f),str(params['apix']),
			str(params['rate']),str(params['kV']),str(params['bfact']),str(params['patch']),str(params['gpu']),str(params['group']),
			str(params['throw']),str(params['trunc']))
		subprocess.Popen(cmd,shell=True).wait()
		doneList.append(f)
		with open('done_list.txt','a') as doneListFile:		
			doneListFile.write("%s\n"%f)