assembly_handbook/ahb_techdraw_extensions.py

import FreeCAD as App
import FreeCADGui as Gui

from PySide.QtCore import QTimer

import TechDraw, TechDrawGui
from PySide import QtGui, QtCore
TDG = TechDrawGui

class CursorItem(QtGui.QGraphicsItem):
	def __init__(self, parent = None, view = None):
		super().__init__(parent)
		self.Type = QtGui.QGraphicsItem.UserType + 501
		self.setZValue(500)
		self.margin = 10.0
		self.size = 100.0
		self.view = view
		
	def removeSceneEventFilter(self, a, b):
		print('removeSceneEventFilter', a, b)
		
	def onViewPosChange(self, callback):
		self.viewPosChangeCallback = callback

	def boundingRect(self):
		return QtCore.QRectF(-self.size/2 - self.margin, -self.size/2 - self.margin,
					  self.size + 2.0*self.margin, self.size + 2.0*self.margin)

	def paint(self, painter, option, widget):
		#print("paint")
		painter.setBrush(QtCore.Qt.darkRed)
		#painter.drawRoundedRect(-100, -100, 200, 200, 50, 50)
		h = self.size/2.0
		painter.drawLine(-h, 0, h, 0)
		painter.drawLine(0, -h, 0, h)
		
	def mousePressEvent(self, event):
		#print('mouse press', event)
		self.startMovePos = event.pos()
		
	def mouseMoveEvent(self, event):
		#print('mouse move', event)
		offset = event.pos() - self.startMovePos
		self.moveBy(offset.x(), offset.y())
		
	def mouseReleaseEvent(self, event):
		#print('mouse release', event)
		#print('new pos', self.x(), self.y())
		if self.viewPosChangeCallback is not None:
			self.viewPosChangeCallback(self.getViewPos())
	
	def getViewPos(self):
		scale = self.view.Scale * 10.0
		return App.Vector(self.x() / scale, -self.y() / scale)
	
	def setViewPos(self, p):
		scale = self.view.Scale * 10.0
		self.setPos(p.x * scale, -p.y * scale)

class ViewCache:
	def __init__(self):
		self.reset()
		
	def reset(self):
		self.projected_origin = None

class TechDrawExtensions:
	views_to_repaint = {}
	
	view_cursors = {}
	
	view_cache = {}
	
	updating_balloon = False
	
	edited_view = None
	
	enable_selected_part_highlight = False # disable for now, for performance reasons
	
	def __init__(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		workbench.docObserver.onObjectTypeChanged('balloon_changed', 'TechDraw::DrawViewBalloon', lambda obj, prop: self.onBalloonChanged(obj, prop))
		workbench.docObserver.onObjectTypeSelected('balloon_selected', 'TechDraw::DrawViewBalloon', lambda operation, obj, sub, point: self.onBalloonSelected(operation, obj, sub, point))
	
	def repaint(self, view):
		self.views_to_repaint[view] = True
		QTimer.singleShot(10, self._do_repaint)
		
	def _do_repaint(self):
		from ahb_raster_view import RasterView
		
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		selection = Gui.Selection.getSelection()
		
		to_repaint = self.views_to_repaint.keys()
		self.views_to_repaint = {}
		
		for view in to_repaint:
			if '_overlay' in view.Label:
				continue
			#print("Repainting " + view.Name)
			
			page = self.getViewPage(view)
			
			view_cache = self.getViewCache(view)
			view_cache.reset()
			
			doc = view.Document
			
			if not 'Assembly_handbook_RasterView' in view.PropertiesList:
				view.addProperty("App::PropertyBool", "Assembly_handbook_RasterView", "Assembly_handbook")
				view.Assembly_handbook_RasterView = True
			
			if view.Assembly_handbook_RasterView:
				print("Rasterizing view " + view.Label + "...")
				
				raster_view = RasterView(view)
				raster_view.render()
				
				view.Visibility = False
				overlayName = view.Label + "_overlay"
				overlay = doc.getObject(overlayName)
				if overlay is None:
					overlay = doc.addObject('TechDraw::DrawViewPart', overlayName)
					page.addView(overlay)
				
				overlay_frame_name = view.Label + "_frame"
				overlay_frame = doc.getObject(overlay_frame_name)
				
				if overlay_frame is None:
					import Draft
					'''
					points = [App.Vector(0.0, 0.0, 0.0), App.Vector(1.01, 0, 0)]
					obj1 = Draft.makeWire(points, closed=False, face=False, support=None)
					
					points = [App.Vector(200.0, 0.0, 0.0), App.Vector(201.01, 0, 0)]
					obj2 = Draft.makeWire(points, closed=False, face=False, support=None)

					#overlay_frame = Draft.upgrade([], delete=True)
					#overlay_frame.Label = overlay_frame_name
					overlay_frame = App.ActiveDocument.addObject("Part::Part2DObjectPython", overlay_frame_name)
					Draft.Block(overlay_frame)
					overlay_frame.Components = [obj1, obj2]
					Draft.ViewProviderDraftPart(overlay_frame.ViewObject)
					doc.removeObject(obj1.Name)
					doc.removeObject(obj2.Name)'''
					
					#overlay_frame = Draft.makeWire(points, closed=False, face=False, support=None)
					overlay_frame = App.ActiveDocument.addObject("Part::Part2DObjectPython", overlay_frame_name)
					Draft.Wire(overlay_frame)
					pos = raster_view.projectImageViewPointTo3D(App.Vector(0,0,0))
					pos2 = raster_view.projectImageViewPointTo3D(App.Vector(0.001,0.001,1))
					overlay_frame.Points = [pos, pos2]
					Draft.ViewProviderWire(overlay_frame.ViewObject)
					overlay_frame.recompute()
					
				overlay_frame2_name = view.Label + "_frame2"
				overlay_frame2 = doc.getObject(overlay_frame2_name)
				if overlay_frame2 is None:
					overlay_frame2 = App.ActiveDocument.addObject("Part::Part2DObjectPython", overlay_frame2_name)
					Draft.Wire(overlay_frame2)
					pos = raster_view.projectImageViewPointTo3D(App.Vector(1,1,0))
					pos2 = raster_view.projectImageViewPointTo3D(App.Vector(1.001,1.001,1))
					overlay_frame2.Points = [pos, pos2]
					Draft.ViewProviderWire(overlay_frame2.ViewObject)
					overlay_frame2.recompute()
				
				overlay.Source = [overlay_frame, overlay_frame2]
				
				overlay.X = view.X
				overlay.Y = view.Y
				overlay.Direction = view.Direction
				overlay.XDirection = view.XDirection
				overlay.ScaleType = view.ScaleType
				overlay.Scale = view.Scale
				overlay.ViewObject.LineWidth = 0.01
				
				# migrate balloons from source view to overlay
				for balloon in page.Views:
					if balloon.TypeId == 'TechDraw::DrawViewBalloon' and "Assembly_handbook_Source" in balloon.PropertiesList and balloon.SourceView == view:
						if balloon.SourceView == view:
							old_source = balloon.Assembly_handbook_Source
							old_OriginOffsetX = balloon.Assembly_handbook_OriginOffsetX
							old_OriginOffsetY = balloon.Assembly_handbook_OriginOffsetY
							old_X = balloon.X
							old_Y = balloon.Y
							old_Visibility = balloon.ViewObject.Visibility
							balloonName = balloon.Name
							
							doc.removeObject(balloon.Name)
							
							balloon = doc.addObject("TechDraw::DrawViewBalloon", balloonName)
							balloon.SourceView = overlay
							balloon.addProperty("App::PropertyXLink", "Assembly_handbook_Source", "Assembly_handbook")
							balloon.Assembly_handbook_Source = old_source
							balloon.addProperty("App::PropertyFloat", "Assembly_handbook_OriginOffsetX", "Assembly_handbook")
							balloon.addProperty("App::PropertyFloat", "Assembly_handbook_OriginOffsetY", "Assembly_handbook")
							balloon.Assembly_handbook_OriginOffsetX = old_OriginOffsetX
							balloon.Assembly_handbook_OriginOffsetY = old_OriginOffsetY
							page.addView(balloon)
							self.updateBalloon(balloon)
							balloon.X = old_X
							balloon.Y = old_Y
							balloon.ViewObject.Visibility = old_Visibility
							balloon.recompute()
				
				overlay.recompute()
				page.recompute()
			else:
				fast_rendering = False
				#try:
				#	fast_rendering = view.Assembly_handbook_FastRendering
				#except:
				#	pass
				
				if view.CoarseView:
					fast_rendering = True
				
				selected_balloons = []
				for obj in Gui.Selection.getSelection():
					if obj.TypeId == 'TechDraw::DrawViewBalloon' and obj.SourceView == view and 'Assembly_handbook_Source' in obj.PropertiesList:
						selected_balloons.append(obj)
				
				#view.clearGeomFormats() # for an unknown reason, this will crash freecad
				
				if not fast_rendering:
					is_first_part = True
				
					parts_to_paint = []
				
					# repaint parts that are highlighted by selection
					if self.enable_selected_part_highlight:
						for balloon in selected_balloons:
							part = self.getBalloonSourcePart(balloon)
							if part in view.XSource:
								parts_to_paint.append(part)
							
					# repaint parts that are new in this step (thick line)
					prev_view = None
					if 'Assembly_handbook_PreviousStepView' in view.PropertiesList:
						prev_view = doc.getObject(view.Assembly_handbook_PreviousStepView)
					for part in view.XSource:
						if (prev_view is None or part not in prev_view.XSource) and part not in parts_to_paint:
							parts_to_paint.append(part)
				
					# make sure the list is not empty, so that we reset all lines
					if len(parts_to_paint) == 0:
						parts_to_paint.append(None)
					
					for part in parts_to_paint:
						default_line_thickness = 0.05
						line_thickness = default_line_thickness
						
						default_color = (0.0, 0.0, 0.0) if fast_rendering else (0.5, 0.5, 0.5)
						color = default_color
						
						if part is not None:
							part_view = workbench.partsCache.getPart2DView(view, part)
							
							center = self.computePartCenter(view, part)
						
							if self.isNewPartInView(view, part):
								line_thickness = 0.2
								color = (0, 0, 0)
							
							if self.enable_selected_part_highlight:			
								for balloon in selected_balloons:
									if part == self.getBalloonSourcePart(balloon):
										color = (0.0, 0.85, 0.0) # selection highlighting
						
						# iterate edges of actual view and highlight matching edges
						for edgeIdx in range(10000):
							hasEdge = False
							try:
								edge = view.getEdgeByIndex(edgeIdx)
								hasEdge = True
							except:
								pass
							if not hasEdge:
								break
							
							is_edge_of_part = False
							if part is not None and (not hasattr(edge.Curve, 'Degree') or edge.Curve.Degree == 1) and len(edge.Vertexes) == 2:
								edgeData = [
									edge.Vertexes[0].X - center.x,
									edge.Vertexes[0].Y - center.y,
									edge.Vertexes[1].X - center.x,
									edge.Vertexes[1].Y - center.y
								]
								v0 = App.Vector(edgeData[0], edgeData[1])
								v1 = App.Vector(edgeData[2], edgeData[3])
								
								for line in part_view.cached_lines:
									l0 = App.Vector(line[0], line[1])
									l1 = App.Vector(line[2], line[3])
									#d = abs(edgeData[0] - line[0]) + abs(edgeData[1] - line[1]) + abs(edgeData[2] - line[2]) + abs(edgeData[3] - line[3])
									d = v0.distanceToLineSegment(l0, l1).Length + v1.distanceToLineSegment(l0, l1).Length
									if d < 0.01:
										is_edge_of_part = True
										break
									
							if is_edge_of_part:
								view.formatGeometricEdge(edgeIdx,1,line_thickness,color,True)
							elif is_first_part:
								# reset edge format
								view.formatGeometricEdge(edgeIdx,1,default_line_thickness,default_color,True)
						
						is_first_part = False
				
						view.requestPaint()
	
	def updateBalloonCursor(self, view):
		selected_balloons = []
		for obj in Gui.Selection.getSelection():
			if obj.TypeId == 'TechDraw::DrawViewBalloon' and obj.SourceView == view and 'Assembly_handbook_Source' in obj.PropertiesList:
				selected_balloons.append(obj)
		
		cursor = self.view_cursors.get(view, None)
		if cursor is not None:
			try:
				cursor.x() # this can throw an exception if the Qt item has been deleted (for example when closing the page)
			except Exception as ex:
				print("Re-generating cursor...")
				cursor = None
		
		if cursor is None:
			cursor = CursorItem(view = view)
			TDG.addQGIToView(view, cursor);
			cursor.onViewPosChange(lambda new_pos: self.onCursorMoved(view, new_pos))
			self.view_cursors[view] = cursor
			
		if len(selected_balloons) == 0:
			cursor.setVisible(False)
		else:
			cursor.setViewPos(App.Vector(selected_balloons[0].OriginX, selected_balloons[0].OriginY))
			cursor.setVisible(True)
		
	def setCurrentViewDirection(self, view):
		from pivy import coin
		
		doc = view.Document
		if doc != Gui.ActiveDocument.Document:
			raise Exception("Current view is not for the same document as TechDraw view " + view.Name)
		activeView = Gui.ActiveDocument.ActiveView
		if str(type(activeView)) not in ["<class 'View3DInventorPy'>", "<class 'Gui.View3DInventor'>"]:
			raise Exception("Current view is not a 3D view")
		
		cam = activeView.getCameraNode()
		
		dir = cam.orientation.getValue().multVec(coin.SbVec3f(0,0,1)).getValue()
		xdir = cam.orientation.getValue().multVec(coin.SbVec3f(1,0,0)).getValue()
		view.Direction = App.Vector(dir[0], dir[1], dir[2])
		view.XDirection = App.Vector(xdir[0], xdir[1], xdir[2])
				
	def refreshView(self, view):
		doc = view.Document
		page = self.getViewPage(view)
		for balloon in page.Views:
			if balloon.TypeId == 'TechDraw::DrawViewBalloon' and "Assembly_handbook_Source" in balloon.PropertiesList and balloon.SourceView == view:
				obj = self.getBalloonSourcePart(balloon)
				balloonColor = (0.0, 0.0, 0.0)
				if obj is None or not obj in view.XSource:
					balloonColor = (1.0, 0.0, 0.0)
				balloon.ViewObject.Color = balloonColor
		
		self.view_cache[view] = None
				
	def toggleEditViewSourceParts(self, view):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		button = self.getToolbarButton('AHB_view_edit_source_parts')
		
		if self.edited_view is None:
			workbench.docLinkObserver.select_link_mode = True
			self.edited_view = view
			
			button.setChecked(True)
			button.setText('End source parts edition')
			
			Gui.Selection.clearSelection()
			for obj in view.XSource:
				Gui.Selection.addSelection(obj)
		else:
			workbench.docLinkObserver.select_link_mode = False
			
			Gui.Selection.clearSelection()
			Gui.Selection.addSelection(self.edited_view)
			self.edited_view = None
			
			button.setChecked(False)
			button.setText('Edit view source parts')
			
	def editViewSourceParts(self, parts, add):
		if self.edited_view is None: return
		xsource = self.edited_view.XSource
		modified = False
		for part in parts:
			if (part in xsource) != add:
				if add:
					xsource.append(part)
				else:
					xsource.remove(part)
				modified = True
		if modified:
			self.edited_view.XSource = xsource
		
	def getToolbarButton(self, buttonName):
		mainwin = Gui.getMainWindow()
		toolbar = None
		for tb in mainwin.findChildren(QtGui.QToolBar):
			if tb.objectName()=='Assembly Handbook':
				toolbar = tb
				
		button = None
		if toolbar is not None:
			for action in toolbar.actions():
				if action.objectName() == buttonName:
					button = action
		
		return button
				
	def onCursorMoved(self, view, new_pos):
		if len(Gui.Selection.getSelection()) == 0: return
		balloon = Gui.Selection.getSelection()[0]
		if balloon.TypeId != 'TechDraw::DrawViewBalloon' or not 'Assembly_handbook_Source' in balloon.PropertiesList: return
		if balloon.SourceView != view: return
		
		balloon.OriginX = new_pos.x
		balloon.OriginY = new_pos.y
		
		obj = self.getBalloonSourcePart(balloon)
		view = balloon.SourceView
		center = self.computePartCenter(view, obj)
		
		balloon.Assembly_handbook_OriginOffsetX = new_pos.x - center.x
		balloon.Assembly_handbook_OriginOffsetY = new_pos.y - center.y
	
	def onBalloonSelected(self, operation, balloon, sub, point):
		#print(operation, obj.Name, sub, point)
		if 'Assembly_handbook_Source' in balloon.PropertiesList:
			#print(operation + " " + balloon.Name)
			view = balloon.SourceView
			self.updateBalloonCursor(view)
			if self.enable_selected_part_highlight:
				self.repaint(view)
			
	def onBalloonChanged(self, obj, prop):
		# Avoid reentry
		if self.updating_balloon:
			return
		
		#print('Balloon changed: ' + obj.Name + '.' + prop)
		if prop == 'Y' and 'Assembly_handbook_Source' in obj.PropertiesList:
			self.updating_balloon = True
			self.updateBalloon(obj)
			self.updating_balloon = False
	
	def updateBalloon(self, balloon):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		view = balloon.SourceView
		obj = self.getBalloonSourcePart(balloon)
		
		partDisplayName = 'Inconnu' if obj is None else self.getPartDisplayName(obj)
		
		objectCenterView = workbench.techDrawExtensions.computePartCenter(view, obj)
		
		balloon.OriginX = objectCenterView.x + balloon.Assembly_handbook_OriginOffsetX
		balloon.OriginY = objectCenterView.y + balloon.Assembly_handbook_OriginOffsetY
		
		balloon.Text = partDisplayName
		balloon.ViewObject.Font = 'DejaVu Sans'
		balloon.ViewObject.Fontsize = 4
		balloon.BubbleShape = 'Inspection'
		balloon.EndTypeScale = 1
	
	def getBalloonSourcePart(self, balloon):
		try:
			return balloon.Assembly_handbook_Source[0]
		except:
			return None
	
	def isPartLink(self, obj):
		if obj.TypeId == 'App::Link':
			return True
		if obj.TypeId == 'Part::FeaturePython' and hasattr(obj, 'LinkedObject'): # variant link
			return True
		return False
		
	def getPartDisplayName(self, obj):
		if self.isPartLink(obj):
			linked_obj = obj.LinkedObject
			if 'Assembly_handbook_PartDisplayName' in linked_obj.PropertiesList:
				return linked_obj.Assembly_handbook_PartDisplayName
			else:
				return linked_obj.Document.Name
		return obj.Name
	
	def isNewPartInView(self, view, obj):
		doc = view.Document
		prev_view = None
		if 'Assembly_handbook_PreviousStepView' in view.PropertiesList:
			prev_view = doc.getObject(view.Assembly_handbook_PreviousStepView)
		
		if prev_view is None:
			return True
		else:
			if not obj in prev_view.XSource:
				return True
			return False
		
	def getActivePage(self):
		activeView = Gui.activeView()
		if activeView is None: return None
		activePage = activeView.getPage() if hasattr(activeView, 'getPage') else None
		return activePage
				
	def getViewPage(self, view):
		for obj in view.InList:
			if obj.TypeId == 'TechDraw::DrawPage':
				if view in obj.Views:
					return  obj
		return None
	
	def forceRedrawPage(self, page, callback = None):
		needPageUpdate = False
		for view in page.Views:
			if view.TypeId == 'TechDraw::DrawViewPart' and 'Assembly_handbook_PreviousStepView' in view.PropertiesList:
				if not 'Assembly_handbook_RasterView' in view.PropertiesList:
					view.addProperty("App::PropertyBool", "Assembly_handbook_RasterView", "Assembly_handbook")
					view.Assembly_handbook_RasterView = True
				
				if 'Assembly_handbook_RasterView' not in view.PropertiesList or not view.Assembly_handbook_RasterView:
					needPageUpdate = True
				self.refreshView(view)
			elif view.TypeId == 'TechDraw::DrawViewBalloon':
				if view.ViewObject.Visibility:
					# workaround for a TechDraw bug: sometimes the balloon should be visible but doesn't appear, showing it again fixes the issue
					view.ViewObject.Visibility = False
					def makeRedrawCallback(view):
						def redrawBalloon():
							view.ViewObject.Visibility = True
						return redrawBalloon
					QTimer.singleShot(0, makeRedrawCallback(view))
				else:
					# workaround for a TechDraw bug: sometimes the balloon text is visible even if the balloon is hidden, hiding it again fixes the issue
					view.ViewObject.Visibility = True
					view.ViewObject.Visibility = False
		
		if page.KeepUpdated or not needPageUpdate:
			for view in page.Views:
				if view.TypeId != 'TechDraw::DrawViewPart':
					view.recompute()
			for view in page.Views:
				if view.TypeId == 'TechDraw::DrawViewPart':
					view.recompute()
					self.repaint(view)
			if callback is not None:
				callback()
		else:
			page.KeepUpdated = True
			def restoreKeepUpdated():
				page.KeepUpdated = False
				for view in page.Views:
					if view.TypeId == 'TechDraw::DrawViewPart':
						self.repaint(view)
				if callback is not None:
					callback()
			QTimer.singleShot(10, restoreKeepUpdated)
	
	def getSourceView(self, view):
		if view.Name.endswith('_overlay'):
			view = view.Document.getObject(view.Name[0:-8])
			if view is None:
				raise Exception("Can't find source view of " + view.Name)
		
		return view
	
	def getOverlayView(self, view):
		if view.Name.endswith('_overlay'):
			return view
		
		overlay = view.Document.getObject(view.Name + '_overlay')
		return overlay if overlay is not None else view
	
	def computePartCenter(self, view, obj):
		view = self.getSourceView(view)
		
		if obj.TypeId == 'App::Link':
			partLink = obj
			objectCenterWorld = partLink.LinkPlacement.Matrix.multiply(partLink.LinkedObject.Shape.CenterOfGravity)
		elif obj.TypeId == 'Part::FeaturePython' and hasattr(obj, 'LinkedObject'): # variant link
			partLink = obj
			objectCenterWorld = partLink.Placement.Matrix.multiply(partLink.LinkedObject.Shape.CenterOfGravity)
		else:
			objectCenterWorld = obj.Shape.CenterOfGravity
			
		'''view_cache = self.getViewCache(view)
		
		key = (objectCenterWorld.x, objectCenterWorld.y, objectCenterWorld.z)
		projected_point = view_cache.projected_points.get(key, None)
		
		if projected_point is None:
			# TechDraw does not expose a way to project a 3D point to 2D view coordinates ; this is a hack to get this value indirectly. The view should be hidden before calling this method, to avoid costly repaints.
			vertId = view.makeCosmeticVertex3d(objectCenterWorld)
			vert = view.getCosmeticVertex(vertId)
			projected_point = vert.Point
			view.removeCosmeticVertex(vertId)
			view_cache.projected_points[key] = projected_point
		
		return projected_point'''
		return self.projectPoint(view, objectCenterWorld)
	
	def projectPoint(self, view, point3d):
		if 'Assembly_handbook_RasterView' in view.PropertiesList and view.Assembly_handbook_RasterView:
			from ahb_raster_view import RasterView
			raster_view = RasterView(view)
			if raster_view.init_image_projection():
				return raster_view.project3DPointToSourceView(point3d)
		
		# DrawViewPart::projectPoint should be exposed to python in freecad 0.21, but for 0.20 we have to use a workaround
		view_cache = self.getViewCache(view)
		if view_cache.projected_origin is None:
			vertId = view.makeCosmeticVertex3d(App.Vector(0,0,0))
			vert = view.getCosmeticVertex(vertId)
			view_cache.projected_origin = vert.Point
			view.removeCosmeticVertex(vertId)
		
		YDirection = view.Direction.cross(view.XDirection)
		return App.Vector(view.XDirection.dot(point3d) + view_cache.projected_origin.x, YDirection.dot(point3d) + view_cache.projected_origin.y, 0)
	
	def getViewCache(self, view):
		cache = self.view_cache.get(view, None)
		if cache is None:
			cache = ViewCache()
			self.view_cache[view] = cache
		return cache