import giepy as gp
import gwidgets as gw
import gwidgets_wx as gwx
import sys
import Tkinter as Tk
import tkFileDialog as TkFD
import tkMessageBox as TkMB
import os.path
import tkSimpleDialog as TkSD
import molshow as ms
import giechem as gc

root = Tk.Tk()

# Defaults and globals
# directories to load and save files
load_directory = '.'
save_directory = '.'
# bandwidth to broaden calculated frequencies
bandwidth = 10.
# Each member of this list is a vibrational_spectrum
good_vibrations = gp.RichList()
# The lowest and highest wavenumber in the calculated spectra
low = 1
high = 3600
# maximum intensity of any spectral curve
max_all_spectra = 1.0
# program icon
icon = 'freqs.ico'

# frequency_type = 0 is IR, 1 is Raman
frequency_type = gw.iVar(0)
frequency_lines = gw.iVar(1)

class VibrationalSpectrum(object):
	#
	# Class to handle all the work required to plot a vibrational spectrum
	#
	def __init__(self, filename):
		#
		# filename = name of the Gaussian output file
		#
		# self.name = name of the spectrum
		# self.frequencies = list of calculated frequencies
		# self.spectrum = wavelength vs intensity data
		# self.molecule = Molecule instance
		#
		# Use read_in_file to read in a Gaussian output file
		#
		self.name = os.path.splitext(os.path.basename(filename))[0]
		self.frequencies = gp.RichList()
		self.spectrum = gp.RichList()
		self.molecule = False
		self.read_in_file(filename)
		self.compute_spectrum()

	def read_in_file(self, filename):
		# filename = full or relative path to the Gaussian output file
		input = gp.RichFile(filename, 'r')
		# Read in the geometry in standard orientation
		input_lines = input.next_split_line(splitter="")
		try:
			testmol = self.read_geometry(input_lines, filename)
			while testmol:
				self.molecule = testmol
				testmol = self.read_geometry(input_lines, filename)
		except StopIteration:
			pass
		# Now read in the frequencies
		input.seek(0)
		input_lines = input.next_split_line(splitter="")
		try:
			new_freqs = self.read_next_frequency(input_lines)
			while new_freqs:
				self.frequencies.extend(new_freqs)
				new_freqs = self.read_next_frequency(input_lines)
		except StopIteration:
			pass
		input.close()
		if self.molecule:
			for freq in self.frequencies:
				self.molecule.add_frequency(freq)
			molplot.show_new_molecule(self.molecule)
		
	def read_geometry(self, lines, filename):
		# Reads the geometry out of a gaussian output file
		# lines = iterator that returns the next line as a series of whitespace separated tokens
		# filename = name of the Gaussian output file - used to name the molecule
		name = os.path.splitext(os.path.basename(filename))[0]
		molecule = gc.Molecule(name=name)
		line = lines.next()
		while len(line) != 2 or (line[0] != 'Standard' or line[1] != 'orientation:'):
			line = lines.next()
		# If we get here, we found a standard orientation block.  Read out the useless header lines:
		for x in xrange(4):
			line = lines.next()
		# Now start reading in the atom lines until we hit a line of all '----'
		line = lines.next()
		while len(line) == 6:
			# We want to use the Gaussian Z axis as our X axis, and vice versa
			#  - older formats (G94?) have a different format that will break this
			try:
				molecule.add_atom(id=line[1], coords=line[5:2:-1])
			except IndexError:
				TkMB.showerror('File Read Error', 'Error reading in molecular geometry -\nExpected 6 columns in Standard Orientation specification')
			line = lines.next()
		return molecule

	def read_next_frequency(self, lines):
		# Reads the next frequency out of a gaussian output file
		# lines = iterator that returns the next line as a series of whitespace separated tokens
		def newline(line):
			old_line = line
			line = lines.next()
			return line, old_line
		my_freqs = gp.RichList()	
		line = lines.next()
		# Searching for a block that looks like:
		#                     A2                     B1                     B2    <--Symmetries
 		# Frequencies --   472.2660               472.2663               697.1287
		freq_type = frequency_type.get()
		while len(line) < 2 or (line[0] != 'Frequencies' or line[1] != '--'):
			if len(line) == 3 and line[1] == 'Thermochemistry':
				return False
			line, old_line = newline(line)
		for token in line[2:]:
			my_freqs.append(gc.Frequency(self, symmetry=old_line.pop(0), frequency=float(token), freq_type=freq_type, molecule=self.molecule, molplot=molplot))
		line = lines.next()
		while line[0] != 'Atom':
			# Lines of the format: Red. masses --     2.7034                 2.7034                 6.0818
			# Some lines have a space in the tag name (such as "Red. masses") and some don't (such as (Depolar")
			if line[1] == '--':
				place = 2
			else:
				place = 3
			for ind, freq in my_freqs.enumerate():
				freq.add_data(line[0], line[place+ind], freq_type)
			line = lines.next()
		line = lines.next()
		while len(line) == 2 + 3*len(my_freqs):
			# Lines of the format:  Atom AN      X      Y      Z        X      Y      Z        X      Y      Z
   			#			1   6     0.20   0.00   0.00     0.11   0.00   0.00     0.00   0.22  -0.23
			for ind, freq in my_freqs.enumerate():
				freq.add_displacements(line[4+ind*3:1+ind*3:-1])
			line = lines.next()
		return my_freqs

	def save_spectrum(self, filename):
		# Saves the XY data to a tab-delimited text file named filename
		output = gp.RichFile(filename, 'w')
		for x, y in zip(range(low, high+1), self.spectrum):
			output.write_n(gp.joiner('\t', x, y))
		output.close()

	def compute_spectrum(self):
		# Compute a vibrational spectrum by broadening each calculated frequency
		self.spectrum = [0]*(high - low + 1)
		freq_type = frequency_type.get()
		for freq in self.frequencies:
			self.spectrum = [ a + b for a, b in zip(self.spectrum, freq.calc_peak(freq_type, low, high, bandwidth)) ]

	def display_me(self, *event):
		if self.molecule:
			molplot.show_new_molecule(self.molecule)

	def plot_spectrum(self):
		# Adds the spectrum to the plot
		plot.add_line_points(zip(range(low, high+1), self.spectrum), name=self.name)
		if frequency_lines.get():
			self.plot_frequency_lines()
		plot.tag_bind(self.name, "<Button-3>", self.display_me)

	def plot_frequency_lines(self, line_color=None):
		# Put a vertical line at each frequency
		# line_color = color of the line to plot - defaults to the same color as the spectrum line
		if not line_color:
			line_color = plot.itemcget(self.name, 'fill')
		for freq in self.frequencies:
			if freq.intensity:
				intensity = freq.peak_height*0.67
			else:
				# Plot 0 intensity frequencies upside down
				intensity = -max_all_spectra/15.0
			name = self.name + str(freq.frequency)
			plot.add_line_points([(freq.frequency, 0), (freq.frequency, intensity)], line_color=line_color, name=name, legend=False, line_width=2)
			plot.tag_bind(name, "<Button-3>", freq.animate_me)

	def recolor_frequency_lines(self, line_color=None):
		# Resets all the frequency lines to be the same color as the spectrum curve, or to be line_color if specified
		self.remove_frequency_lines()
		self.plot_frequency_lines(line_color=line_color)

	def remove_frequency_lines(self):
		# Deletes all the frequency lines from the plot
		for freq in self.frequencies:
			name = self.name + str(freq.frequency)
			plot.tag_unbind(name, "<Button-3>")
			plot.delete(name)

	def max_intensity(self):
		# Returns the highest intensity for any frequency
		return max([ freq.intensity for freq in self.frequencies ])

	def max_spectrum(self):
		# Returns the highest intensity of the vibrational spectrum
		return max(self.spectrum)

	def remove_self(self):
		# Removes all the frequency lines and the vibrational spectrum from the plot, and removes self from the list of spectra
		good_vibrations.remove(self)
		try:
			plot.remove_set(self.name)
			for freq in self.frequencies:
				plot.remove_set(self.name+str(freq.frequency))
		except KeyError:
			pass
		# Delete this molecule if it is shown
		molplot.delete_molecule(self.molecule)
		replot_all_spectra()

def open_file():
	# Get the name of the file to open, then open it and plot the spectrum
	global load_directory, save_directory
	load_name = TkFD.askopenfilename(filetypes=[('G03', 'g03out'), ('G98', 'g98out'), ('All files', '*')], initialdir=load_directory, title='Open File')
	if load_name:
		load_directory = os.path.split(load_name)[0]
		if save_directory == '.':
			save_directory = load_directory
		good_vibrations.append(VibrationalSpectrum(load_name))
		replot_all_spectra()

def recompute_all_spectra(replot=True):
	for vib in good_vibrations:
		vib.compute_spectrum()
	if replot:
		replot_all_spectra()

def replot_all_spectra():
	global max_all_spectra
	# max_intensity = maximum intensity (or activity) of a calculated frequency
	try:
		max_intensity = max([ vib.max_intensity() for vib in good_vibrations ])
	except ValueError:
		max_intensity = 1.
	# max_all_spectra = maximum height of any peak in the calculated spectrum curve
	try:
		max_all_spectra = max([ vib.max_spectrum() for vib in good_vibrations ])
	except ValueError:
		max_all_spectra = 1.
	try:
		height_scale = max_intensity/max_all_spectra
	except:
		height_scale = 1.
	plot.reset()
	if frequency_type.get():
		plot.change_y_axis(ystart=0., yend=max(max_all_spectra*1.05, 1.0))
	else:
		plot.change_y_axis(ystart=max(max_all_spectra*1.05, 1.0), yend=-max_all_spectra/14.)
	for vib in good_vibrations:
		vib.plot_spectrum()

def save_spectrum():
	# Write out the xy data to a tab-delimited file
	if len(good_vibrations) > 0:
		choice = choose_spectrum()
		if choice < 0:
			return
		save_name = TkFD.asksaveasfilename(defaultextension='txt', filetypes=[('Tab Delimited Text', 'txt'), ('All files', '*')], initialdir=save_directory, title='Save Spectrum')
		if save_name:
			good_vibrations[choice].save_spectrum(save_name)

def choose_spectrum():
	# Put up a popup window that allows the user to pick which spectrum to do something to - return the index of that spectrum in 
	# good_vibrations
	choice = gw.iVar(0)
	if len(good_vibrations) > 1:
		choice.set(-1)
		def grab_choice():
			choice.set(rbuttons.get())
		spec_chooser = gw.ModalWindow(root, title='Choose Spectrum', user_accept=grab_choice, buttons=True, coords='mouse', icon=icon)
		rbuttons = spec_chooser.append(gw.EasyRadiobuttons(spec_chooser, [vib.name for vib in good_vibrations]))
		spec_chooser.pack_and_display()
	return choice.get()

def remove_spectrum():
	# Revmoes a spectrum curve, the frequency lines and the molecule (everything)
	if len(good_vibrations) > 0:
		choice = choose_spectrum()
		if choice < 0:
			return
		good_vibrations[choice].remove_self()

def get_bandwidth():
	global bandwidth
	newband = TkSD.askfloat('New Bandwidth', 'Enter new bandwidth in cm-1', initialvalue=bandwidth, parent=root, minvalue=0.01, maxvalue=1000.)
	if newband:
		bandwidth = newband
		recompute_all_spectra()

def get_freq_type():
	def grab_choice():
		frequency_type.set(rbuttons.get())
	old_freq = frequency_type.get()
	type_chooser = gw.ModalWindow(root, title='Spectrum Type', user_accept=grab_choice, buttons=True, coords='mouse', icon=icon)
	rbuttons = type_chooser.append(gw.EasyRadiobuttons(type_chooser, ['IR', 'Raman'], init=frequency_type.get()))
	type_chooser.pack_and_display()
	if frequency_type.get() != old_freq:
		recompute_all_spectra()
		# frequency_type = 0 is IR, 1 is Raman
		if frequency_type.get():
			plot.ytitle = 'Activity'
		else:
			plot.ytitle = 'Intensity'
		plot.axis_titles()

def copy_plot():
	gwx.copy_to_clipboard(widget=plot, offset1=-2, offset2=-2)

def get_vib_amp():
	def set_amp():
		molplot.frequency_scale = float(vib_scale.get())
	vib_win = gw.ModalWindow(root, title='Vibrational Amplitude', buttons=True, user_accept=set_amp, coords='mouse', icon=icon)
	vib_scale = vib_win.append(gw.EasyScale(vib_win, from_=1, to=20, length=200, orient=Tk.HORIZONTAL, value=min(molplot.frequency_scale, 10)))
	vib_win.pack_and_display()

def get_vib_speed():
	def set_speed():
		molplot.frequency_speed = float(vib_scale.get())
	vib_win = gw.ModalWindow(root, title='Vibrational Speed', buttons=True, user_accept=set_speed, coords='mouse', icon=icon)
	vib_scale = vib_win.append(gw.EasyScale(vib_win, from_=1, to=100, length=300, orient=Tk.HORIZONTAL, value=min(molplot.frequency_speed, 100)))
	vib_win.pack_and_display()

def get_mol_size():
	def set_size():
		plot_height = height.get()
		plot_width = width.get()
		molplot.configure(height=plot_height, width=plot_width)
		molplot.plot_molecule()
	vib_win = gw.ModalWindow(root, title='Molecule Size', buttons=True, user_accept=set_size, coords='mouse', icon=icon)
	height = vib_win.append(gw.SuperEntry(vib_win, label_text='Height (pixels)', vartype='int', varvalue=int(molplot.cget('height'))))
	width = vib_win.append(gw.SuperEntry(vib_win, label_text='Width (pixels)', vartype='int', varvalue=int(molplot.cget('width'))))
	vib_win.pack_and_display()

def toggle_frequency_lines():
	# Turn the plotting of the vertical frequency lines on or off
	if frequency_lines.get():
		for vib in good_vibrations:
			vib.plot_frequency_lines()
	else:
		for vib in good_vibrations:
			vib.remove_frequency_lines()
frequency_lines = gw.iVar(1)

def initiate_text():
	plot.initiate_user_created_text()

def give_help():
	readme = gp.RichFile('Readme.txt')
	info = gp.joiner('', readme.readlines())
	window = gw.HelpWindow(parent=root, title='Visualizer Help', icon=icon, info=info)

def show_about():
	info = """Frequency_Visualizer version 1.0

Written by David Giesen
Eastman Kodak Company
david.giesen@kodak.com

Using Python with Pmw and PIL extensions

This program and the code found in the files

color.py
Frequency_Visualizer.py
giechem.py
giepy.py
gwidgets.py
gwidgets_wx.py
molshow.py
readxyz.py

are placed in the public domain with no 
restrictions on their use or modification.

All code originating from the Python, PMW,
and PIL packages retains its original 
licensing.
"""
	window = gw.HelpWindow(parent=root, title='About Frequency_Visualizer', icon=icon, info=info, height=12, width=40)

menubar = gw.EasyMenuBar(parent=root, order=['File', 'Edit', 'Preferences', 'Help'], 
		menus={'File': [('Open Gaussian Output', open_file),
				('Save XY Data', save_spectrum),
				('Remove Spectrum', remove_spectrum),
				('Quit', root.quit)],
			'Edit': [('Copy Plot', copy_plot),
				 ('Add Text', initiate_text)],
			'Preferences': [('Bandwidth', get_bandwidth),
				('Frequency Type', get_freq_type),
				('Vibration Amplitude', get_vib_amp),
				('Vibration Speed', get_vib_speed),
				('Molecule Size', get_mol_size),
				('Frequency Lines', toggle_frequency_lines, 'checkbutton', frequency_lines, 1)],
			'Help': [('Help', give_help),
				 ('About', show_about)]
				})

plot = gw.GenericPlot(root, height=500, width=800, xstart=high, xend=low-1, xtitle='wavenumber', xmajor=19, ytitle='Intensity', usey=True, ydigits=1)
molplot = ms.MolPlot(root, height=225, width=225, highlightthickness=1, stopsign=True)
molplot.plot_molecule()
mol_window = gw.EasyCanvasObject(Tk.Canvas.create_window, plot, 425, 300, window=molplot, moveable=True, button=3)
plot.pack()

root.title('Frequency Visualizer')
if sys.platform == 'win32':
	root.iconbitmap(icon)

root.mainloop()