#!BPY

"""
Name: 'Displace Mesh from UV Image'
Blender: 248
Group: 'UV'
Tooltip: 'Unselect meshs that dont match  '
"""

# $Id: clean_mesh.py,v 1.1 2004/08/04 06:16:45 ianwill Exp $
#
# -------------------------------------------------------------------------- 
# Mesh Cleaner 1.0 By Campbell Barton (AKA Ideasman)
# -------------------------------------------------------------------------- 
# ***** BEGIN GPL LICENSE BLOCK ***** 
# 
# 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. 
# 
# You should have received a copy of the GNU General Public License 
# along with this program; if not, write to the Free Software Foundation, 
# Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA. 
# 
# ***** END GPL LICENCE BLOCK ***** 
# -------------------------------------------------------------------------- 


# Made by Ideasman/Campbell 2005/06/10 - ideasman@linuxmail.org


from Blender import *

def NMeshCopy(mesh):
	copy = NMesh.New()
	copy.setMode( mesh.getMode() )
	copy.setMaterials(mesh.getMaterials(1)) # Empty == None
	copy.hasVertexUV(mesh.hasVertexUV())
	copy.hasFaceUV(mesh.hasFaceUV())
	copy.hasVertexColours(mesh.hasVertexColours())
	copy.verts = mesh.verts
	copy.faces = mesh.faces
	return copy
	



def main():
	try:
		scn = Scene.GetCurrent()
		active_ob = scn.objects.active
		me = active_ob.getData()
		
	except:
		Draw.PupMenu('Select 1 mesh object.')
		return
		
	if active_ob.getType() != 'Mesh':
		Draw.PupMenu('Select 1 mesh object.')
		return
	
	#USER OPTIONS
	
	# Axis - what axis to use
	axis = Draw.PupMenu('Displacement Axis%t|Vertex Normal|Objects Local X|Objects Local Y|Objects Local Z')
	if axis == -1:
		return
	
	
	# iterations
	iterations = Draw.PupIntInput('Iterations: ', 1, 1, 20)
	if not iterations:
		return
	
	# Invert- black high or whit high?
	invert = Draw.PupMenu('Invert tone?%t|White is Higher%x0|Black is Higher%x1')
	if invert == -1:
		return	
	
	# Normalize?
	normalize = Draw.PupMenu('Normalize Displace Range?%t|Defulat Off%x0|Stretch to Min/Max Values%x1')
	if normalize == -1:
		return	
	
	# How far do we displace allong our axis?.
	dist = Draw.PupFloatInput('displace: ', 0.5, 0.001, 20.0, 0.1, 3)
	if not dist:
		return
	
	dist = dist / iterations # Scale down the dist per iteration
	
	# zero is the level of colour that has no affect.
	# default is 0 wich means all will be raised
	# 1 measn all will be depresed.
	zero = Draw.PupFloatInput('zero level: ', 0.5, 0.0, 1.0, 0.1, 3)
	if zero == None: # zero may be 0 so use different if.
		return
	
	print '\nStarting displacement...'
	time1 = sys.time()
	is_editmode = Window.EditMode()
	if is_editmode:
		Window.EditMode(0)
	
	# Generate weights once
	weights = [[] for i in range(len(me.verts)) ]
	
	#------------------------------------------------------#
	# Weight from UV texture ------------------------------#
	#------------------------------------------------------#
	for f in [f for f in me.faces if len(f) > 2]:
		if f.image != None:
			# Get the pixel dimensions for the image
			maxx, maxy = f.image.getMaxXY()
			for i in range(len(f)):
				vert = f[i]
				u, v = f.uv[i]
				# Clamping the UV values to be inside the image.
				u %= 1.0; v %= 1.0
				
				# scale down by 1 pixel to avoid goin off teh edge of the image.				
				uPixLoc =int((u * maxx) - (1.0/maxx) ) # 1.0/maxy is 1 pixel in float
				vPixLoc =int((v * maxy) - (1.0/maxy) )
				
				col = f.image.getPixelF(uPixLoc, vPixLoc)
				
				
				avCol = (col[0] + col[1] + col[2]) / 3
				if invert:
					avCol = 1-avCol
				
				weights[vert.index].append(avCol)
	
	
	#------------------------------------------------------#
	# Weight from MTex ------------------------------------#
	#------------------------------------------------------#
	
	# First lets see of tw have any materials with textures, a fast way to see of we can displace based on textures.
	has_mtex_colour = False # assume no textures.
	for mat in me.getMaterials(0):
		if mat.getTextures():
			has_mtex_colour = True
			break
	
	# We at lease 1 MTex for the mesh.
	if has_mtex_colour:
		
		# Backup material shade value.
		materialShadeBackup = {}
		for mat in me.getMaterials(0): # we havnt modified the materials used so we can get from blenders mesh direct.
			if mat.name not in materialShadeBackup.keys(): # The mat might be applied twice, so we dont want to backup the modified Value.
				materialShadeBackup[mat.name] = mat.mode # & Material.Modes['SHADELESS']
				mat.mode |= Material.Modes['SHADELESS'] # We dont want lights to effect the texture.
		
		# Make a new mesh with vertex colours
		
		temp_tex_mesh = NMeshCopy(me)
		
		temp_tex_object = NMesh.PutRaw(temp_tex_mesh)
		
		temp_tex_object.setMatrix(active_ob.matrixWorld)
		temp_tex_mesh = temp_tex_object.getData()
		# me = active_ob.getData() # PutRaw removes the link so lets get the link back
		
		# Apply the texture to vert colours
		temp_tex_mesh.update(1,0,1)
		temp_tex_mesh = temp_tex_object.getData()
		
		# Slow and crap way to check if we have a useable texture.
		# check for colour variation! break once we find it.
		
		''' #  COMMENTED OUT FOR SPEED, Just assume it has colour
		has_mtex_colour = False
		colourVariarionDict = {}
		for f in [f for f in me.faces if len(f) > 2]:
			for col in f.col:
				colourVariarionDict[(col.r, col.g, col.b)] = None
				if len(colourVariarionDict) > 1:
					# He have mroe then 1 colour, break
					has_mtex_colour = True
		
		if has_mtex_colour:
		'''
		# Wight from vertex colours!
		for f in [f for f in temp_tex_mesh.faces if len(f) > 2]:
				for i in range(len(f)):
					vert = f[i]
					col = f.col[i]
					avCol = (col.r + col.g + col.b) / 765.0 # 255 each / 3 ... 255 * 3
					if invert:
						avCol = 1-avCol
					weights[vert.index].append(avCol)
		
		
		# Restore the material values
		for k in materialShadeBackup.keys():
			mat = Material.Get(k)
			mat.mode = materialShadeBackup[k]
		
		# Remove the vert col object
		# Comment out for testing	
		scn.objects.unlink(temp_tex_object)
		
	
	# Now combine all the weights
	max_weight = 0.0; min_weight = 1.0 # set to oppistes
	i = len(weights)
	while i:
		i-=1
		length = len(weights[i]) # Can we add below?
		if length:
			new_weight = reduce(lambda a,b: a+b, weights[i]) / length
			weights[i] = new_weight
			
			# Record max/min weights, olny used if normalize enabled
			max_weight = max(max_weight, new_weight)
			min_weight = min(min_weight, new_weight)
			
		else:
			weights[i] = None # 0?
		

	weight_range = max_weight - min_weight
	
	if weight_range == 0:
		Draw.PupMenu('There is no tonal variation from the displacement source')
		return
	
	
	
	# Apply Weights (iterations) times.
	while iterations:
		print 'iteration counting down to 1:', iterations
		iterations -=1 # Just a counter	
		
		orig_normals = [v.no for v in me.verts]
		
		# Now apply weights to verts?
		for i, v in enumerate(me.verts):
			w = weights[i]
		
			if w != None: # was some weight.
				
				# Normalize weights?
				if normalize:
					w = (w - min_weight) / weight_range
				
				if axis == 1: # Displace allong normals
					v.co.x, v.co.y, v.co.z = v.co + (dist * (orig_normals[i]*(w-zero)))  # apply the original normal location as well weighted.
				elif axis == 2: # Displace allong X axis
					v.co.x = v.co.x + (dist * (w-zero))
				elif axis == 3: # Displace allong Y axis
					v.co.y = v.co.y + (dist * (w-zero))				
				elif axis == 4: # Displace allong Z axis
					v.co.z = v.co.z + (dist * (w-zero))
		
		me.update(1, (me.edges != []), 0)
		if iterations: # Do this to get the new normals if we have more then 1 iteration to go
			me = active_ob.getData()
			
	# Reset the selection
	active_ob.sel = 0 # makes active again.		
	active_ob.sel = 1
	
	# Put us back into the right mode
	if is_editmode:
		Window.EditMode(1)
	print 'displaced in,',  sys.time() - time1

main()