#####################################################
#  ICARUS Camera Calibration Importer for Blender   #
#####################################################
#  Version 1.07d: Sep 17, 2004                      #
#                                                   #
# --------------------------------------------------#
#  ALT-P TO START                                   #
#                                                   #
#####################################################
#
# Compatible ONLY with Blender 2.34, possibly newer too, depending on python API
#
# import1.03: ICARUS >1.07 adaptation, feature points data ignored
#
# import1.04: ICARUS v2.00 adaptation, new file format, not backwards compatible.
#		rot_quaternion might be useful, but ignored for now, matrix still works as it should, nothing to fix.
#		focal_length has changed. not used anymore, uses field of view instead which is now also available.
#		The ICARUS example files don't seem to have the same origin point for calibration & reconstruction,
#		so I could not really test if the fov is correct, it does not seem to be quite right.
#
# import1.05: Simon Gibson of the ICARUS team helped out with the FOV calculation, it is now correct, thanks Simon!
#	included the option of importing the feature points data as a mesh
#
# import1.06: Apparently the fileformat has changed, one of the (very) few users of the script (Gregory Harbin)
#			  found that it couldn't import the tracking data any longer.
#			  There is an extra camera projection matrix added to the fileformat.
#
# import1.07: Special 2.32 version on request
# import1.07b: Another request, updated for >=233, though still no ipo lens curve support it seems.
# import1.07c: Finally, Blender2.34, Camera Lens curve can be created from script.
# import1.07d: Modified to make it work without math module, no longer needs external python modules

import Blender

# Get Blender version
BL_VERSION = Blender.Get('version')
if BL_VERSION<234:
	raise "Sorry, this version needs Blender version 234..."

from Blender import NMesh, BGL, Draw, Window
from Blender.Mathutils import *

#################################
# GLOBAL VARS                   #
#################################

# pi factors needed for tangent approximation
pi_1 = 3.14159265359
pi_2 = 6.28318530718
pi_q = 0.785398163397
pi_h = 1.57079632679
pi_3q = 2.35619449019
deg2rad = 0.01745329252 # for fov calc.

# number of data text lines for a single frame (>=ICARUS v2.08)
DATA_BLOCK_SIZE = 18

ERRMSG		= 	[0, 'Ready!']
AVG_ASPR	= 	0.0
TOT_FRAMES	=	0
AVG_FOCLEN	=	0.0
SCR_SIZE	=	(0, 0)
INFO_TEXT	=	None
LAST_FILENAME	=	''

#---------------
# EVENT NUMBERS
#---------------
evt_EXIT	=	1
evt_fselect	=	2
evt_ignore	=	3
evt_IPO		=	4
evt_singimg	=	5
evt_imp_fpt	=	6

#---------
# BUTTONS
#---------
Tfilename	=	Draw.Create('')
Timp_fpt	=	Draw.Create(0)


#################################
# FUNCTIONS						#
#################################

# tangent approximation
def fTan(x):
	t1 = (0.999999986 - 0.0958010197*x*x)*x
	t2 = 1.0 + (0.00971659383*x*x - 0.429135022)*x*x
	return t1/t2

def tan_apx(x):
	t = abs(x) / pi_1
	t = pi_1 * (t - float(int(t)))
	if t>pi_3q:
		v = -fTan(pi-t)
	elif t>pi_h:
		v = -1.0/fTan(t-pi_h)
	elif t>pi_q:
		v = 1.0/fTan(pi_h-t)
	else:
		v = fTan(t)
	if x<0.0: return -v
	return v

# converts Field Of View to pixel focal length to Camera Lens value
def FOV2Lens(FOV, width, height):
	Lens = 16.0 / tan_apx(0.5 * deg2rad * FOV[0])
	if width<height: Lens *= (width/height)
	return Lens

# Find the best integer fraction representation to approximate a real value
# Used for determining a value for AspX & AspY from aspectratio
def fracApprox(frac, maxintval):
	bot = [int(frac), 1]
	top = [int(frac+1.0), 1]
	while 1:
		mid = [bot[0]+top[0], bot[1]+top[1]]
		if (mid[0] / float(mid[1])) < frac:
			bot = mid
		else:
			top = mid
		if max(max(bot, top)) > maxintval: break
	return tuple(min(bot, top))

# opens file and returns data if succesful
def LOADFILE():
	global ERRMSG, LAST_FILENAME

	FILENAME = Tfilename.val
	LAST_FILENAME = FILENAME

	if FILENAME=='':
		ERRMSG = [1, "No file specfied!"]
		return

	try:
		fp = open(FILENAME)
		data = fp.read()
		fp.close()
	except IOError:
		ERRMSG = [1, "Cannot open %s !!" % FILENAME]
		return
	else:
		# consider the first comment as a file ID, of course probably not future compatible
		if data.splitlines()[0][:14]!="# Camera track":
			ERRMSG = [1, "Unexpected start of file, nothing done"]
			return

	return data


# the main function, converts the data to IPO curve(s)
def MakeIpoCurves():
	global ERRMSG, AVG_ASPR, AVG_FOCLEN, TOT_FRAMES, SCR_SIZE, INFO_TEXT

	dt = LOADFILE()
	if dt==None:
		return

	# Get the current scene camera object, for <=223 use Blender210 module
	camobj = Blender.Scene.getCurrent().getCurrentCamera()
	camobj_ipo = camobj.ipo
	cam = camobj.getData()
	camipo = cam.getIpo()

	if not camobj_ipo:
		camobj_ipo = Blender.Ipo.New('Object', 'camobj_ipo')
		camobj.setIpo(camobj_ipo)
		# add LocXYZ & RotXYZ curves
		LocX = camobj_ipo.addCurve('LocX')
		LocY = camobj_ipo.addCurve('LocY')
		LocZ = camobj_ipo.addCurve('LocZ')
		RotX = camobj_ipo.addCurve('RotX')
		RotY = camobj_ipo.addCurve('RotY')
		RotZ = camobj_ipo.addCurve('RotZ')
	else:
		OBJ_ipoc = camobj_ipo.getCurves()
		# get available curvenames
		camob_curvename_list = [cr.getName() for cr in OBJ_ipoc]
		# get or create LocX
		try:
			LocX = OBJ_ipoc[camob_curvename_list.index('LocX')]
		except:
			LocX = camobj_ipo.addCurve('LocX')
		# get or create LocY
		try:
			LocY = OBJ_ipoc[camob_curvename_list.index('LocY')]
		except:
			LocY = camobj_ipo.addCurve('LocY')
		# get or create LocZ
		try:
			LocZ = OBJ_ipoc[camob_curvename_list.index('LocZ')]
		except:
			LocZ = camobj_ipo.addCurve('LocZ')
		# get or create RotX
		try:
			RotX = OBJ_ipoc[camob_curvename_list.index('RotX')]
		except:
			RotX = camobj_ipo.addCurve('RotX')
		# get or create RotY
		try:
			RotY = OBJ_ipoc[camob_curvename_list.index('RotY')]
		except:
			RotY = camobj_ipo.addCurve('RotY')
		# get or create RotY
		try:
			RotZ = OBJ_ipoc[camob_curvename_list.index('RotZ')]
		except:
			RotZ = camobj_ipo.addCurve('RotZ')

	# get curvenames again, since some curves might be new
	OBJ_ipoc = camobj_ipo.getCurves()
	camob_curvename_list = [cr.getName() for cr in OBJ_ipoc]

	INFO_TEXT = None

	# test camera ipo, from 234 finally possible to add a Lens curve
	if not camipo:
		camipo = Blender.Ipo.New('Camera', 'camipo')
		cam.setIpo(camipo)
		CamLens = camipo.addCurve("Lens")
	else:
		CamLens = camipo.getCurve("Lens")

	# Bezier can now be set for interpolation,
	# but cannot set the control points like in the old api yet
	OBJ_ipoc[camob_curvename_list.index('LocX')].setInterpolation('Bezier')
	OBJ_ipoc[camob_curvename_list.index('LocY')].setInterpolation('Bezier')
	OBJ_ipoc[camob_curvename_list.index('LocZ')].setInterpolation('Bezier')
	OBJ_ipoc[camob_curvename_list.index('RotX')].setInterpolation('Bezier')
	OBJ_ipoc[camob_curvename_list.index('RotY')].setInterpolation('Bezier')
	OBJ_ipoc[camob_curvename_list.index('RotZ')].setInterpolation('Bezier')
	if CamLens: CamLens.setInterpolation('Bezier')

	# Parse file
	# block order is expected to be (not counting empty lines):
	# Frame int
	#   projection_matrix
	#		matrix4x3
	# 	translation
	#		vector3
	#	rot_matrix
	#		matrix3x3
	# 	rot_quaternion
	#		quaternion
	# 	focal_length float
	# 	field_of_view float float
	# 	pixel_aspect_ratio float
	# 	skew float
	# 	principal_point vector3

	# parse data, start at line 2, remove empty lines
	dtlines = [line.strip() for line in dt.splitlines()[2:] if line!='']

	# for ICARUS >1.07: find start of feature points data and remove, not used here
	# can be imported seperately though
	fpdt_start = 0
	for line_num in range(len(dtlines)):
		if dtlines[line_num].find('features')!=-1:
			fpdt_start = line_num
			break
	if fpdt_start: dtlines = dtlines[:fpdt_start]

	# read in all frames, setting points in camera ipo curves
	# Only uses LocXYZ & RotXYZ, other data ignored

	# the average aspect ratio, focal length & screensize (from principal point) values
	tot_aspr = 0.0
	tot_foclen = 0.0
	tot_scrsize = [0.0, 0.0]
	# number of frames from number of datablocks
	TOT_FRAMES = len(dtlines) / DATA_BLOCK_SIZE
	# now either create or modify Ipo curve points
	# Y & Z axes are swapped
	for i in range(TOT_FRAMES):
		# catch parse faults when file is not what it was expected to be
		try:
			# read data block per frame
			DATA_PT = i * DATA_BLOCK_SIZE
			dtblock = dtlines[DATA_PT : DATA_PT + DATA_BLOCK_SIZE]
			# frame number
			curframe = float(dtblock[0].split()[1])
			# camera_projection block skipped
			# translation
			t = dtblock[6].split()
			curloc = [float(t[0]), float(t[1]), float(t[2])]
			# rotation matrix
			mtx = [0,0,0]
			t = dtblock[8].split()
			mtx[0] = [float(t[0]), float(t[1]), float(t[2])]
			t = dtblock[9].split()
			mtx[1] = [float(t[0]), float(t[1]), float(t[2])]
			t = dtblock[10].split()
			mtx[2] = [float(t[0]), float(t[1]), float(t[2])]
			# focal length, not used, not in pixels anymore
			focal_length = float(dtblock[13].split()[1])
			# so use field of view instead
			t = dtblock[14].split()
			field_of_view = (float(t[1]), float(t[2]))
			# pixel aspect ratio
			aspect_ratio = float(dtblock[15].split()[1])
			# sum for average
			tot_aspr += aspect_ratio
			# skew is ignored
			# principal point, which is assumed to be the center of the screen
			# calculate width & height from this
			t = dtblock[17].split()
			principal_point = [float(t[1]), float(t[2])]
			scrwidth = principal_point[0] * 2.0
			scrheight = principal_point[1] * 2.0
			# sumtotal to calculate average screen size from principal point
			tot_scrsize[0] += scrwidth
			tot_scrsize[1] += scrheight
		except:
			ERRMSG = [1, "File parse error! Is this really an ICARUS file?"]
			return

		# Matrix element swap (L<->R) as adviced by Hos
		mtx2 = [[mtx[0][0], mtx[2][0], mtx[1][0]],
				[mtx[0][2], mtx[2][2], mtx[1][2]],
				[mtx[0][1], mtx[2][1], mtx[1][1]]]
		currot = Matrix(mtx2[0], mtx2[1], mtx2[2]).toEuler()
		currot = Vector(list(currot)) * 0.1	# ipo corr.

		# Set the the data in ipo curve
		# the Blender current frame curframe+1.0
		Bcurframe = curframe+1.0

		# LocX
		try:
			nbzt = LocX.getPoints()[i]
			nbzt.setPoints((Bcurframe, curloc[0]))
 		except:
			LocX.addBezier((Bcurframe, curloc[0]))
		# LocY
		try:
			nbzt = LocY.getPoints()[i]
			nbzt.setPoints((Bcurframe, curloc[2]))
 		except:
			LocY.addBezier((Bcurframe, curloc[2]))
		# LocZ
		try:
			nbzt = LocZ.getPoints()[i]
			nbzt.setPoints((Bcurframe, curloc[1]))
 		except:
			LocZ.addBezier((Bcurframe, curloc[1]))
		# RotX
		try:
			nbzt = RotX.getPoints()[i]
			nbzt.setPoints((Bcurframe, 9+currot[0]))
 		except:
			RotX.addBezier((Bcurframe, 9+currot[0]))
		# RotY
		try:
			nbzt = RotY.getPoints()[i]
			nbzt.setPoints((Bcurframe, currot[1]))
 		except:
			RotY.addBezier((Bcurframe, currot[1]))
		# RotZ
		try:
			nbzt = RotZ.getPoints()[i]
			nbzt.setPoints((Bcurframe, currot[2]))
 		except:
			RotZ.addBezier((Bcurframe, currot[2]))

		# now set the camera lens value if possible
		if CamLens:
			try:
				nbzt = CamLens.getPoints()[i]
				nbzt.setPoints((Bcurframe, FOV2Lens(field_of_view, scrwidth, scrheight)))
			except:
				CamLens.addBezier((Bcurframe, FOV2Lens(field_of_view, scrwidth, scrheight)))

	# calculate average of aspect ratio & focal length
	t = float(TOT_FRAMES)
	if t!=0.0: t = 1.0/t
	AVG_ASPR = tot_aspr * t
	AVG_FOCLEN = tot_foclen * t
	# average screen size rounded
	SCR_SIZE = [round(tot_scrsize[0]*t), round(tot_scrsize[1]*t)]

	# return ipo(s) to Blender
	OBJ_ipoc[camob_curvename_list.index('LocX')].points = LocX
	OBJ_ipoc[camob_curvename_list.index('LocY')].points = LocY
	OBJ_ipoc[camob_curvename_list.index('LocZ')].points = LocZ
	OBJ_ipoc[camob_curvename_list.index('RotX')].points = RotX
	OBJ_ipoc[camob_curvename_list.index('RotY')].points = RotY
	OBJ_ipoc[camob_curvename_list.index('RotZ')].points = RotZ
	if CamLens: CamLens.points = CamLens

	# Set the current frame back to the first to make changes visible
	sc = Blender.Scene.getCurrent()
	rc = sc.getRenderingContext()
	rc.startFrame(1)
	rc.endFrame(Bcurframe)
	rc.currentFrame(1)
	# set render window size
	rc.imageSizeX(SCR_SIZE[0])
	rc.imageSizeY(SCR_SIZE[1])
	# aspect ratio
	ar = fracApprox(AVG_ASPR, 200)
	rc.aspectRatioX(ar[0])
	rc.aspectRatioY(ar[1])

	ERRMSG = [0, "Done! Ipo(s) modified/created."]


# Create mesh from Feature Points
def MakeFPTMesh():
	global ERRMSG

	dt = LOADFILE()
	if dt==None: return

	# parse data, start at line 2, remove empty lines
	dtlines = [line.strip() for line in dt.splitlines()[2:] if line!='']

	# find start of feature points data
	fpdt_start = 0
	for line_num in range(len(dtlines)):
		if dtlines[line_num].find('features')!=-1:
			fpdt_start = line_num
			break
	if fpdt_start:
		dtlines = dtlines[fpdt_start:]

		# number of features
		fpt_num = int(dtlines[0].split()[0])

		if fpt_num!=0:

			me = NMesh.GetRaw()
			# y/z swapped
			for line in dtlines[1:]:
				t = line.split()
				me.verts.append(NMesh.Vert(float(t[2]), float(t[4]), float(t[3])))

			NMesh.PutRaw(me, 'Features')

			ERRMSG = [0, 'Done! Feature point mesh created.']
			return

	ERRMSG = [1, 'No feature point data found']


#################################
# GUI FUNCTIONS					#
#################################

# FILE SELECT (back again for 231)
def callback(filename):
	global Tfilename, ERRMSG
	Tfilename.val = filename
	ERRMSG = [0, 'Ready !']
	#Draw.Register(draw, event, bevent)

def FileSelect():
	Window.FileSelector(callback, "Load Icarus file")


# Window size
def GetWindowFactors():
	# Get the viewport size to make the buttons fit the window
	wsize = BGL.Buffer(BGL.GL_INT, 4)
	BGL.glGetIntegerv(BGL.GL_VIEWPORT, wsize)
	return wsize[2]*0.5, wsize[2]/425.0, wsize[3]/514.0


def draw():
	global Tfilename, Timp_fpt

	# clear screen
	BGL.glClearColor(0.35, 0.35, 0.35, 1)
	BGL.glClear(BGL.GL_COLOR_BUFFER_BIT)

	# get the window coordinate multiply factors
	midx, mulx, muly = GetWindowFactors()

	# background title box
	BGL.glColor3f(0.1, 0.1, 0.1)
	BGL.glRecti(0, 484*muly, 426*mulx, 514*muly)

	# Draw title in 3d style
	Title = "ICARUS Camera Calibration Import"
	BGL.glColor3f(0.5, 0.5, 0.5)
	BGL.glRasterPos2i(midx - len(Title)*3, 494*muly)
	Draw.Text(Title)
	BGL.glColor3f(1, 1, 1)
	BGL.glRasterPos2i(midx - len(Title)*3-2, 494*muly+2)
	Draw.Text(Title)

	# FileSelector button (back again)
	Draw.Button("FSEL", evt_fselect, 10*mulx, 450*muly, 50*mulx, 18*muly, "Use fileselector")

	Tfilename = Draw.String("Filename: ", evt_ignore, 62*mulx, 450*muly, 415*mulx-62*mulx, 18*muly, Tfilename.val, 128,
							"Name of Camera Calibration textfile")

	# The main work buttons
	Draw.Button("Create Curves", evt_IPO, 10*mulx, 420*muly, 150*mulx, 18*muly, "Create IPO curves from filedata")
	# Single image button not used, probably just as well handled, if not better, with a single value Ipo
	#Draw.Button("Create Loc&Rot", evt_singimg, 170*mulx, 420*muly, 150*mulx, 18*muly, "Create Loc & Rot values for single frame image")
	# Import Feature points as mesh option
	Draw.Button("Feature Points Mesh", evt_imp_fpt, 170*mulx, 420*muly, 150*mulx, 18*muly, "Creates mesh from feature point data")

	# information text, info about file
	if (AVG_ASPR!=0.0) and not ERRMSG[0]:
		BGL.glColor3f(0.5, 1, 0.5)
		BGL.glRasterPos2i(10*mulx, 380*muly)
		Draw.Text("Info for file:")
		BGL.glColor3f(1, 1, 0)
		BGL.glRasterPos2i(10*mulx, 360*muly)
		# not from Tfilename.val, as the name will change when the button is used without creating curves yet
		Draw.Text(LAST_FILENAME)
		BGL.glColor3f(0.5, 1, 0.5)
		BGL.glRasterPos2i(10*mulx, 340*muly)
		Draw.Text("Total number of frames: %d" % TOT_FRAMES)
		BGL.glColor3f(1, 0.5, 0)
		BGL.glRasterPos2i(10*mulx, 310*muly)

		# Since in 234, now all parameters can be set from script,
		# the INFO_TEXT is not used anymore, always None
		if INFO_TEXT:
			BGL.glColor3f(1, 0.75, 0)
			i = 0
			for st in INFO_TEXT:
				BGL.glRasterPos2i(10*mulx, (220-20*i)*muly)
				Draw.Text(st)
				i += 1

	BGL.glColor3f(0, 0, 0)
	BGL.glRecti(10*mulx, 33*muly, 415*mulx, 53*muly)
	if ERRMSG[0]:
		# Real error
		BGL.glColor3f(1, 0, 0)
	else:
		BGL.glColor3f(0, 1, 0)
	BGL.glRasterPos2i(14*mulx, 40*muly)
	Draw.Text(ERRMSG[1])

	# Exit button
	Draw.Button("Exit", evt_EXIT, 10*mulx, 10*muly, 100*mulx, 18*muly)


def event(evt, val):
	if (evt==Draw.QKEY and not val):
		Draw.Exit()


def bevent(evt):
	if evt == evt_EXIT:
		Draw.Exit()
	elif evt == evt_fselect:
		FileSelect()
	elif evt == evt_IPO:
		MakeIpoCurves()
		Blender.Window.RedrawAll()
	elif evt == evt_imp_fpt:
		MakeFPTMesh()
		Blender.Window.RedrawAll()
	#elif evt == evt_singimg:
	#	MakeSingleImageCam(Tfilename.val)


Draw.Register(draw, event, bevent)