#!BPY

"""
Name: 'Onion_Skin'
Blender: 248
Group: 'Animation'
Tooltip: 'Adds Onion Skinning Spacehandler in the 3d view.'
"""
__author__ = ["crouch"]
__url__ = ("http://wiki.blender.org/index.php/Scripts/Manual/Animation/Onion_Skin")
__version__ = "0.2"

__bpydoc__ = """\
This script brings onion skinning to the 3d view. 

Enable Scriplinks. 

Open onionskin_spacehandler in the Blender text editor. 

Run Onion_Skin_gui script. 

"""

# ***** BEGIN GPL LICENSE BLOCK *****
#
# Copyright (C) 2009 Bart Crouch
#
# 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 *
import bpy

# Default buttons
alpha = Draw.Create(0.4)
but_ob_act = Draw.Create(0)
but_ob_sel = Draw.Create(0)
but_ob_all = Draw.Create(1)

# Get window coordinates
def getCoords():
	global xmin, ymin, width, height
	
	windows = Window.GetScreenInfo()
	for w in windows:
		if w['type'] == Window.Types.VIEW3D:
			xmin, ymin, xmax, ymax = w['vertices']
			width = xmax-xmin
			height = ymax-ymin

# Get z and colour buffers
def getBuffers():
	global zbuf, colbuf
	
	# Trick so the current onionskin is temporarily not displayed
	img.name = "OnionSkinningTemporary"
	# Hide objects according to the restriction toggle buttons
	if but_ob_act.val==1 or but_ob_sel.val==1:
		sce = bpy.data.scenes.active
		if but_ob_act.val==1:
			display = [sce.objects.active]
		if but_ob_sel.val==1:
			display = [ob for ob in sce.objects.selected]
		for ob in sce.objects:
			if ob not in display:
				ob.restrictDisplay = 1
	Blender.Redraw()
	# Get the buffers
	zbuf = Buffer(GL_FLOAT, [width*height])
	glReadPixels(xmin, ymin, width, height, GL_DEPTH_COMPONENT, GL_FLOAT, zbuf)
	colbuf = Buffer(GL_FLOAT, [width*height, 4])
	glReadPixels(xmin, ymin, width, height, GL_RGBA, GL_FLOAT, colbuf)
	# Restore old situation
	if but_ob_act.val==1 or but_ob_sel.val==1:
		for ob in sce.objects:
			ob.restrictDisplay = 0
	img.name = "OnionSkinning"

# Get image and saved data
def getImgData():
	global img, frames, curframe, view, offset
	
	try:
		img = Image.Get("OnionSkinning")
		frames = [f for f in img.properties['OnionSkinning']['frames']]
		view = [q for q in img.properties['OnionSkinning']['view']]
		offset = [o for o in img.properties['OnionSkinning']['offset']]
	
	except:
		img = Image.New("OnionSkinning", width, height, 32)
		frames = []
		view = Window.GetViewQuat()
		offset = Window.GetViewOffset()
	curframe = Blender.Get('curframe')

# Draw new image
def drawImage(frame, total_frames):
	progress = 0
	progress_text = "Drawing onion skin: %i/%i"%(frame, total_frames)
	Window.DrawProgressBar(progress, progress_text)
	frames.append(curframe)
	newfac = 1.0/len(frames)
	
	for y in range(height):
		#progress bar
		current_progress = round((float(y)/(height-1))*100)
		if current_progress-10>=progress*100:
			progress = current_progress/100.0
			Window.DrawProgressBar(progress, progress_text)
		
		for x in range(width):
			r_old, g_old, b_old, a_old = img.getPixelF(x,y)
			i = y*width + x
			z = zbuf[i]
			if z != 1:
				r, g, b, a = colbuf[i]
				if a_old == 0:
					a_old_fac = 0
				else:
					a_old_fac = 1
				if newfac != 1:
					r = (r_old + r)/(1.0+a_old_fac)
					g = (g_old + g)/(1.0+a_old_fac)
					b = (r_old + b)/(1.0+a_old_fac)
				img.setPixelF(x,y,(r,g,b,alpha.val))
			else:
				if newfac == 1:
					r_old, g_old, b_old, a_old = colbuf[i]
					r_old = g_old = b_old = 0.0
					a_old = 0.0
				img.setPixelF(x,y,(r_old,g_old,b_old,a_old))
	Window.DrawProgressBar(1,progress_text)

# Save data along with image
def saveData():
	# Save data as ID properties
	zoom = Window.GetPerspMatrix().scalePart().length
	img.properties['OnionSkinning'] = {}
	img.properties['OnionSkinning']['frames'] = frames
	img.properties['OnionSkinning']['size'] = [width, height]
	img.properties['OnionSkinning']['view'] = Window.GetViewQuat()
	img.properties['OnionSkinning']['offset'] = Window.GetViewOffset()
	img.properties['OnionSkinning']['zoom'] = zoom
		
	# Free old image from memory
	img.glFree()

# Change the size of the image (3d window size has changed)
def newSize():
	global img
	img = Image.Get("OnionSkinning")
	all_frames = [f for f in img.properties['OnionSkinning']['frames']]
	getCoords()
	img.name = "OnionSkinningOld"
	img = Image.New("OnionSkinning", width, height, 32)
	saveData()
	updateOnion(True)

# Test if an onion skin has already been created
def onionExist():
	try:
		img = Image.Get("OnionSkinning")
		return True
	except:
		Window.WaitCursor(0)
		choice = Draw.PupMenu("No onion skin present yet%t|Onionize!|Cancel")
		Window.WaitCursor(1)
		if choice == 1:
			addFrame(1,1)
		return False

# Test if the current view matches that of the onion skin
def testView():
	getImgData()
	if len(frames)==0:
		return True
	if Window.GetViewQuat() != view or Window.GetViewOffset() != offset:
		Window.WaitCursor(0)
		choice = Draw.PupMenu("Current view doesn't match onion skin%t|Match view|Cancel")
		Window.WaitCursor(1)
		if choice == 1:
			resetView()
		else:
			return False
	return True

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

# Add the current frame to the image
def addFrame(frame, total_frames):
	getCoords()
	
	# Error checking
	if not testView():
		return
	if curframe in frames:
		Window.WaitCursor(0)
		choice = Draw.PupMenu("Frame already skinned%t|Update onion skin|Cancel")
		Window.WaitCursor(1)
		if choice == 1:
			updateOnion(True)
		return
	try:
		old_zoom = Image.Get("OnionSkinning").properties['OnionSkinning']['zoom']
		old_size = Image.Get("OnionSkinning").size
	except:
		old_zoom = Window.GetPerspMatrix().scalePart().length
		old_size = [width, height]
	if Window.GetPerspMatrix().scalePart().length != old_zoom:
		if curframe not in frames:
			frames.append(curframe)
			img.properties['OnionSkinning']['frames'] = frames
		updateOnion(True)
		return
	if old_size[0]!=width or old_size[1]!=height:
		if curframe not in frames:
			frames.append(curframe)
			img.properties['OnionSkinning']['frames'] = frames
		newSize()
		return
	
	getBuffers()
	drawImage(frame, total_frames)
	saveData()
	Blender.Redraw()

# Change the transparency of the onion skin
def changeAlpha():
	global alpha
	Window.WaitCursor(0)
	alpha_new = Draw.PupFloatInput("Alpha", alpha.val, 0.0, 1.0, 1, 2)
	Window.WaitCursor(1)
	if alpha_new:
		alpha = Draw.Create(alpha_new)
	else:
		return
	try:
		img = Image.Get("OnionSkinning")
	except:
		return
	width, height = img.size
	progress = 0
	progress_text = "Changing transparency"
	Window.DrawProgressBar(progress, progress_text)
	for y in range(height):
		current_progress = round((float(y)/(height-1))*100)
		if current_progress-20>=progress*100:
			progress = current_progress/100.0
			Window.DrawProgressBar(progress, progress_text)
		for x in range(width):
			r, g, b, a = img.getPixelF(x,y)
			if a != 0.0:
				img.setPixelF(x, y, (r, g, b, alpha.val))
	Window.DrawProgressBar(1, progress_text)
	Blender.Redraw()

# Reset the view to the one used by the onion skin
def resetView():
	if not onionExist():
		return
	
	getImgData()
	Window.SetViewQuat(view)
	Window.SetViewOffset(offset)
	
	Blender.Redraw()

# Update the onion skin
def updateOnion(true_update):
	global width, height, frames, view, offset

	# Error checking
	if not onionExist():
		return
	if not true_update:
		if not testView():
			return
	else:
		getImgData()
	
	# Clear image
	progress = 0
	progress_text = "Clearing old onion skin"
	Window.DrawProgressBar(progress, progress_text)
	width, height = img.size
	all_frames = frames
	for y in range(height):
		# Progress bar
		current_progress = round((float(y)/(height-1))*100)
		if current_progress-20>=progress*100:
			progress = current_progress/100.0
			Window.DrawProgressBar(progress, progress_text)
		
		for x in range(width):
			img.setPixelF(x,y,(0.0,0.0,0.0,0.0))
	frames = []
	saveData()
	Window.DrawProgressBar(1, progress_text)
	
	# Add the old frames
	original_frame = Blender.Get('curframe')
	for f in all_frames:
		Blender.Set('curframe', f)
		Blender.Redraw()
		addFrame(all_frames.index(f)+1,len(all_frames))
	Blender.Set('curframe', original_frame)
	Blender.Redraw(-1)

# Remove the current frame from the onion skin
def removeOnion(removeAll):
	global frames
	
	# Error checking
	if not onionExist():
		return
	if not testView():
		return
	if not removeAll and curframe not in frames:
		Window.WaitCursor(0)
		error = Draw.PupMenu("Error%t|Current frame not in onion skin")
		Window.WaitCursor(1)
		return

	if removeAll:
		frames = []
	else:
		frames.remove(curframe)
	img.properties['OnionSkinning']['frames'] = frames
	updateOnion(False)
	
#########################

# Draw the Graphical User Interface
def drawGui():
	global but_ob_act, but_ob_sel, but_ob_all
	but_add = Draw.PushButton("Onionize!", 1, 5, 190, 80, 20, "Add the current frame to the onion skin")
	but_ob_act = Draw.Toggle("act", 2, 5, 155, 27, 20, but_ob_act.val, "Add only the active object to the onion skin")
	but_ob_sel = Draw.Toggle("sel", 3, 32, 155, 26, 20, but_ob_sel.val, "Add all selected objects to the onion skin")
	but_ob_all = Draw.Toggle("all", 4, 58, 155, 27, 20, but_ob_all.val, "Add everything to the onion skin")
	but_alpha = Draw.PushButton("Alpha", 5, 5, 135, 80, 20, "Change the transparency of the onion skin")
	but_view = Draw.PushButton("Go to view", 6, 5, 100, 80, 20, "Go to the 3D view that contains the onion skin")
	but_refresh = Draw.PushButton("Update", 7, 5, 80, 80, 20, "Update the onion skin to the current situation")
	but_del = Draw.PushButton("Delete", 8, 5, 60, 80, 20, "Remove the current frame from the onion skin")
	but_delall = Draw.PushButton("Delete all", 9, 5, 40, 80, 20, "Remove all frames from the onion skin")
	but_quit = Draw.PushButton("Quit", 99, 5, 5, 80, 20, "Exit this script")

# Mouse and keyboard processing
def event(evt, val):
	if evt in [Draw.QKEY, Draw.XKEY, Draw.ESCKEY]:
		Draw.Exit()

# Make the restriction toggle buttons work
def changeRestriction(evt):
	global but_ob_act, but_ob_sel, but_ob_all
	but_ob_act = but_ob_sel = but_ob_all = Draw.Create(0)
	if evt==2:
		but_ob_act = Draw.Create(1)
	if evt==3:
		but_ob_sel = Draw.Create(1)
	if evt==4:
		but_ob_all = Draw.Create(1)
	Draw.Redraw()

# Button processing
def bevent(evt):
	Window.WaitCursor(1)
	if evt == 1:
		addFrame(1,1)
	elif evt in [2,3,4]:
		changeRestriction(evt)
	elif evt == 5:
		changeAlpha()
	elif evt == 6:
		resetView()
	elif evt == 7:
		updateOnion(True)
	elif evt == 8:
		removeOnion(False)
	elif evt == 9:
		removeOnion(True)
	Window.WaitCursor(0)
	if evt == 99:
		Draw.Exit()

Draw.Register(drawGui, event, bevent)
