import giepy as gp

element_names = ['XX', 'H','HE','LI','BE','B','C','N','O','F','NE','NA','MG','AL','SI','P','S','CL','AR',
		'K','CA','SC','TI','V','CR','MN','FE','CO','NI','CU','ZN','GA','GE','AS','SE','BR','KR',
		'RB','SR','Y','ZR','NB','MO','TC','RU','RH','PD','AG','CD','IN','SN','SB','TE','I','XE',
		'CS','BA','LA','CE','PR','ND','PM','SM','EU','GD','TB','DY',
		'HO','ER','TM','YB','LU','HF','TA','W','RE','OS','IR','PT',
		'AU','HG','TL','PB','BI','PO','AT','RN','FR','RA','AC','TH','PA','U','NP','PU','AM','CM','BK','CF','XX',
		'FM','MD','CB']

# Table row gives which row of the periodic table the atom number appears in.  All elements
table_row = [0] + [1]*2 + [2]*8 + [3]*8 + [4]*18 + [5]*18 + [6]*32 + [7]*32

# Gives the anticipated maximum distance between two atoms that is considered to be a bond based on the row of the periodic table the two 
# atoms are from.  bond distance = bonding_distances[RowA][RowB].  Lists are padded to account for the fact that there is no row 0 on the table
bonding_distances = [[], 
		[0, 0.9, 1.3, 1.7, 1.9, 2.0, 2.0, 2.0], 
		[0, 1.3, 1.85, 2.2, 2.4, 2.4, 2.4, 2.4],
		[0, 1.7, 2.2, 2.7, 2.9, 2.9, 2.9, 2.9],
		[0, 1.9, 2.4, 2.9, 2.9, 2.9, 2.9, 2.9],
		[0, 2.0, 2.4, 2.9, 2.9, 2.9, 2.9, 2.9],
		[0, 2.0, 2.4, 2.9, 2.9, 2.9, 2.9, 2.9],
		[0, 2.0, 2.4, 2.9, 2.9, 2.9, 2.9, 2.9]]

def is_bond(a, b, dist):
	# a and b should be the atomic numbers of the two atoms involved in the possible bond
	# dist should be the distance between a and b
	# Returns True if a bond exists, False if not
	row_a = table_row[a]
	row_b = table_row[b]
	return dist <= bonding_distances[row_a][row_b]

class Atom(object):
	#
	# Stores the data for a single atom
	#
	def __init__(self, name="", number=0, coords=None, index=0):
		# name = should be the element name
		# number = atomic number
		# coords = tuple or list of x, y, z coordinates
		# index = atom number (place of this atom in list of atoms
		gp.params_to_attributes(self, ['name', 'number', 'index'], locals())
		self.name = self.name.upper()
		if self.name and not self.number:
			self.number = self.get_number(self.name)
		elif self.number and not self.name:
			self.name = self.get_name(self.number)
		self.coords = gp.RichList()
		if coords:
			self.new_coords(coords)
		# self.bonds = list of atoms this atom is bonded to
		self.bonds = gp.RichList()

	def add_bond(self, partner):
		# Adds a bond to atom number "partner"
		self.bonds.append(partner)

	def new_coords(self, coords):
		# Change the XYZ coordinates of this atom
		# coords = non-string iterable of length 3 (X,Y,Z)
		self.coords = gp.RichList([ float(x) for x in coords ])
		if len(self.coords) != 3:
			self.coords = gp.RichList()

	def get_name(self, number):
		# Gets the element name for atomic number 'number'
		try:
			self.name = element_names[number]
		except IndexError:
			self.name = ""

	def get_number(self, name):
		# Gets the atomic number for element 'name'
		try:
			self.number = element_names.index(name.upper())
		except ValueError:
			self.number = 0

	def get_name_or_number(self, value):
		# Figures out whether this is an element name or atomic number and if so, finds the corresponding atomic number or element name
		try:
			self.number = int(value)
			self.name = self.get_name(self.number)
			return True
		except ValueError:
			# This isn't an integer, try to treat it as an element name
			try:
				self.get_number(value)
				if self.number:
					self.get_name(self.number)
					return True
			except ValueError:
				return False

class Molecule(object):
	#
	# Stores the data for a molecule
	# 	- currently only lists of atoms and calculated frequencies
	#
	def __init__(self, name="", atoms=None, frequencies=None):
		#
		# name = name of the molecule
		# atoms = optional list of Atom objects
		# frequencies = optional list of Frequency objects
		#
		gp.params_to_attributes(self, ['name'], locals())
		# self.atoms = list of class Atom objects
		if atoms:
			self.atoms = gp.RichList(atoms)
		else:
			self.atoms = gp.RichList()
		self.bonds = gp.RichDict()
		if frequencies:
			self.frequencies = gp.RichList(frequencies)
		else:
			self.frequencies = gp.RichList()

	def add_atom(self, id=None, coords=None):
		#
		# Adds the atom to the molecule and finds any bonds to that atom
		#
		# id = atomic number or element name of the atom
		# coords = XYZ coordinates of the atom
		#
		atom_number = self.num_atoms()
		self.atoms.append(Atom(coords=coords, index=atom_number))
		# Find the atomic number and element name
		self.atoms[atom_number].get_name_or_number(id)
		# Find bonds
		for ind in xrange(atom_number):
			if self.find_bonds(ind, atom_number):
				self.atoms[ind].add_bond(atom_number)
				self.atoms[atom_number].add_bond(ind)

	def add_frequency(self, frequency):
		# Frequency should by a Frequency class object or have a similar interface
		self.frequencies.append(frequency)

	def distance(self, a, b):
		# Returns the distance between atom numbers a and b
		return gp.distance(self.atoms[a].coords, self.atoms[b].coords)

	def find_bonds(self, a, b):
		# Calculates whether atom numbers a and b are bonded based on the distance between them
		# Returns True if a bond exists, False if not
		dist = self.distance(a, b)
		return is_bond(self.atoms[a].number, self.atoms[b].number, dist)

	def are_bonded(self, a, b):
		# Returns True if a bond exists between atoms a and b - pulls the data from already calculated lists of bonds
		return b in self.atoms[a].bonds

	def num_atoms(self):
		return len(self.atoms)

	def center_it(self):
		# Places the "center of mass" of the molecule at the origin.  "center of mass" is calculated as if each atom has a weight of 1
		for axis in xrange(3):
			vals = [ atom.coords[axis] for atom in self.atoms ]
			my_min = min(vals)
			my_max = max(vals)
			ave = (my_min + my_max)/2.0
			for atom in self.atoms:
				atom.coords[axis] = atom.coords[axis] - ave

class Frequency(object):
	#
	# class to contain the data for a calculated frequency
	#
	def __init__(self, spectrum, symmetry='?', frequency=0., raman=0., ir=0., mass=0., force=0., depolar=0., molplot=None, freq_type=0, molecule=None):
		#
		# spectrum = parent VibrationalSpectrum class object
		# symmetry = symmetry of the vibration
		# frequency = wavenumber of the vibration
		# raman = Raman intensity of the vibration
		# ir = IR intensity of the vibration
		# mass = reduced mass of the vibration
		# force = force constant of the vibration
		# depolar = depolarization ratio of the vibration
		# molplot = MolPlot instance for animation of this frequency
		# freq_type = 0 for IR, 1 for Raman
		# molecule = should be object with similar interface to class Molecule
		#
		# self.displacements = list of XYZ displacements for each atom in the molecule
		#
		gp.params_to_attributes(self, ['spectrum', 'symmetry', 'frequency', 'raman', 'ir', 'mass', 'force', 'depolar', 'molplot', 'molecule'], locals())
		self.displacements = gp.RichList()
		if freq_type:
			self.intensity = self.raman
		else:
			self.intensity = self.ir
		
	def add_data(self, type, value, freq_type):
		# Add some type of calculated frequency data other than coordinates
		# This is the data in the header of each calculated frequency.
		# type = First whitespace delimited string on this line of the Gaussian output file
		# value = value for this frequency
		# freq_type = 0 if IR, 1 if Raman intensities/activities are desired
		val = float(value)
		if type == 'Red.':
			self.mass = val
		elif type == 'Frc':
			self.force = val
		elif type == 'IR':
			self.ir = val
			if not freq_type:
				self.intensity = self.ir
		elif type == 'Raman':
			self.raman = val
			if freq_type:
				self.intensity = self.raman
		elif type == 'Depolar':
			self.depolar = val

	def add_displacements(self, vals):
		# Add the displacements (XYZ displacement) for an atom - need to be added in the same order as the Atoms in the molecule
		self.displacements.append([ float(x) for x in vals])

	def calc_peak(self, freq_type, low, high, bandwidth):
		# calculate the contribution of this frequency to the overall vibrational spectrum by using a Lorentzian curve to broaden the
		# calculated frequency.
		# freq_type = 0 if IR, 1 if Raman intensities/activities are desired
		# low, high = highest and lowest wavenumbers desired
		# bandwidth = bandwidth to use in the Lorentzian to broaden the frequency line
		if freq_type:
			self.intensity = self.raman
		else:
			self.intensity = self.ir
		scaled_band_over_pi = self.intensity*bandwidth/3.14159265
		band_squared = bandwidth * bandwidth
		intensities = [ scaled_band_over_pi / ((wavenumber - self.frequency)**2 + band_squared) for wavenumber in xrange(low, high+1)]
		self.peak_height = max(intensities)
		return intensities

	def animate_me(self, event):
		# Animates this frequency on a MolPlot instance
		if self.spectrum.molecule and self.molplot:
			self.molplot.change_molecule(self.spectrum.molecule)
			self.molplot.animate_frequency(self)