# Plots a 3D tensor as a volume by showing its first slice along each mode

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import imageio
from copy import deepcopy

dpi = 200

def plot_tensor(input_paths, dtype, sizes, output_paths, selected_modes=[0, 1, 2], logarithmic=False):
	
	# selected_modes corresponds to [right_slice, left_slice, top_slice]
	
	if type(input_paths) is not list:
		input_paths = [input_paths,]
		output_paths = [output_paths,]
	
	slices_per_tensor = []
	min_value = None
	max_value = None
	
	print("Plotting %s versions of tensor ..."%len(input_paths))
	
	for input_path in input_paths:
		
		# Load file and extract transposed 3D slice
		data = np.fromfile(input_path, dtype).reshape(sizes)
		data_shape = deepcopy(data.shape)
		modes = list(range(len(sizes)))
		for selected_mode in selected_modes:
			modes.remove(selected_mode)
		data = data.transpose(selected_modes + modes)
		data = data[tuple([slice(0, data.shape[i]) for i in range(3)] + [0,]*max(0, len(sizes) - 3))]
		
		# Extract slices and calculate data bounds
		slices = []
		for i in range(3):
			slices.append(data[tuple([slice(0, data.shape[j]) for j in range(i)] + [0,] + [slice(0, data.shape[j]) for j in range(i + 1, 3)])])
			if logarithmic:
				# Consider logarithms of absolute value instead
				slices[-1] = np.log(np.abs(slices[-1]))
			if min_value is None:
				min_value = np.amin(slices[-1])
				max_value = np.amax(slices[-1])
			else:
				min_value = min(min_value, np.amin(slices[-1]))
				max_value = max(max_value, np.amax(slices[-1]))
		
		slices_per_tensor.append(slices)
		
		print("Extracted slices from %s"%input_path)
		
		del data
	
	for tensor_index in range(len(input_paths)):
		
		# Draw plot
		fig = plt.figure()
		ax = plt.axes(projection="3d")
		ax.set_box_aspect([data_shape[i] for modes in selected_modes])
		for i in range(3):
			
			# Draw an image on the surface orthogonal to e_i
			current_slice = (slices_per_tensor[tensor_index][i] - min_value)/(max_value - min_value)
			
			# create a 21 x 21 vertex mesh
			x_coords_in_image, y_coords_in_image = np.meshgrid(np.linspace(0, current_slice.shape[0], current_slice.shape[0] + 1), np.linspace(0, current_slice.shape[1], current_slice.shape[1] + 1), indexing="ij")
			
			# create vertices for a rotated mesh (3D rotation matrix)
			zero_array = np.zeros(x_coords_in_image.shape)
			X = zero_array if i == 0 else x_coords_in_image
			Y = x_coords_in_image if i == 0 else (zero_array if i == 1 else y_coords_in_image)
			Z = zero_array if i == 2 else y_coords_in_image
			
			ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=plt.cm.coolwarm(current_slice), shade=False, vmin=0, vmax=1)
			
			# Draw edges around face
			#line_coords = [[0, meshgrid[0, -1], meshgrid[-1, -1], meshgrid[-1, 0]] for meshgrid in (X, Y, Z)]
			#ax.plot(line_coords[0], line_coords[1], line_coords[2])
		
		# Finalize plot
		ax.grid(False) # Hide grid lines
		ax.set_axis_off()
		ax.invert_xaxis()
		ax.invert_zaxis()
		output_path = output_paths[tensor_index]
		plt.savefig(output_path, bbox_inches="tight", dpi=dpi, transparent=True)
		print("Saved image as %s"%output_path)
		trim_image(output_path)
		
		# Delete redundant data
		slices_per_tensor[tensor_index] = None

# Trims whitespace from all slides of an image
def trim_image(path):
	image = imageio.imread(path)
	not_transparent = image[:, :, 3] != 0
	not_white = np.any(image[:, :, :3] != 255, axis=2)
	boolean_image = np.logical_and(not_transparent, not_white)
	slicing = [slice(0, image.shape[j]) for j in range(3)]
	for i in range(2):
		collapsed_image = np.any(boolean_image, axis=(i + 1)%2)
		nonzeros = np.nonzero(collapsed_image)[0]
		slicing[i] = slice(nonzeros[0], nonzeros[-1] + 1)
		
	imageio.imwrite(path, image[tuple(slicing)])

if __name__ == "__main__":
	#plot_tensor("../../Research/compression-comparison/data/hyperspectral_images/moffett_field_float32", "float32", [512, 1024, 224], "../../Presentations/ATC/Poster/images/moffett_field_coolwarm_rotated.png", selected_modes=[0, 1, 2])
	#plot_tensor("../../Research/compression-comparison/data/turbulence/isotropic1024fine-pressure-time-128", "float32", [100, 128, 128, 128], "../../Presentations/ATC/Poster/images/isotropic_pressure_time_coolwarm.png", selected_modes=[1, 2, 3])
	plot_tensor("../../Manuscripts/atc-paper/data/isotropic1024-fine-pressure-time-128-core_128x128x128x100", "float64", [128, 128, 128, 100], "../../Presentations/ATC/Poster/images/isotropic_pressure_time_untruncated_core_log_coolwarm.png", selected_modes=[0, 1, 2], logarithmic=True)

