import FreeCAD as App
import FreeCADGui as Gui

class RasterView:
	def __init__(self, view):
		self.source_view = view
		doc = view.Document
		self.image_file_name = doc.FileName.replace('.FCStd', '') + '_raster/' + view.Name + '.png'
		
	def init_image_projection(self):
		doc = self.source_view.Document
		image_name = self.source_view.Label + "_raster"
		image = doc.getObject(image_name)
		if image is None:
			return False
		
		self.image_view = image
		
		if image.Assembly_handbook_ViewVolumeWidth > 0:
			self._precompute_image_projection()
			return True
		
		return False
		
	def init_image(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		doc = self.source_view.Document
		page = workbench.techDrawExtensions.getViewPage(self.source_view)
		
		image_name = self.source_view.Label + "_raster"
		image = doc.getObject(image_name)
		if image is None:
			image = doc.addObject('TechDraw::DrawViewImage', image_name)
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeWidth", "Assembly_handbook")
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeHeight", "Assembly_handbook")
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeDepth", "Assembly_handbook")
			image.addProperty("App::PropertyVector", "Assembly_handbook_ViewVolumeOffset", "Assembly_handbook")
		
		if not image in page.Views:
			page.addView(image)
			
			new_views_list = page.Views
			new_views_list.remove(image)
			view_idx = new_views_list.index(self.source_view)
			new_views_list.insert(view_idx, image)
			page.Views = new_views_list
		
		self.image_view = image
		
		if image.Assembly_handbook_ViewVolumeWidth > 0:
			self._precompute_image_projection()
			
	def _precompute_image_projection(self):
		YDirection = self.source_view.Direction.cross(self.source_view.XDirection)
		
		self.image_x_dir = self.source_view.XDirection / self.image_view.Assembly_handbook_ViewVolumeWidth
		self.image_y_dir = YDirection / self.image_view.Assembly_handbook_ViewVolumeHeight
		self.image_z_dir = self.source_view.Direction / self.image_view.Assembly_handbook_ViewVolumeDepth
		
		self.image_x_dir_inv = self.source_view.XDirection * self.image_view.Assembly_handbook_ViewVolumeWidth
		self.image_y_dir_inv = YDirection * self.image_view.Assembly_handbook_ViewVolumeHeight
		self.image_z_dir_inv = self.source_view.Direction * self.image_view.Assembly_handbook_ViewVolumeDepth
		
	def project3DPointToImageView(self, point3d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		return App.Vector(self.image_x_dir.dot(point3d) + offset.x, self.image_y_dir.dot(point3d) + offset.y, self.image_z_dir.dot(point3d) + offset.z)
	
	def project3DPointToSourceView(self, point3d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		offset = App.Vector((offset.x-0.5) * self.image_view.Assembly_handbook_ViewVolumeWidth, (offset.y-0.5) * self.image_view.Assembly_handbook_ViewVolumeHeight, (offset.z-0.5) * self.image_view.Assembly_handbook_ViewVolumeDepth)
		#image_view_point = App.Vector(self.image_x_dir.dot(point3d), self.image_y_dir.dot(point3d), self.image_z_dir.dot(point3d))
		#return App.Vector(image_view_point.x * self.image_view.Assembly_handbook_ViewVolumeWidth, image_view_point.y * self.image_view.Assembly_handbook_ViewVolumeHeight, 0)
		
		YDirection = self.source_view.Direction.cross(self.source_view.XDirection)
		return App.Vector(self.source_view.XDirection.dot(point3d) + offset.x, YDirection.dot(point3d) + offset.y, self.image_z_dir.dot(point3d) + offset.z)
	
	def projectImageViewPointTo3D(self, point2d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		p = point2d - offset
		return self.image_x_dir_inv * p.x + self.image_y_dir_inv * p.y + self.image_z_dir_inv * p.z
	
	def _flatten_objects_tree(self, obj_list):
		result = []
		for obj in obj_list:
			if obj.TypeId == 'Part::FeaturePython' and hasattr(obj, 'LinkedObject'): # variant link
				result.extend(self._flatten_objects_tree(obj.Group))
			elif obj.TypeId in ['App::Link']:
				result.extend(self._flatten_objects_tree([obj.LinkedObject]))
			elif obj.TypeId in ['App::Part', 'App::DocumentObjectGroup']:
				result.extend(self._flatten_objects_tree(obj.Group))
			elif self._should_render(obj) or self._should_render_as_is(obj) or obj.TypeId in ['PartDesign::CoordinateSystem', 'PartDesign::Line']:
				result.append(obj)
				if hasattr(obj, 'Group'):
					result.extend(self._flatten_objects_tree(obj.Group))
			
		return result
	
	def _should_render(self, obj):
		return obj.TypeId in ['Part::Feature', 'Part::FeaturePython', 'PartDesign::Body', 'Part::Mirroring', 'Part::Cut', 'Part::Part2DObjectPython', 'Part::MultiFuse', 'Part::Loft', 'Part::Torus', 'Part::Cylinder']
	
	def _should_render_as_is(self, obj):
		return obj.TypeId in ['App::FeaturePython']
		
	def render(self, fast_render = True):
		from pivy import coin
		import os
		from PIL import Image, ImageDraw, ImageChops
		import Part
		
		Image.MAX_IMAGE_PIXELS = 9999999999 # allow very high resolution images
		
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		view = self.source_view
		
		self.init_image()
		
		print('Rasterizing ' + view.Label + " to " + self.image_file_name + "...")
		
		dir = os.path.dirname(self.image_file_name)
		if not os.path.exists(dir):
			os.makedirs(dir)
			
		tmp_doc = App.newDocument('tmp_raster', hidden=False, temp=False)
			
		objects_to_reset = {}
		duplicated_parts = {}
		try:
			# construct new scene with links to the parts we want
			sceneGroup = tmp_doc.addObject('App::DocumentObjectGroup', 'Scene')
			prev_parts = []
			new_parts = []
			all_parts = view.XSource + view.Source
			for part in all_parts:
				link = tmp_doc.addObject('App::Link', part.Name)
				link.Label = part.Label
				if part.TypeId == 'App::Link':
					link.LinkedObject = part.LinkedObject
					link.Placement = part.Placement
				elif part.TypeId == 'Part::FeaturePython' and hasattr(part, 'LinkedObject'): # variant link
					link.LinkedObject = part.LinkedObject
					link.Placement = part.Placement
				else:
					link.LinkedObject = part
					if part.TypeId in ['Part::Part2DObjectPython']:
						link.Placement = part.Placement
					
				is_new_part = workbench.techDrawExtensions.isNewPartInView(view, part)
				
				if not fast_render:
					# check if another part with different render settings will conflict with ours
					# a conflict occurs when two parts link to the same object (directly or indirectly), because render settings (such as color) are set at the object level
					is_conflicting = False
					if link.LinkedObject in duplicated_parts.keys():
						link.LinkedObject = duplicated_parts[link.LinkedObject]
					else:
						other_parts = prev_parts if is_new_part else new_parts
						for other_part in other_parts:
							other_objects = self._flatten_objects_tree([other_part])
							for obj in self._flatten_objects_tree([link]):
								if self._should_render(obj) and obj in other_objects:
									is_conflicting = True
								
						if is_conflicting:
							# We must copy the part because otherwise we can't control the emissive color (link material override does not work for emissive color)
							#print("conflict: " + link.LinkedObject.Document.Name + '#' + link.LinkedObject.Label)
							shape_copy = Part.getShape(link.LinkedObject,'',needSubElement=False,refine=False)
							part_copy = tmp_doc.addObject('Part::Feature','ShapeCopy')
							part_copy.Shape = shape_copy
							part_copy.Label = part.Label
							duplicated_parts[link.LinkedObject] = part_copy
							link.LinkedObject = part_copy
							part_copy.ViewObject.Visibility = False
				
				sceneGroup.addObject(link)
				
				if is_new_part:
					new_parts.append(link)
				else:
					prev_parts.append(link)
					
				# hide objects that we don't want to display ; also make a backup of properties we want to reset after we're done
				for obj in self._flatten_objects_tree([link]):
					if obj in objects_to_reset.keys():
						continue
					
					if self._should_render(obj):
						if not fast_render:
							objects_to_reset[obj] = (
								obj.ViewObject.Visibility,
								obj.ViewObject.LineColor,
								obj.ViewObject.ShapeMaterial.AmbientColor,
								obj.ViewObject.ShapeMaterial.DiffuseColor,
								obj.ViewObject.ShapeMaterial.SpecularColor,
								obj.ViewObject.ShapeMaterial.EmissiveColor,
								obj.ViewObject.LineWidth,
								obj.ViewObject.DisplayMode
							)
							
							if not obj.ViewObject.Visibility:
								obj.ViewObject.ShapeMaterial.AmbientColor = (0, 0, 0)
								obj.ViewObject.ShapeMaterial.DiffuseColor = (0, 0, 0)
								obj.ViewObject.ShapeMaterial.SpecularColor = (0, 0, 0)
								obj.ViewObject.ShapeMaterial.EmissiveColor = (0, 0, 0)
					else:
						objects_to_reset[obj] = (
							obj.ViewObject.Visibility,
						)
							
						obj.ViewObject.Visibility = False
						
			tmp_doc_view = Gui.getDocument(tmp_doc.Name).mdiViewsOfType('Gui::View3DInventor')[0]
		
			cam = tmp_doc_view.getCameraNode()
			
			rot = coin.SbRotation(coin.SbVec3f(1,0,0), coin.SbVec3f(view.XDirection.x,view.XDirection.y,view.XDirection.z))
			rot *= coin.SbRotation(coin.SbVec3f(0,0,1), coin.SbVec3f(view.Direction.x,view.Direction.y,view.Direction.z))
			cam.orientation.setValue(rot)
			
			targetViewVolume = None
			try:
				targetViewVolume = view.Assembly_handbook_ViewVolume
			except:
				pass
			
			if targetViewVolume is None:
				tmp_doc_view.fitAll()
			else:
				sceneGroup.ViewObject.Visibility = False
				viewVolumeLink = tmp_doc.addObject('App::Link', 'ViewVolume')
				viewVolumeLink.LinkedObject = targetViewVolume
				viewVolumeLink.Placement = targetViewVolume.Placement
				tmp_doc_view.fitAll()
				tmp_doc.removeObject(viewVolumeLink.Name)
				sceneGroup.ViewObject.Visibility = True
			
			viewVolume = cam.getViewVolume(0.0)
			self.image_view.Assembly_handbook_ViewVolumeWidth = viewVolume.getWidth()
			self.image_view.Assembly_handbook_ViewVolumeHeight = viewVolume.getHeight()
			self.image_view.Assembly_handbook_ViewVolumeDepth = viewVolume.getDepth()
			
			max_res = 3200
			#max_res = 1500
			resolution = [
				int(viewVolume.getWidth() * view.Scale * 10),
				int(viewVolume.getHeight() * view.Scale * 10)
			]
			if resolution[0] > max_res:
				resolution[1] = int(resolution[1] * max_res / resolution[0])
				resolution[0] = int(max_res)
			if resolution[1] > max_res:
				resolution[0] = int(resolution[0] * max_res / resolution[1])
				resolution[1] = int(max_res)
			
			if fast_render:
				composite_img = self._render_lines(tmp_doc, resolution, prev_parts + new_parts, (0.0, 0.0, 0.0), [])
			else:
				# render old parts in gray lines
				prev_parts_img = self._render_lines(tmp_doc, resolution, prev_parts, (0.6, 0.6, 0.6), [], fast_render)
				
				# render new parts in black lines (old parts can mask them)
				new_parts_img = self._render_lines(tmp_doc, resolution, new_parts, (0.0, 0.0, 0.0), prev_parts, fast_render)
				
				# create the composite image
				composite_img = prev_parts_img.copy()
				composite_img.paste(new_parts_img, None, new_parts_img)
			
			# Optimize the image to reduce storage size
			if not fast_render:
				num_colors = 32
				
				# All-or-nothing alpha: we use a white background and only make pixels fully transparent where alpha is zero, to not loose antialiasing
				bg_img = Image.new(composite_img.mode, composite_img.size, color = '#ffffff')
				bg_img.paste(composite_img.convert('RGB'), composite_img)
				final_alpha = composite_img.split()[3].point(lambda p: 0 if p <= int(255/num_colors+0.5) else 255)
				composite_img = bg_img
				composite_img.putalpha(final_alpha)
				
				# Convert to indexed colors
				composite_img = composite_img.quantize(colors=num_colors, dither=Image.Dither.NONE)
			
		finally:
			#raise Exception("test")
			# restore properties on objects we have modified
			for obj, props in objects_to_reset.items():
				obj.ViewObject.Visibility = props[0]
				if self._should_render(obj):
					obj.ViewObject.LineColor = props[1]
					obj.ViewObject.ShapeMaterial.AmbientColor = props[2]
					obj.ViewObject.ShapeMaterial.DiffuseColor = props[3]
					obj.ViewObject.ShapeMaterial.SpecularColor = props[4]
					obj.ViewObject.ShapeMaterial.EmissiveColor = props[5]
					obj.ViewObject.LineWidth = props[6]
					obj.ViewObject.DisplayMode = props[7]
			
			# remove the temporary document
			App.closeDocument(tmp_doc.Name)
		
		# Crop the image, which is also used to deduce the center of the source view
		original_size = composite_img.size
		
		diff_source_img = composite_img.split()[-1]
		bg = Image.new(diff_source_img.mode, diff_source_img.size, '#000000') # fills an image with the background color
		diff = ImageChops.difference(diff_source_img, bg) # diff between the actual image and the background color
		bbox = diff.getbbox() # finds border size (non-black portion of the image)
		composite_img = composite_img.crop(bbox)
		
		'''draw = ImageDraw.Draw(composite_img)
		
		def debugPoint(p3d):
			p2d = self.project3DPointToImageView(p3d)
			pp = App.Vector(p2d.x * original_size[0] - bbox[0], (1.0-p2d.y) * original_size[1] - bbox[1])
			#print('pp', pp)
			
			len = 100
			draw.line([(pp.x, pp.y-len), (pp.x, pp.y+len)], fill=128, width = 7)
			draw.line([(pp.x-len, pp.y), (pp.x+len, pp.y)], fill=128, width = 7)
		
		debugPoint(App.Vector(-12.5, 37.5, 25.0))
		debugPoint(App.Vector(-12.5, -1387.5, 25.0))
		debugPoint(App.Vector(131.23702882966705, -655.0000021095163, 145.21130178331268))'''
		
		composite_img.save(self.image_file_name)
		
		sb_offset = viewVolume.projectToScreen(coin.SbVec3f(0,0,0))
		crop_offset = App.Vector(((bbox[0] + bbox[2])/2 - original_size[0]/2)/original_size[0], ((bbox[1] + bbox[3])/2 - original_size[1]/2)/original_size[1], 0)
		self.image_view.Assembly_handbook_ViewVolumeOffset = App.Vector(sb_offset[0] - crop_offset.x, sb_offset[1] + crop_offset.y, sb_offset[2])
		
		self._precompute_image_projection()
		
		p2dA = self.project3DPointToImageView(App.Vector(0,0,0))
		p2dB = self.project3DPointToImageView(view.XDirection)
		image_scale = view.Scale / (p2dB.x - p2dA.x) / original_size[0] * 10
		
		# display the image in the view
		image = self.image_view
		image.ImageFile = ""
		image.Scale = image_scale
		image.X = view.X
		image.Y = view.Y
		image.ImageFile = self.image_file_name
		image.ViewObject.Crop = True
		image.Width = composite_img.size[0] * image_scale / 10.0 * 1.01
		image.Height = composite_img.size[1] * image_scale / 10.0 * 1.01
		image.recompute()
			
	def _render_lines(self, doc, resolution, parts, line_color, masking_parts, fast_render = True):
		import tempfile
		from PIL import Image, ImageDraw, ImageFilter
		
		doc_view = Gui.getDocument(doc.Name).mdiViewsOfType('Gui::View3DInventor')[0]
		
		# render lines in black, background in red, fill shapes in green
		# the green band contains the lines images, the red band contains the inverted alpha layer
		for link in doc.findObjects():
			if link in parts or link in masking_parts:
				link.ViewObject.Visibility = True
				
				# in current version of freecad, link override material does not allow to override all material properties, for example emissive color, so we have to change material of the linked object
				for obj in self._flatten_objects_tree([link]):
					if self._should_render(obj) and not fast_render:
						obj.ViewObject.LineColor = (0.0, 0.0, 0.0, 0.0) if link in parts else (1.0, 0.0, 1.0)
						obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
						obj.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0)
						obj.ViewObject.ShapeMaterial.SpecularColor = (0.0, 0.0, 0.0, 0.0)
						obj.ViewObject.ShapeMaterial.EmissiveColor = (0.0, 1.0, 0.0, 0.0) if link in parts else (1.0, 0.0, 1.0)
						
						# We need to set two different values otherwise freecad does not always update LineWidth of sub-elements
						obj.ViewObject.LineWidth = 1.0
						obj.ViewObject.LineWidth = 2.0
			else:
				link.ViewObject.Visibility = False
		
		temp_file_name = tempfile.gettempdir() + "/ahb_temp_image.png"
		doc_view.saveImage(temp_file_name, resolution[0]+2, resolution[1]+2, "#ff0000") # we add 1 pixel border that we will need to crop later
		lines_bands_img = self._read_image(temp_file_name)
		
		lines_bands = lines_bands_img.split()
		lines_img = lines_bands[1]
		alpha_img = lines_bands[0].point(lambda p: 255 - p)
		
		generate_outlines = not fast_render
		if generate_outlines:
			# Render all shapes with different colors, in order to extract outlines (where color changes)
			# This is needed because FreeCAD does not render lines on the boundary of curve shapes, such as spheres or cylinders
			# The technique could be improved by using the depth buffer instead, in order to detect boundaries within the same object
			step = 8
			r = step
			g = step
			b = step
			for link in doc.findObjects():
				if link in parts or link in masking_parts:
					for obj in self._flatten_objects_tree([link]):
						if self._should_render(obj) and obj.TypeId != 'Part::Part2DObjectPython':
							obj.ViewObject.DisplayMode = 'Shaded'
							obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
							obj.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0)
							obj.ViewObject.ShapeMaterial.SpecularColor = (0.0, 0.0, 0.0, 0.0)
							obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0) if link in parts else (1.0, 1.0, 1.0, 0.0) 
							
							r = r + step
							if r >= 256 - step:
								r = step
								g = g + step
								if g >= 256 - step:
									g = step
									b = b + step
									if b >= 256 - step:
										b = step
						else:
							obj.ViewObject.Visibility = False
			
			doc_view.saveImage(temp_file_name, (resolution[0]+2)*2, (resolution[1]+2)*2, "#ffffff") # shapes are rendered at twice the resolution for antialiasing
			shapes_img = self._read_image(temp_file_name)
			
			outlines_img = None
			for x in range(0, 3):
				for y in range(0, 3):
					if x == 1 and y == 1: continue
					kernel = [0, 0, 0, 0, 1, 0, 0, 0, 0]
					kernel[y * 3 + x] = -1
					partial_outlines = shapes_img.filter(ImageFilter.Kernel((3, 3), kernel, 1, 127))
					partial_outlines = partial_outlines.point(lambda p: 255 if p == 127 else 0)
					partial_outlines = partial_outlines.convert("L")
					partial_outlines = partial_outlines.point(lambda p: 255 if p == 255 else 0)
					if outlines_img is None:
						outlines_img = partial_outlines
					else:
						outlines_img.paste(partial_outlines, None, partial_outlines.point(lambda p: 0 if p == 255 else 255))
		
			lines_fullres = lines_img.resize(outlines_img.size, Image.NEAREST)
			lines_fullres.paste(outlines_img, None, outlines_img.point(lambda p: 255 if p == 0 else 0))
			#lines_fullres.paste(255, alpha_fullres.point(lambda p: 255 if p == 0 else 0))
			all_lines = lines_fullres.resize(lines_img.size, Image.BILINEAR)
			#all_lines = lines_img.copy()
			
			alpha_fullres = alpha_img.resize(outlines_img.size, Image.NEAREST)
			alpha_fullres.paste(outlines_img.point(lambda p: 255), None, outlines_img.point(lambda p: 255 if p == 0 else 0))
			alpha_img = alpha_fullres.resize(all_lines.size, Image.BILINEAR)
		else:
			all_lines = lines_img
			alpha_img = alpha_img.point(lambda p: 0 if p == 0 else 255)
		
		# colorize final image
		fill_color = (1.0, 1.0, 1.0)
		result = Image.merge("RGBA", [
			all_lines.point(lambda p: int(fill_color[0] * p + line_color[0] * (255.0 - p))),
			all_lines.point(lambda p: int(fill_color[1] * p + line_color[1] * (255.0 - p))),
			all_lines.point(lambda p: int(fill_color[2] * p + line_color[2] * (255.0 - p))),
			alpha_img
		])
		
		# crop 1px borders
		result = result.crop((1, 1, result.size[0] - 1, result.size[1] - 1))
		
		return result
	
	def _read_image(self, file_name):
		from PIL import Image
		with Image.open(file_name) as image:
			return image.copy()