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, sigma, dt, iters):
	
	log("---")
	for image in images:		
		
		log("SessionID: %s ImageID: %s" % (client.getSessionId(), image.getId().getValue())) 
			
		#-------------------------------------------------
		#call volviewer and render the image
		#-------------------------------------------------
		envarg = environ.copy()
		envarg["DISPLAY"] = ":0.0"

		#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"

		#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("compute_filter_AnisotropicDiffusion(0, "+str(sigma)+", "+str(dt)+", "+str(iters)+")\n")
		script_file.write("compute_filter_AnisotropicDiffusion(1, "+str(sigma)+", "+str(dt)+", "+str(iters)+")\n")
		script_file.write("compute_filter_AnisotropicDiffusion(2, "+str(sigma)+", "+str(dt)+", "+str(iters)+")\n")
		script_file.write("save_omero_image('Anisotropic Difusion Filtered - sigma:"+str(sigma)+" dT:"+str(dt)+" iterations:"+str(iters)+" Parent Image ID:"+str(image.getId().getValue())+"')\n")
		

		script_file.write("close_omero_connection()\n")
		script_file.write("set_close()\n");
		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) )

		#-------------------------------------------------
		#be clean and delete the script_filename
		#-------------------------------------------------
		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) )

	log("---")

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

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

	#-------------------------------------------------
	#parse IDs for projects, datasets, images 
	#-------------------------------------------------
	dataType = commandArgs["Data_Type"]
	log("dataType: %s" % dataType);

	sigma = commandArgs["Sigma"]
	dt = commandArgs["dT"]
	iters = commandArgs["Iterations"]

	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, sigma, dt, iters)
		
	#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, sigma, dt, iters)

	#IMAGES
	if len(datasetIds) == 0:
		images = []
		for imageId in imageIds:
			images.append(gateway.getImage(imageId))
		#-------------------------------------------------
		doSomething(images, sigma, dt, iters)
	
if __name__ == '__main__':
	dataTypes = [rstring('Project'),rstring('Dataset'),rstring('Image')]

	client = scripts.client('Anisotopic_Diffusion', """Filters 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.Float("Sigma", optional=False, grouping="3",
		description="The difussion coefficient", default=5.0),	
	scripts.Float("dT", optional=False, grouping="4",
		description="The diffusion rate", default=0.1),	
	scripts.Int("Iterations", optional=False, grouping="5",
		description="The number of iterations", default=5),	
	
	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()
