#!/usr/bin/python

###############################

    #Compute average layer v1.1
    #python gimp plugin
    #2008, Elmar Hoefner
    #Licensed under GPLv3
    #see www.fsf.org for details

###############################

# Version history
# [V1.0 - not officially numbered]  First release
# V1.1 - Work on (rectangular) selections


		#In a 4d array (like superarray) access layers, rows, columns and pixels as follows:
		#array[:,x] = Row x in all layers
		#array[:,:,y] = Column y in all layers
		#array[:,x,y] = Pixel xy in all layers
		#array[:,x,y,rgba] = Channel rgba (0-3) in all layers

		# to access meanarray colors: 
		# meanarray.T[0] = R
		# meanarray.T[1] = G
		# meanarray.T[2] = B
		# meanarray.T[3] = A
		
		# If missing an alpha channel is added
import scipy
from scipy.stats.stats import mode, gmean, hmean, median
from gimpfu import *
import os
from math import floor

class Averager:
	def __init__(self, image, meantype, as_new_image, only_visible_layers, cutoff, side):
		"""Constructor"""
		self.image=image	
		self.meantype=meantype
		self.as_new_image=as_new_image
		if only_visible_layers:
			self.layerlist=[]
			for layer in image.layers:
				if layer.visible:
					self.layerlist.append(layer)
			#if self.layerlist==[] # no layer is visible, choose all
				#self.layerlist=img.layers
		else: self.layerlist=image.layers
		
		# Enough layers must be left after cutting extreme values...
		self.values_to_cut=int(floor(cutoff*len(self.layerlist)/100.0))
		self.cutoff_side=side
		
		# Create a value to check how many layers/values are left:
		if self.cutoff_side < 2: 
			remove = 1
		else: remove = 2
		self.values_used=len(self.layerlist)-remove*self.values_to_cut
		
		#image format or selection format, if one:
		non_empty, x1, y1, x2, y2 = pdb.gimp_selection_bounds(image)
		self.selection_width=x2-x1
		self.selection_height=y2-y1
		self.image_width=image.width
		self.image_height=image.height
		self.origin_x=x1
		self.origin_y=y1
		self.endpoint_x=x2
		self.endpoint_y=y2
		
		self.bpp=4 # always assume alpha (see read_in!)	
		#self.sarray=None # sarray is created by calling self.superarray()
		# How much should the progress indicator advance?
		self.step=1/(self.values_used)
		
	def read_in(self,drawable):
		"""Convert pixel region to scipy array. Taken (& modified) from John Fremlin's retinex implementation."""		
		# always add alpha to save hassles later...
		drawable.add_alpha()
		#get a pixel region
		pr = drawable.get_pixel_rgn(self.origin_x, self.origin_y, self.selection_width, self.selection_height, False)
		# Convert pr to array, float is needed to keep precision in calculations
		a = scipy.fromstring(pr[self.origin_x:self.endpoint_x,self.origin_y:self.endpoint_y],"B")
		assert(a.size == self.selection_width * self.selection_height * self.bpp)
		imagearr = scipy.array(a.reshape(self.selection_height,self.selection_width,self.bpp),scipy.float32)[:,:,0:self.bpp]
		return imagearr
	
	def write_out(self,drawable, imagearr):
		"""Convert scipy array to pixel region. Taken (& modified) from John Fremlin's retinex implementation."""
		# Convert back to unsigned integer
		byte_image = scipy.array((imagearr).round(0),scipy.uint8)
		self.pr = drawable.get_pixel_rgn(self.origin_x, self.origin_y, self.selection_width, self.selection_height, True)
		assert(byte_image.size == self.selection_width * self.selection_height * self.bpp)
		# pixel region needs absolute coordinates!
		self.pr[self.origin_x:self.endpoint_x,self.origin_y:self.endpoint_y] = byte_image.tostring()
		
	def superarray(self):
		"""Create a 'superarray' containing all layers in an additional axis"""
		# axis (0) is layer index		
		# create a superarray for all layers, always assume alpha channel (-> '4')
		# acccess the full layer on axis 0	
		sarr=scipy.zeros((len(self.layerlist),self.selection_height,self.selection_width,4))
		gimp.progress_init("Reading layers...")
		local_step=self.step
		for layer in self.layerlist:
			# read each layer into a subarray of superarray
			sarr[self.layerlist.index(layer)]=self.read_in(layer)
			gimp.progress_update(self.step)
			local_step+=self.step		
		self.sarray=sarr
		return 
	
	def winsorize_superarray(self):
		"""Cut extreme values from color channels"""
		gimp.progress_init("Winsorizing...")
		self.sarray.sort(0)
		if self.cutoff_side==2:
			self.sarray=self.sarray[self.values_to_cut:-self.values_to_cut]
		elif self.cutoff_side==0:
			self.sarray=self.sarray[self.values_to_cut:]
		elif self.cutoff_side==1:
			self.sarray=self.sarray[:-self.values_to_cut]
		gimp.progress_update(1)
		return		
		
	def calculate_mean(self):
		"""Return mean array depending on mean type"""
		
		
		if self.values_to_cut>0: self.winsorize_superarray()
		
		if self.meantype == "arith":
			self.layername="Arithmetical Mean"
			gimp.progress_init("Calculating "+self.layername)
			self.meanarray=scipy.mean(self.sarray,0)
			gimp.progress_update(1)
			return 
		elif self.meantype == "geom":
			self.layername="Geometrical Mean"
			gimp.progress_init("Calculating "+self.layername)
			self.meanarray=gmean(self.sarray)
			gimp.progress_update(1)
			return 
		elif self.meantype == "harm":
			self.layername="Harmonic Mean"
			gimp.progress_init("Calculating "+self.layername)
			self.meanarray=hmean(self.sarray)
			gimp.progress_update(1)
			return 
		elif self.meantype == "median":
			self.layername="Median"
			gimp.progress_init("Calculating "+self.layername)
			self.meanarray=median(self.sarray)
			gimp.progress_update(1)
			return 
		elif self.meantype == "mode":
			self.layername="Mode"
			gimp.progress_init("Calculating "+self.layername)
			self.meanarray=mode(self.sarray)[0]
			gimp.progress_update(1)
			return
			
			
	def export(self):
		"""Export new layer or new image"""
		gimp.progress_init("Exporting...")
		if self.as_new_image:
			new_image=gimp.Image(self.image_width,self.image_height,RGB)
			# Convenience: Set filename of new image to layer name and use same file extension & path as original file
			if self.image.filename: # this only works if the original image has a file name, i.e. has been saved earlier
				path=os.path.split(self.image.filename)[0]
				ext=self.image.filename.split('.')[-1]
				new_image.filename=os.path.join(path,self.layername+'.'+ext)
			else:
				new_image.filename=self.layername
			#Create layer, image and display
			meanlayer=gimp.Layer(new_image, self.layername, self.image_width, self.image_height)
			meanlayer.add_alpha()
			#fill with transparency
			pdb.gimp_drawable_fill(meanlayer, 3)
			self.write_out(meanlayer, self.meanarray)
			new_image.add_layer(meanlayer)
			disp=gimp.Display(new_image)
		else:
			meanlayer=gimp.Layer(self.image, self.layername, self.image_width, self.image_height)
			meanlayer.add_alpha()
			#fill with transparency
			pdb.gimp_drawable_fill(meanlayer, 3)
			self.write_out(meanlayer, self.meanarray)
			self.image.add_layer(meanlayer)
			self.image.raise_layer_to_top(meanlayer)
		gimp.progress_update(1)
		return

def do_the_work(image, drawable, meantype, as_new_image, only_visible_layers, cutoff, side):
	"""Main function, should be named 'main', but this name is already used by the gimp-python interface"""
	working_class=Averager(image,meantype,as_new_image,only_visible_layers, cutoff, side)
	pdb.gimp_message("%i out of %i values per Pixel are used." % (working_class.values_used, len(working_class.layerlist)))
	working_class.superarray()
	#pdb.gimp_message("superarray shape: %s" % working_class.sarray.shape)
	working_class.calculate_mean()
	#pdb.gimp_message("new superarray shape: %s" % working_class.sarray.shape)
	working_class.export()
	return

#def tester():
	#"""make testing easier"""
	#img=gimp.image_list()[-1]
	#draw=img.layers[0]
	#do_the_work(img,draw,"arith",1,1,30,0)
	
register(
        "average_layer",
        "Compute average layer v1.1",
        "Compute average layer offering different statistical approaches.\n Arithmetical mean with cutoff equals winsorized mean.\n Try to use cutoff on both sides with different means \n to remove unwanted objects in a series of photographs.\n Mode is quite slow, but works (Take your time!).",
        "Elmar W. Hoefner",
        "Licensed under GPLv3",
        "May 2008",
        "<Image>/Filters/Layer Effects/A_verage Layer",
        "RGB*, GRAY*",
	[
	[PF_RADIO, "meantype", "Choose mean method", "arith", (
				("A_rithmetical", "arith"),
				("_Geometrical", "geom"),
				("_Harmonic","harm"),
				("M_edian","median"),
				("_Mode", "mode"))],
	[PF_TOGGLE, "as_new_image", "Paste as new image", 1],
	[PF_TOGGLE, "only_visible_layers", "Only use visible layers", 1],
	[PF_SPINNER, "cutoff", "Cut off extreme values (percent).\n At least one value is left, \npercentage is rounded (floor) to whole layers", 10, (0,100,5)],
	[PF_SPINNER, "side", "Cut off side (dark (0), bright (1), both(2))", 2, (0,2,1)]
	],
        [],
        do_the_work)

main()