import omero.scripts as scripts
import omero.util.script_utils as scriptUtil
from omero.rtypes import *
import omero.gateway
import omero_api_Gateway_ice
import omero.util.imageUtil as imgUtil
import subprocess
from os import environ
import numpy
import random
import string

#-------------------------------------------------
#keep these global saves parsing them all the time	
#-------------------------------------------------
gateway = None
session = None
rawPixelStore = None
pixelsService = None
updateService = None
rawFileStore = None
queryService = None
volviewcmd = '/opt/VolViewer/VolViewer'

#-------------------------------------------------
#helper	function to log messages/debug
#-------------------------------------------------
logLines = []
def log(text):
	""" Adds lines of text to the logLines list"""
	print text
	logLines.append(text)
    
def doSomething(images, fps, rotang, step, xaxis, yaxis, zaxis):
	
	log("---") 
	for image in images:		
		
		log("SessionID: %s ImageID: %s" % (client.getSessionId(), image.getId().getValue())) 
			
		#-------------------------------------------------
		#call volviewer and render the movie
		#-------------------------------------------------
		envarg = environ.copy()
		envarg["DISPLAY"] = ":0.0"

		#capture movie settings in filename		
		movie_filename = "RotateMovie_ImageID"+str(image.getId().getValue())+"_angle"+str(rotang)+"_step"+str(step)+".avi"

		#generate unique script filename
		script_filename = "" .join(random.choice(string.ascii_uppercase + string.digits) for x in range(10))
		script_filename = "script_" + str(script_filename) + ".txt"

		#generate unique temp folder filename for movie frames
		movie_tempfoldername = "" .join(random.choice(string.ascii_uppercase + string.digits) for x in range(10))	
		movie_tempfoldername = "RotateMovie_" + str(movie_tempfoldername)+"/"

		#create the temp folder
		cmd = 'mkdir '+ str(movie_tempfoldername)
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )

		#create VolViewer script
		script_file = open("/opt/VolViewer/"+str(script_filename), "w")
		script_file.write("set_hide()\n");
		script_file.write("open_omero_connection('localhost', 4064, '" + str(client.getSessionId()) + "')\n")
		script_file.write("open_omero_image("+str(image.getId().getValue())+")\n")
		script_file.write("movie_rotate("+str(step)+", '/opt/VolViewer/"+str(movie_tempfoldername)+"', "+str(xaxis)+", "+str(yaxis)+", "+str(zaxis)+", "+str(rotang)+")\n")
		script_file.write("close_omero_connection()\n")
		script_file.write("set_close()");
		script_file.close()

		#execute the script
		cmd = volviewcmd + ' "open_script(/opt/VolViewer/'+str(script_filename)+')"';
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )
		
		#-------------------------------------------------
		#encode the movie file
		#-------------------------------------------------
		cmd = 'mencoder mf:///opt/VolViewer/'+str(movie_tempfoldername)+'*.png -o '+str(movie_filename)+' -ovc lavc -oac lavc'
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )

		#-------------------------------------------------
		# uploads the file to the server, attaching it to the 'parent' Project/Dataset as an OriginalFile annotation,
		# with the figLegend as the description. Returns the id of the originalFileLink child. (ID object, not value)
		#-------------------------------------------------
		format = "movie/avi"
		output = str(movie_filename)
		fileAnnotation = scriptUtil.uploadAndAttachFile(queryService, updateService, rawFileStore, image, output, format, "Volume rendering.")

		#-------------------------------------------------
		#be clean and delete the:
		# movie_tempfoldername, 
		# script_filename, 
		# movie_filename
		#-------------------------------------------------
		cmd = 'rm -rf /opt/VolViewer/'+str(movie_tempfoldername)
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )

		cmd = 'rm -f /opt/VolViewer/'+str(script_filename)
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )

		cmd = 'rm -f ./'+str(movie_filename)
		run = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=envarg)
		returncode = run.wait()
		stdout = run.stdout.readlines()
		stderr = run.stderr.readlines()		
		log("stdout: %s stderr: %s" % (stdout,stderr) )

	log("---")

def parseCommandArguments(commandArgs):    
	log("---")

	parent = None
	imageIds = []
	datasetIds = []
	projectIds = []

	#-------------------------------------------------
	#parse IDs for projects, datasets, images 
	#-------------------------------------------------
	dataType = commandArgs["Data_Type"]
	fps = commandArgs["FPS"]
	rotang = commandArgs["Rotation_Angle"]
	step = commandArgs["Angular_Step_Size"]
	xaxis = commandArgs["X-axis"]
	yaxis = commandArgs["Y-axis"]
	zaxis = commandArgs["Z-axis"]

	log("dataType: %s" % dataType);
	log("---");
	if dataType == "Image": 						#IMAGES
		for imageId in commandArgs["IDs"]:
			try:
				iId = long(imageId.getValue())
				imageIds.append(iId)
			except: pass
	elif dataType == "Dataset": 					#DATASETS
		for datasetId in commandArgs["IDs"]:
			try:
				dId = long(datasetId.getValue())
				datasetIds.append(dId)
			except: pass
	else: 											#PROJECTS
		for projectId in commandArgs["IDs"]:
			try:
				pId = long(projectId.getValue())
				projectIds.append(pId)
			except: pass
	log("---")
	
	#-------------------------------------------------
	#check actually we have something
	#-------------------------------------------------
	if len(imageIds) == 0 and len(datasetIds) == 0 and len(projectIds) == 0:
		print "No image IDs, dataset IDs found or proeject IDs found."       
	log("Found: %d projects, %d datasets, %d images" % (len(projectIds), len(datasetIds), len(imageIds)))
	
	#-------------------------------------------------
	#now retrieve all the imageIDs we wish to process
	#-------------------------------------------------
	
	#PROJECTS
	projects = []
	if len(projectIds) != 0:
		projects = gateway.getProjects(projectIds, False)

	for project in projects:
		if project == None:
			print "No project found for ID: %s" % projectId
			continue
		projectName = project.getName().getValue()
		images = gateway.getImages(omero.api.ContainerClass.Project, [project.getId().getValue()])
		log("Project: %s     ID: %d     Images: %d" % (projectName, project.getId().getValue(), len(images)))
		#-------------------------------------------------
		doSomething(images, fps, rotang, step, xaxis, yaxis, zaxis)
		
	#DATASETS
	for datasetId in datasetIds:
		dataset = gateway.getDataset(datasetId, False)
		if dataset == None: 
			print "No dataset found for ID: %s" % datasetId
			continue
		datasetName = dataset.getName().getValue()
		images = gateway.getImages(omero.api.ContainerClass.Dataset, [datasetId])
		log("Dataset: %s     ID: %d     Images: %d" % (datasetName, datasetId, len(images)))
		#-------------------------------------------------
		doSomething(images, fps, rotang, step, xaxis, yaxis, zaxis)

	#IMAGES
	if len(datasetIds) == 0:
		images = []
		for imageId in imageIds:
			images.append(gateway.getImage(imageId))
		#-------------------------------------------------
		doSomething(images, fps, rotang, step, xaxis, yaxis, zaxis)
	
if __name__ == '__main__':
	dataTypes = [rstring('Project'),rstring('Dataset'),rstring('Image')]

	client = scripts.client('Rotate_Movie', """Generates a rotation moview for all image(s) selected, in a dataset or in a project""",
	
	scripts.String("Data_Type", optional=False, grouping="1",
		description="The data you want to work with.", values=dataTypes, default="Image"),
	
	scripts.List("IDs", optional=False, grouping="2",
		description="List of Project IDs, Dataset IDs or Image IDs").ofType(rlong(0)),
	
	scripts.Int("FPS", optional=False, grouping="3",
		description="Frames per second", default=10),	

	scripts.Float("Rotation_Angle", optional=False, grouping="4",
		description="The angle of rotation", default=45.0),	

	scripts.Int("Angular_Step_Size", optional=False, grouping="5",
		description="The angle step size", default=2),	

	scripts.Int("Z-axis", optional=False, grouping="6",
		description="The angle step size", default=0),	
	scripts.Int("Y-axis", optional=False, grouping="7",
		description="The angle step size", default=1),	
	scripts.Int("X-axis", optional=False, grouping="8",
		description="The angle step size", default=0),	

	version = "1.01",
	authors = ["Jerome Avondo", "JIC"],
	institutions = ["John Innes Centre"],
	contact = "jerome.avondo@bbsrc.ac.uk",
	)

	try:
	
		session = client.getSession()
	
		gateway = session.createGateway()
		rawPixelStore= session.createRawPixelsStore()
		pixelsService= session.getPixelsService()
		queryService = session.getQueryService()
		updateService = session.getUpdateService()
		rawFileStore = session.createRawFileStore()
		
		commandArgs = {}
		
		for key in client.getInputKeys():
			if client.getInput(key):
				commandArgs[key] = client.getInput(key).getValue()
		
		log("---");
		log("commands: %s" % commandArgs)
		
		parseCommandArguments(commandArgs)
		
		client.setOutput("Message", rstring("Script was successfull"))
		
	finally:	
		client.closeSession()