#!BPY

""" Registration info for Blender menus:
Name:		 'Ma_Baker'
Blender:	248
Group:		'UV'
Tooltip:	 'Create uv colour according to the angle of each edge.'
"""

__author__ = "The Blender Community"
__url__ = ("http://www.blenderartists.org")
__version__ = "0.3 29/06/06"

__bpydoc__ = """\

Usage:

Open the script in your script window.
Or open it in your text window and hit ALT + P;

Select a mesh object and run the script.

Set the max angle lower for more severe colouring.
Set the min angle higher for the colouring to affect less verts.

"""
# Code history
"""
Version 0.1 28-june-2006
	Based on an idea by endi.
	Original code by macouno.

Version 0.2 29-june-2006
	Code optimised by tedi.
	Some bits removed, comments added,
	convex and concave only modes added,
	by macouno.
	
Version 0.3 03-juli-2006
	Spaces converted to tabs,
	added directional function,
	by macouno.

"""
# ***** BEGIN GPL LICENSE BLOCK *****
#
# Copyright (C) 2006: The Blenderartists Community
#
# 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 *****

import Blender
from Blender import*
from Blender.BGL import*
from Blender.Draw import*

## Global variables.
VERSION = '0.3'
BUTTON = {
	'MINANGLE': Create(0),
	'MAXANGLE': Create(180),
	'INVERT': Create(0),
	'CK': Create(0.0),
	'AVCOL': Create(0),
	'CONVEX': Create(1),
	'CONCAVE': Create(1),
	'WIND': Create(1),
	'INVWIND': Create(0),
	'WINDDIR': Create(6),
	'WINDCOL': Create(255)
}

DATA = {}
DATA['MENAME'] = "a"
DATA['MESSAGE'] = "Select a mesh object and run the script"

# Apply a matrix to a normal and return a vector.
def normal_transform(normal, matrix):
	matrix = matrix.rotationPart()
	x, y, z = normal
	return Mathutils.Vector(
	x*matrix[0][0] + y*matrix[1][0] + z*matrix[2][0],
	x*matrix[0][1] + y*matrix[1][1] + z*matrix[2][1],
	x*matrix[0][2] + y*matrix[1][2] + z*matrix[2][2])

# Function for getting the midpont of a face.
def GetFaceMid(face):
	Loc = [0] * 3
	for v in face.v:
		for a in range(3): Loc[a] += v.co[a]
	return Loc

# Function for getting the vector between two points
def PtoPVector(v1, v2):
	return Mathutils.Vector(v1[0] - v2[0], v1[1] - v2[1], v1[2] - v2[2])

# The script itself
def RunScript():
	# Make sure a mesh is selected first
	try:
		if Object.GetSelected()[0].getType() == 'Mesh': RunScript = True
		else:
			RunScript = False
			DATA['MESSAGE'] = 'You need to select a mesh object!'
	except:
		RunScript = False
		DATA['MESSAGE'] = 'You need to select a mesh object!'

	# Continue if a mesh has been selected
	if RunScript:

		# Set Base colours
		if BUTTON['CONCAVE'].val is 0 and BUTTON['CONVEX'].val is 0:
			if BUTTON['WIND'].val is 1:
				if BUTTON['INVWIND'].val is 1:
					BUTTON['CK'].val = 255
					BUTTON['AVCOL'].val = 255			
				else:
					BUTTON['CK'].val = 0
					BUTTON['AVCOL'].val = 0
			else:
				BUTTON['CK'].val = 127.5
				BUTTON['AVCOL'].val = int(127.5)
				RunScript = False
				DATA['MESSAGE'] = 'You should select a shading mode!'
		elif BUTTON['CONVEX'].val is 0 and BUTTON['CONCAVE'].val is 1:
			BUTTON['CK'].val = 255
			if BUTTON['INVERT'].val is 1:
				 BUTTON['AVCOL'].val = 0
			else:
				BUTTON['AVCOL'].val = 255
		elif BUTTON['CONCAVE'].val is 0 and BUTTON['CONVEX'].val is 1:
			BUTTON['CK'].val =255
			if BUTTON['INVERT'].val is 1:
				 BUTTON['AVCOL'].val = 255
			else:
				BUTTON['AVCOL'].val = 0	
		else:
			BUTTON['CK'].val = 127.5
			BUTTON['AVCOL'].val = int(127.5)
			
		if BUTTON['WIND'].val is 1:		
			if BUTTON['WINDDIR'].val is 1:
				WindVec = [0.0, 0.0, -1.0];
			elif BUTTON['WINDDIR'].val is 2:
				WindVec = [0.0, 1.0, 0.0];
			elif BUTTON['WINDDIR'].val is 3:
				WindVec = [1.0, 0.0 ,0.0];
			elif BUTTON['WINDDIR'].val is 4:
				WindVec = [-1.0, 0.0, 0.0];
			elif BUTTON['WINDDIR'].val is 5:
				WindVec = [0.0, -1.0, 0.0];
			else:
				WindVec = [0.0, 0.0, 1.0];

		# Get mesh data
		tarObj = Object.GetSelected()[0]
		DATA['MENAME'] = tarObj.getData(1)
		me = NMesh.GetRaw(DATA['MENAME'])
		if BUTTON['WIND'].val is 1:
			ObMatrix = Mathutils.Matrix(tarObj.getMatrix('worldspace'));

		# Make sure the mesh has vertex colours enabled
		if not me.hasVertexColours():
			me.hasVertexColours(1); me.update()
			sc = [BUTTON['AVCOL'].val] * 3 + [255]
			[f.col.append(NMesh.Col(sc[0], sc[1], sc[2], sc[3])) for f in me.faces for v in f.v]
			me.update()

		# Get the face normals and midpoints of each face.
		for v in xrange(len(me.verts)):
			Normals = []; MidPoints = []
			Myfaces = [f for f in me.faces if me.verts[v] in f.v]

			# Add midpoints and normals to their respective lists.
			for f in Myfaces:
				MidPoints.append(GetFaceMid(f))
				Normals.append(f.no)			

			# Calculate the angle only if more than one face is connected to the vert.
			if len(Normals) >= 1:
				TheAngle = 0; TheWind = 0; Dot = 0; Nr1 = 0; Nr2 = 0
				# Get the average angle between the faces and whether it's concave or convex.
				for a in xrange(len(Normals)):
				
					# Find out the angle of this face relative to the wind direction.
					if BUTTON['WIND'].val is 1:
						WindAngle = 180 - Mathutils.AngleBetweenVecs(Mathutils.Vector(WindVec), normal_transform(Normals[a], ObMatrix))
					
						if WindAngle > BUTTON['MINANGLE'].val:
						
							if WindAngle > BUTTON['MAXANGLE'].val: WindAngle = BUTTON['MAXANGLE'].val
							TheWind += WindAngle - BUTTON['MINANGLE'].val
							Nr2 += 1	
				
					if len(Normals) > 1:
						for b in xrange(len(Normals)):
							if a != b:
								Angle = Mathutils.AngleBetweenVecs(Mathutils.Vector(Normals[a]), Mathutils.Vector(Normals[b]))

								# Only get the details if the angle is bigger than the minimum
								if Angle > BUTTON['MINANGLE'].val:

									# Find out whether the angle is concave or convex
									MidVector = PtoPVector(MidPoints[a], MidPoints[b])
									DotVec = Mathutils.DotVecs(MidVector, Mathutils.Vector(Normals[a]))

									if Angle > BUTTON['MAXANGLE'].val: Angle = BUTTON['MAXANGLE'].val

									TheAngle += Angle - BUTTON['MINANGLE'].val
									Dot += DotVec
									Nr1 += 1
									
				if TheAngle > 0:
					TheAngle /= Nr1
					
				# Get the wind factor if it applies.
				if BUTTON['WIND'].val is 1:
					if TheWind > 0:
						TheWind /= Nr2
						WindFactor = (BUTTON['MAXANGLE'].val - BUTTON['MINANGLE'].val) / TheWind
					else: 
						WindFactor = 0

				# Set the vertex colour for the vert
				for f in Myfaces:
					for vNr in xrange(len(f.v)):
						if f.v[vNr] is me.verts[v]:
						
							if TheAngle is 0:
								ColVal = BUTTON['AVCOL'].val

							else:
								AngleFactor = (BUTTON['MAXANGLE'].val - BUTTON['MINANGLE'].val) / TheAngle

								# Calculate the colour for concave or convex only bake
								if BUTTON['CONVEX'].val is 0 and BUTTON['CONCAVE'].val is 1 and Dot < 0:
									if BUTTON['INVERT'].val is 1:
										ColVal = int(round(BUTTON['CK'].val / AngleFactor))
									else:
									 ColVal = int(round(BUTTON['CK'].val - (BUTTON['CK'].val / AngleFactor)))
								elif BUTTON['CONCAVE'].val is 0 and BUTTON['CONVEX'].val is 1 and Dot > 0:
									if BUTTON['INVERT'].val is 1:
										ColVal = int(round(BUTTON['CK'].val - (BUTTON['CK'].val / AngleFactor)))
									else:
										ColVal = int(round(BUTTON['CK'].val / AngleFactor))
								elif BUTTON['CONVEX'].val is 0 or BUTTON['CONCAVE'].val is 0:
									ColVal = BUTTON['AVCOL'].val

								# Calculate the colour for regular bake
								else:
									if Dot >= 0 and BUTTON['INVERT'].val is 0 or Dot <= 0 and BUTTON['INVERT'].val is 1:
										ColVal = int(round(BUTTON['CK'].val + (BUTTON['CK'].val / AngleFactor)))

									else: ColVal = int(round(BUTTON['CK'].val - (BUTTON['CK'].val / AngleFactor)))
									
							
							# Add the wind if it's a factor.
							if BUTTON['WIND'].val is 1:
								if WindFactor > 0.0:

									if BUTTON['CONVEX'].val is 1 or BUTTON['CONCAVE'].val is 1:
										if BUTTON['INVWIND'].val is 1:
											WindCol = 255 - (255 / WindFactor)
											ColVal = int(round((ColVal + WindCol) * 0.5))
										else:
											WindCol = 255 / WindFactor
											ColVal = int(round((ColVal + WindCol) * 0.5))
									else:
										if BUTTON['INVWIND'].val is 1:
											ColVal = 255 - int(round(255 / WindFactor))
										else:
											ColVal = int(round(255 / WindFactor))

							# Set the colour of the vert
							f.col[vNr].r = f.col[vNr].g = f.col[vNr].b = ColVal
							f.col[vNr].a = 255

			# Set the colour of the vert to the average if it has one or less faces connected
			else:
				for f in Myfaces:
					for vNr in xrange(len(f.v)):
						if f.v[vNr] is me.verts[v]:
							f.col[vNr].r = f.col[vNr].g = f.col[vNr].b = BUTTON['AVCOL'].val
							f.col[vNr].a = 255

		me.update()
		if RunScript is True:
			DATA['MESSAGE'] = str(DATA['MENAME']) + ' has been baked!'

# Draw the Gui
def DrawGui():

	# The background
	glClearColor(0.5, 0.5, 0.5, 0.0)
	glClear(GL_COLOR_BUFFER_BIT)

	# The black border
	glColor3f(0, 0, 0)
	glRectf(1, 1, 269, 274)

	# The highlight border
	glColor3f(0.7, 0.7, 0.8)
	glRectf(3, 3, 268, 273)

	# The dark bottom block
	glColor3f(0, 0, 0)
	glRectf(4, 4, 267, 165)

	# The settings header
	glRectf(4, 166, 267, 186)

	# The	messages
	glRectf(4, 187, 267, 230)

	# The messages header
	glRectf(4, 231, 267, 251)

	# The header
	glRectf(4, 252, 267, 272)

	# The header
	glColor3f(0.7, 0.7, 0.8)
	glRasterPos2d(20, 258)
	HeadText = "MaBaker " + str(VERSION) + "(c) 2006 endi, tedi & macouno"
	Text(HeadText)

	# Messages
	glColor3f(0.8, 0.9, 1.0)
	glRasterPos2d(20, 237)
	Text("messages")
	glColor3f(1.0, 1.0, 1.0)
	glRasterPos2d(20, 210)
	Text(str(DATA['MESSAGE']))

	# Settings
	glColor3f(0.8, 0.9, 1.0)
	glRasterPos2d(20, 172)
	Text("settings")

	# Button for inverting the shading
	BUTTON['INVERT'] = Toggle("invert shading", 2, 20, 135, 230, 20, BUTTON['INVERT'].val, "Invert the shading.")

	# Button for colouring only concave angles
	BUTTON['INVWIND'] = Toggle("invert wind", 2, 20, 105, 110, 20, BUTTON['INVWIND'].val, "Invert the wind Colour")
	
	types = "Wind origin %t|bottom %x1|rear %x2|right %x3|left %x4|front %x5|top %x6";	
	BUTTON['WINDDIR'] = Menu(types, 2, 140, 105, 110, 20, BUTTON['WINDDIR'].val, "The direction from which the directional abrasion originates.");

	# Button for colouring only convex angles
	BUTTON['CONVEX'] = Toggle("convex", 2, 20, 75, 70, 20, BUTTON['CONVEX'].val, "Shade convex angles.")

	# Button for colouring only concave angles
	BUTTON['CONCAVE'] = Toggle("concave", 2, 100, 75, 70, 20, BUTTON['CONCAVE'].val, "Shade concave angles.")

	# Button for colouring only concave angles
	BUTTON['WIND'] = Toggle("wind", 2, 180, 75, 70, 20, BUTTON['WIND'].val, "Shade directional wear")

	# Number input for maximum angle
	BUTTON['MAXANGLE'] = Number("max angle: ", 2, 140, 45, 110, 20, BUTTON['MAXANGLE'].val, BUTTON['MINANGLE'].val, 180.0, "The maximum angle")

	# Number input for minimum angle
	BUTTON['MINANGLE'] = Number("min angle: ", 2, 20, 45, 110, 20, BUTTON['MINANGLE'].val, 0.0, BUTTON['MAXANGLE'].val, "The minimum angle ")

	# Bake and exit script buttons
	Button("BAKE", 1, 20, 15, 110, 20, "Bake your currently selected mesh object.")
	Button("EXIT", 5, 140, 15, 110, 20, "Exit the script!")

# Close the script when the escape key is pressed.
def event(evt, val):
	if evt == QKEY and not val: Exit()

# The events as triggered by the gui.
def bevent(evt):
	if evt is 1:
		RunScript(); Redraw()
	elif evt is 2: Redraw()
	elif evt is 3 :
		BUTTON['CONCAVE'].val = 0
		Redraw()
	elif evt is 4:
		BUTTON['CONVEX'].val = 0
		Redraw()
	elif evt is 5: Exit()

# Register the events.
Register(DrawGui, event, bevent)