'''

        Tool for reading of element and isotope yield tables
	
	Christian Ritter 10/2013

	
	This tool contains a class which can deal
	with ascii tables. The following shows examples
	in how to extract data.

	In the following is shown how to extract data with examples.
    	(We prefer to load ipython with matplotlib and numpy support
    	via the alias
    	alias mpython='ipython -pylab -p numpy -editor emacsclient')


	mpython

        Python 2.7.3 (default, Jul 24 2012, 10:05:38) 
        Type "copyright", "credits" or "license" for more information.

        IPython 0.12 -- An enhanced Interactive Python.
        ?         -> Introduction and overview of IPython's features.
        %quickref -> Quick reference.
        help      -> Python's own help system.
        object?   -> Details about 'object', use 'object??' for extra details.

        IPython profile: numpy

        Welcome to pylab, a matplotlib-based Python environment [backend: GTKAgg].
        For more information, type 'help(pylab)'.



        ################################################################
        #for working with mesa set - e.g. from inside the set directory

        In [1]: import read_yields_table as y #Make sure script is where PYTHONPATH variable
					      # is pointing to
        In [2]: y1=y.yields('.')  #execute inside the dir containing a yield table file
				  #optional with file as filename 
        In [3]: y1.get(1.65,0.02,'Yields') #return yields for a star of mass 1.65Msun 
					   #and metallicity Z=0.02
		y1.get(1.65,0.0001,'age')  #return header quantity of certain table




'''

import matplotlib.pyplot as plt

color=['r','k','b','g']
marker_type=['o','p','s','D']
line_style=['--','-','-.',':']

class yields():

        def __init__(self,dir='element_table_winds_and_preSN.txt'):

		'''
			dir : specifing the filename of the table file

		'''


                import os
                if '/' in dir:
                        self.label=dir.split('/')[-1]
                else:
                        self.label=dir
                self.path=dir
	
		file1=open(dir)
		lines=file1.readlines()
		file1.close()
		header1=[]
		table_header=[]
		age=[]
		yield_data=[]
		kin_e=[]
		lum_bands=[]
		m_final=[]
		for line in lines:
			if 'H' in line[0]:
				if not 'Table' in line:
					header1.append(line.strip())
				else: 
					table_header.append(line.strip())
					yield_data.append([])
					lum_bands.append([])
					m_final.append([])
				if 'Lifetime' in line:
					age.append(float(line.split(':')[1]))
				if 'kinetic energy' in line:
					kin_e.append(float(line.split(':')[1]))
				if 'band' in line:
					lum_bands[-1].append(float(line.split(':')[1]))
				if 'Mfinal' in line:
					m_final[-1].append(float(line.split(':')[1]))
				continue
		
			if '&Isotopes &Yields' in line or '&Elements &Yields' in line:
				title_line=line.split('&')[1:]
				column_titles=[]
				for t in title_line:
					yield_data[-1].append([])	
					column_titles.append(t.strip())
				#print column_titles
				continue
			#iso ,name and yields
			yield_data[-1][0].append(line.split('&')[1].strip())
			yield_data[-1][1].append(float(line.split('&')[2].strip()))
			# for additional data
			datal=len(yield_data[-1])
			for t in range(2,datal-2):
				yield_data[-1][t].append(float(line.split('&')[t+1].strip()))
			#for A and Z
			for t in range(datal-2,datal):
				yield_data[-1][t].append(int(line.split('&')[t+1].strip()))
				
		self.yield_data=yield_data
		#table header points to element in yield_data
		self.table_idx={}
		i=0
		for table in table_header:
			self.table_idx[table.split(':')[1].strip()]=i
			i+=1
		#define  header
		self.header={}
		#print 'header1: ',header1
		for h in header1:
			self.header[h.split(':')[0][1:].strip()]=h.split(':')[1].strip()
		self.table_header=table_header
		#print self.table_header
		self.cols=column_titles
		#print self.header
		#print self.table_idx
		self.age=age
		self.kin_e=kin_e
		self.lum_bands=lum_bands
		self.m_final=m_final
		#print self.age
		#print m_final
	def get(self,M,Z,quantity,specie=''):

		'''
			Extracting quantity
			for star of mass M and metallicity Z.
			quantity (if available): 
				columns (string), check by using yields.cols
				  or 'age' for stellar lifetime
				  or 'kin E' for kinetic energy
				  or 'lum bands' to get the three luminosity bands
				  or 'Mfinal' to get the final stellar mass

			specie : optional, return certain specie

		'''
		
		inp='(M='+str(float(M))+',Z='+str(float(Z))+')'	
		idx=self.table_idx[inp]
		#print 'len tableidx:',len(self.table_idx)
		#print 'len age',len(self.age)
		if quantity=='age':
			data=self.age[idx]
			return data
		if quantity =='kin E':
                        data=self.kin_e[idx]
                        return data
		if quantity == 'lum bands':
                        data=self.lum_bands[idx]
                        return data
		if quantity == 'Mfinal':
			#print idx, self.m_final
			data=self.m_final[idx][0]	
			return data		
		else:
			data=self.yield_data[idx]
			if specie=='':
				idx_col=self.cols.index(quantity)
				set1=data[idx_col]
				return set1
			else:
				idx_col=self.cols.index('Yields')
				set1=data[idx_col]
				specie_all= data[0]
				for k in range(len(set1)):
					if specie in specie_all[k]:
						return set1[k]