"""
Author: David J. Giesen
	Eastman Kodak Company

Shows a molecule on a Tk canvas.  The molecule can be rotated, and if it has frequencies, the frequencies can be animated
"""

import Tkinter as Tk
import giepy as gp
import math

# Colors of each element ball, padded to account for no element 0
element_colors = ['grey', 'light grey',					'light blue',
		'grey', 'blue', 'green', 'black', 'blue', 'red', 'green', 'light green',
		'brown', 'dark grey', 'grey', 'tan', 'pink', 'yellow', 'dark green', 'dark red',
		'yellow', 'white', 'grey', 'grey', 'grey', 'green', 'grey', 'dark red', 'dark blue', 'grey', 'orange', 'purple', 'yellow', 'dark green', 'pink', 'green', 'brown', 'green',
		'red', 'yellow', 'grey', 'grey', 'grey', 'grey', 'grey', 'grey', 'grey', 'grey', 'grey', 'white', 'brown', 'grey', 'pink', 'green', 'purple', 'green',
		'green', 'white'] + ['grey']*14 + ['grey', 'grey', 'grey', 'black', 'grey', 'blue', 'green', 'grey', 'gold', 'grey', 'brown', 'grey', 'cyan', 'red', 'purple', 'green'] + ['grey'] * 32

# Ball radii for ball and stick (padded to have a 0 for element zero)
ball_radii = [0, 0.20,							0.30,
	0.24, 0.26,			0.30, 0.33, 0.33, 0.33, 0.30, 0.40,
	0.35, 0.38,			0.42, 0.45, 0.45, 0.45, 0.40, 0.50] + [0.58]*18 + [0.63]*32 + [0.70]*32

# horizon is the z-coordinate where everything would converge to the center of the screen.  This helps provide some perspective to the image
# by making things further away seem smaller and slanted slightly to the center of the image.
horizon = -50.

# To find a vector perpendicular to vectors a and b, take the cross product (|a||b|sin(n)) where n is the angle between a and b
# |a| is just the magnitude of a
#
# Rotation about an arbitrary axis: http://inside.mines.edu/~gmurray/ArbitraryAxisRotation/

class MolPlot(Tk.Canvas):
	def __init__(self, parent, molecule=None, height=500, width=800, bg='white', bd=0, highlightthickness=0, stopsign=False, **kwargs):
		#
		# parent = container for this widget
		# molecule = molecule to be displayed.  Should be of a class that has an interface equivalent to the giechem.Molecule() class
		# height, width = height & width of the widget in pixels
		# bg = color of the background
		# bd = width of the canvas border
		# highlightthickness = highlightthickness of this widget (passed to Tk.Canvas)
		# stopsign = True to include a stopsign icon in the upper left of the widget.  The stopsign can be used to stop a frequency
		#			animation
		#
		# Any additional keyword parameters are passed to the Tk.Canvas widget
		#
		# To change the molecule that is displayed, use MolPlot.change_molecule(molecule)
		# To clear the widget, use MolPlot.delete_molecule()
		# To animate a frequency, use MolPlot.animate_frequency(frequency)
		#	The frequency class is found in Frequency Visualizer.py
		#
		self.parent = parent
		Tk.Canvas.__init__(self, self.parent, height=height, width=width, bg=bg, bd=bd, highlightthickness=highlightthickness, **kwargs)
		if molecule:
			self.change_molecule(molecule)
		else:
			self.molecule = None
		self.define_center()
		self.track_start_x = None
		self.track_start_y = None
		self.bind("<B1-Motion>", self.rotation_tracking)
		self.bind("<ButtonRelease-1>", self.stop_rotation_tracker)
		self.frames = 10.
		self.frequency_running = False
		self.frequency_scale = 5.0
		self.frequency_speed = 10.
		self.after_id = 0
		if stopsign:
			stopsign = self.create_polygon((3,6), (3,12), (6,16), (11,16), (15,12), (15,6), (12,3), (6,3), fill='red')
			self.tag_bind(stopsign, '<Button-1>', self.stop_frequency)

	def provide_perspective(self, vibrating=False):
		# Provides visual perspective to the image by converging everything towards the center of the screen and reducing its size
		# based on how close it is to the horizon point.
		# First, find the most +'ve Z coordinate
		# If vibrating, use the vibration coordinates
		total_z = self.molecule.foreground - self.molecule.horizon
		# Now every atom gets perspective based on where it lies between pos_z and the horizon
		for atom in self.molecule.atoms:
			if vibrating:
				coords = atom.vib_coords
			else:
				coords = atom.coords
			perspective = (coords[2] - self.molecule.horizon) / total_z
			atom.perspective_radius = atom.ball_radius * perspective
			atom.perspective_coords = [coords[0] * perspective, coords[1] * perspective, coords[2] ]

	def start_rotation_tracker(self, event):
		self.track_start_x = event.x
		self.track_start_y = event.y
		self.bind("<Motion>", self.rotation_tracking)
		self.bind("<ButtonRelease-1>", self.stop_rotation_tracker)

	def stop_rotation_tracker(self, event):
		self.track_start_x = None
		self.track_start_y = None

	def rotation_tracking(self, event):
		if self.track_start_x != None:
			self.track_stop_x = event.x
			self.track_stop_y = event.y
			# We set the rotation vector to go through the origin (centerx, centery), and be perpendicular to the line the user 
			# dragged with the mouse and in the X, Y plane.
			# To find a vector perpendicular to vectors a and b, take the cross product c = a x b.
			# In this case, we want a to be the a vector starting at the origin and in the direction the user dragged and b to 
			# be a unit vector in the +Z direction starting at the origin.
			# Since b is (0, 0, 1), the cross product math works out so that if a = (x, y, 0), then c = (y, -x, 0)
			deltax = self.track_stop_x - self.track_start_x
			deltay = self.track_stop_y - self.track_start_y
			length = gp.distance((0., 0.), (deltax, deltay))
			if length > 0.:
				angle = math.pi * length/min(self.centerx, self.centery)
				rotation_vector = (-deltay, deltax)
				self.rotate_molecule(rotation_vector, angle)
				self.track_start_x = self.track_stop_x
				self.track_start_y = self.track_stop_y
		else:
			self.track_start_x = event.x
			self.track_start_y = event.y

	def rotate_molecule(self, vector, theta):
		# Rotates the molecule by theta around vector
		# vector should start at the origin and be defined by a tuple that gives the X, Y displacement of the endpoint from the origin
		# Now calculate the rotation matrix
		# Rotation about an arbitrary axis: http://inside.mines.edu/~gmurray/ArbitraryAxisRotation/
		u = vector[0]
		u2 = u*u
		v = vector[1]
		v2 = v*v
		uv = u*v
		w = 0. # This simplifies things, and the w terms drop out of the rotation matrix
		L2 = u2 + v2
		L = math.sqrt(L2)
		sinth = math.sin(theta)
		costh = math.cos(theta)
		one_m_costh = 1.0 - costh
		vsinth_over_L = v * sinth/L
		usinth_over_L = u * sinth/L
		uv_one_m_costh_over_L2 = uv * one_m_costh/L2
		# We don't use the 4th row or column of the rotation matrix on that webpage because that is only for translations
		self.rotation_matrix = gp.RichList()
		self.rotation_matrix.append(gp.RichList())
		self.rotation_matrix[0].append((u2 + v2 * costh)/L2)
		self.rotation_matrix[0].append(uv_one_m_costh_over_L2)
		self.rotation_matrix[0].append(vsinth_over_L)
		self.rotation_matrix.append(gp.RichList())
		self.rotation_matrix[1].append(uv_one_m_costh_over_L2)
		self.rotation_matrix[1].append((v2 + u2*costh)/L2)
		self.rotation_matrix[1].append(-usinth_over_L)
		self.rotation_matrix.append(gp.RichList())
		self.rotation_matrix[2].append(-vsinth_over_L)
		self.rotation_matrix[2].append(usinth_over_L)
		self.rotation_matrix[2].append((u2+v2)*costh/L2)
		# Now multiply each atom's coordinates times the rotation matrix to get the new rotated coordinates
		for atom in self.molecule.atoms:
			dotprods = gp.RichList()
			for ind in xrange(3):
				dotprods.append(gp.RichList([ atom.coords[x] * self.rotation_matrix[ind][x] for x in xrange(3)]))
				dotprods[-1] = dotprods[-1].sum()
			atom.coords = [x for x in dotprods]
		# Same thing for the frequencies
		for freq in self.molecule.frequencies:
			for move in freq.displacements:
				dotprods = gp.RichList()
				for ind in xrange(3):
					dotprods.append(gp.RichList([ move[x] * self.rotation_matrix[ind][x] for x in xrange(3)]))
					dotprods[-1] = dotprods[-1].sum()
				for coord in xrange(3):
					move[coord] = dotprods[coord]
		if not self.frequency_running:
			self.plot_molecule()

	def width(self):
		return int(self.cget('width'))

	def height(self):
		return int(self.cget('height'))

	def define_center(self):
		self.centerx = self.width()/2.0
		self.centery = self.height()/2.0

	def show_new_molecule(self, molecule):
		# Loads a new molecule into the canvas and displays it
		# molecule should have an interface similar to giechem.py Molecule class object
		self.change_molecule(molecule)
		self.plot_molecule()

	def change_molecule(self, molecule):
		# Loads a new molecule into the canvas, but DOES NOT DISPLAY it - use show_new_molecule for that
		# molecule should have an interface similar to giechem.py Molecule class object
		if self.frequency_running:
			self.stop_frequency()
		self.molecule = molecule
		for atom in molecule.atoms:
			atom.ball_radius = ball_radii[atom.number]
			atom.color = element_colors[atom.number]
		self.molecule.center_it()
		self.find_largest_distance()
	
	def find_largest_distance(self):
		# Finds the atom that most deviates from the center of the molecule - used to define the closest approach to the user in the
		# perspective routines
		max_distance = 0.
		for atom in self.molecule.atoms:
			max_distance = max(max_distance, gp.distance(atom.coords, (0., 0., 0.)))
		self.molecule.horizon = min(horizon, -max_distance * 1.10)
		self.molecule.foreground = max_distance

	def erase_molecule(self):
		self.delete('molecule')

	def delete_molecule(self, molecule=False):
		# External code can use molecule to pass in a molecule object - the shown molecule will only be deleted if it is the same 
		# object as the passed in molecule object
		if not molecule or molecule == self.molecule:
			if self.frequency_running:
				self.stop_frequency()
			self.delete('molecule')

	def plot_molecule(self, vibrating=False):
		# Plots the current molecule on the canvase
		if self.molecule:
			self.erase_molecule()
			self.molecule.center_it()
			self.determine_scale()
			self.define_center()
			self.provide_perspective(vibrating=vibrating)
			self.create_plot_list()
			self.plot_items()

	def atom_center(self, atom):
		# Returns the X, Y center of the atom in canvas pixel coordinates
		# atom should be Atom class object
		centerx = self.centerx + self.pixels_per_ang * atom.perspective_coords[0]
		centery = self.centery + self.pixels_per_ang * atom.perspective_coords[1]
		return centerx, centery

	def ball_coordinates(self, atom):
		# Returns the coordinates for the sphere that represents atom
		# atom should be Atom class object
		centerx, centery = self.atom_center(atom)
		radius = atom.perspective_radius * self.pixels_per_ang
		return centerx-radius, centery-radius, centerx+radius, centery+radius

	def get_atom_max_z(self, atom):
		# atom should be Atom class object
		# Returns the most positive Z coordinate of the atom's sphere
		return atom.perspective_coords[2] + atom.perspective_radius

	def create_plot_list(self):
		# Create a list of all items that need to be plotted, sorted by most positive/least negative Z-coordinate so that items
		# appear on top of each other properly.  List items are either an object of the atom class (for atoms) or a tuple with
		# two objects of the atom class (for bonds).
		self.plot_list = gp.RichList()
		for atom in self.molecule.atoms:
			self.plot_list.append([atom])
			for ind in atom.bonds:
				if ind < atom.index:
					self.plot_list.append([atom, self.molecule.atoms[ind]])
		def z_sort(plot_item):
			if len(plot_item) == 1:
				# atom
				return self.get_atom_max_z(plot_item[0])
			else:
				# bond
				return max(plot_item[0].perspective_coords[2], plot_item[1].perspective_coords[2])
		self.plot_list.sort(key=z_sort)

	def clip_bond(self, a, b, clip_ratio):
		# Clips off the part of the bond that would be hidden by an atom sphere.  If we don't do this, the bond is visible all the
		# way to the center of the sphere instead of coming out at an angle like it should
		# a, b should be the current endpoints of the bond, and the portion clipped will come from the b end
		# returns the new coordinates of b
		rise = b[1] - a[1]
		run = b[0] - a[0]
		new_b = (a[0] + run - run*clip_ratio, a[1] + rise - rise*clip_ratio)
		return new_b

	def plot_items(self):
		# plots all the atoms & bonds for the molecule
		# use self.create_plot_list before calling this
		for item in self.plot_list:
			if len(item) == 1:
				# atom
				self.create_oval(self.ball_coordinates(item[0]), fill=item[0].color, tag='molecule')
			else:
				# bond
				centerx, centery = self.atom_center(item[0])
				partnerx, partnery = self.atom_center(item[1])
				distance = gp.distance(item[0].perspective_coords, item[1].perspective_coords)
				clip_ratio = item[1].perspective_radius/distance
				endx, endy = self.clip_bond((centerx, centery), (partnerx, partnery), clip_ratio)
				clip_ratio = item[0].perspective_radius/distance
				startx, starty = self.clip_bond((partnerx, partnery), (centerx, centery), clip_ratio)
				self.create_line(startx, starty, endx, endy, width=self.bondwidth, fill='#C0C0C0', tag='molecule')

	def determine_scale(self):
		# Figures out the correct scale so that the molecule fills most, but not all of the canvas
		xcoords = [ atom.coords[0] for atom in self.molecule.atoms ]
		xmin = min(xcoords)
		xmax = max(xcoords)
		xdelta = max(xmax - xmin, 5.) + 2.
		xpixels_per_ang = self.width()/xdelta
		ycoords = [ atom.coords[1] for atom in self.molecule.atoms ]
		ymin = min(ycoords)
		ymax = max(ycoords)
		ydelta = max(ymax - ymin, 5.) + 2.
		ypixels_per_ang = self.height()/ydelta
		self.pixels_per_ang = min(xpixels_per_ang, ypixels_per_ang)
		if self.pixels_per_ang > 25.:
			self.bondwidth = 5
		elif self.pixels_per_ang > 10.:
			self.bondwidth = 3
		else:
			self.bondwidth = 1

	def stop_frequency(self, *event):
		# Stops any current animation
		self.after_cancel(self.after_id)
		self.frequency_running = False
		self.plot_molecule()

	def animate_frequency(self, freq, frame=0, count=1):
		# freq = the frequency to animate
		# frame = the frame (from 0 (ground state) to self.frames (fully vibrated)) currently displayed
		# count = the direction (+1 or -1) to increment frame
		# frequencies should be objects of giechem.py class Frequency
		self.frequency_running = True
		# Cancel any rogue frequency animation out there - could happen if the user asks for a new animation while I'm
		# animating an old one.
		self.after_cancel(self.after_id)
		frame = frame + count
		if frame > self.frames:
			frame = self.frames - 1
			count = -1
		elif frame < 0:
			frame = 1
			count = 1
		frame_factor = float(frame) / self.frames
		scale = self.frequency_scale * frame_factor
		for move, atom in zip(freq.displacements, self.molecule.atoms):
			atom.vib_coords = [ atom.coords[x] + move[x] * scale for x in xrange(3) ]
		self.plot_molecule(vibrating=True)
		self.after_id = self.after(int(1000.0/self.frequency_speed), self.animate_frequency, freq, frame, count)

if __name__ == '__main__':
	import sys
	import readxyz as xyz
	# sys.argv[1] should be the name of an .xyz file of the format:
	# #_of_atoms
	# junk line
	# A X Y Z
	# B X Y Z
	# ...
	# C X Y Z
	#
	# where A-C are element names (or atomic numbers) and X Y Z are coordinates
	molecule = xyz.read_xyz_file(sys.argv[1])
	root = Tk.Tk()
	canvas = Tk.Canvas(root, height=600, width=850)
	canvas.pack()
	plot = MolPlot(root, molecule=molecule, height=100, width=100, borderwidth=0)
	plot.plot_molecule()
	plot.pack(pady=0, ipady=0)
	plot_window = canvas.create_window(425, 300, window=plot)
	plotmove_x = plotmove_y = None
	def move_plot(event):
		global plotmove_x, plotmove_y
		if plotmove_x == None:
			plotmove_x = event.x
			plotmove_y = event.y
		else:
			movex = event.x - plotmove_x
			movey = event.y - plotmove_y
			canvas.move(plot_window, movex, movey)
	def move_plot_end(event):
		global plotmove_x, plotmove_y
		plotmove_x = plotmove_y = None
	plot.bind("<B3-Motion>", move_plot)
	plot.bind("<ButtonRelease-3>", move_plot_end)
	root.mainloop()