#-------------------------------------------------------------------
#Utility script to read exoplanet and parent star
#data from the California Carnegie Planet Search tables
#(http://exoplanets.org), and compute various derived
#quantities and simplified tables for plotting
#This is the original version used to make the graphs
#in Chapter 1 of the book.  Since that time, the exoplanets.org
#web site has been slightly re-organized and the table columns
#have changed a bit.  See ProcessExoplanets.py in the
#ChapterScripts directory for the version
#that works with the version of the table that can currently be
#downloaded from exoplanets.org.
#-------------------------------------------------------------------

#The input tables contain non-numeric data,
#so the readTable utility can't be used

#**ToDo:
#         *Compile this data into an exoplanets.py
#          data handbook, analogous to planets.py

import string,math
from ClimateUtilities import *

#Dictionary  of star table position of useful information
#
#     HD        HD number
#     Hipp      Hipparcos number
#     name      Star alternate name
#     BminusV   B magnitude minus V magnitude (color index)
#     Vmag      V magnitude (visual apparent magnitude)
#     dist      Distance in parsecs
#     Teff      Effective temperature
#     FeToH     Fe to H ratio (index of metallicity)
#     M         Mass in units of Solar masses
starDict = {'HD':0,'Hipp':1,'name':2,'BminusV':5,'Vmag':6,
            'dist':7,'Teff':8,'FeToH':10,'M':12}
#Dictionary  of planet table position of useful information
#
#     name      Planet name
#     period    orbital period in Earth days
#     ecc       eccentricity
#     mass      Planet maximum mass in Jupiter masses
#     rsm       semi-major axis in A.U.
#
planetDict = {'name':1,'period':2,'ecc':4,'mass':8,'rsm':9}

#Functions needed for magnitude and T calculation
def T(BV):
    logT = (14.551 - BV)/3.684
    if logT > 3.961:
        a,b,c = [.344,-3.402,8.037]
        logT = (-b - (b*b-4*a*c)**.5)/(2.*a)
    return 10.**logT

def Bolo(T):
    x = math.log10(T) - 4.
    return -8.499 * x**4 + 13.421* x**3- 8.131*x**2 - 3.901 *x - 0.438

starFile = open('StarsOrig.txt')
planetFile = open('PlanetsOrig.txt')

#Read in and process stellar data
lines = starFile.readlines()

start = False
header = True
ColorIndex = []
TList = []
IBoloList = []
starNames = []
MstarList = []
starHDIndex ={} #Dictionaries linking star ID to index
starHippIndex={}
starNameIndex={}
count = 0
for line in lines:
    if start:
        buff = line.split(',')
        if header:
            header = False
        else:
            #First get a star catalog number or name so that
            #we can make a dictionary of star data indexed by
            #catalog number/name
            HD = buff[starDict['HD']]
            Hipp = buff[starDict['Hipp']]
            starName = buff[starDict['name']]
            #Put in the dictionary entry
            if len(HD.strip())>0 and (not HD.strip() == '-1'):
                starKey = 'HD '+HD.strip()
                starHDIndex[starKey] = count
            if len(Hipp.strip())>0 and (not Hipp.strip() == '-1'):
                starKey = 'HIP '+Hipp.strip()
                starHippIndex[starKey]= count
            if len(starName.strip())>0:
                starNameIndex[starName.strip()] = count
            starNames.append(starName)
            #
            BminusV = string.atof(buff[starDict['BminusV']])
            #
            #Now look for temperature data. If it's missing,
            #recompute temperature from B-V
            try:
                Teff = string.atof(buff[starDict['Teff']])
            except:
                Teff = T(BminusV)
            TList.append(Teff)
            #Convert magnitude to absolute magnitude
            d = string.atof(buff[starDict['dist']])
            Vmag = string.atof(buff[starDict['Vmag']])
            AbsVmag = Vmag - 5.*math.log10(d/10.)
            #Add in bolometric correction
            AbsBoloMag = AbsVmag + Bolo(Teff)
            #Convert to V luminosity (relative to Sun)
            I = 10.**((4.72-AbsBoloMag)/2.5)
            IBoloList.append(I)
            MstarList.append(atof(buff[starDict['M']]))
            #
            count += 1
    elif 'EVERYTHING BELOW CONFORMS' in line:
        start = True

#Now read in the planet data
lines = planetFile.readlines()
start = False
header = True
MPlanetList = []
rsmList = []
eccList = []
yearList = []
IPlanetList = []
TstarPlanetList = []
planetNames = []
count = 0
for line in lines:
    if start:
        buff = line.split(',')
        if header:
            header = False
        else:
            star = buff[planetDict['name']].strip()[:-1]
            star = star.strip()
            try:
                ecc = atof(buff[planetDict['ecc']].split('(')[0])
            except:
                ecc = 0.
            m = atof(buff[planetDict['mass']].split('(')[0])
            rsm = atof(buff[planetDict['rsm']].split('(')[0])
            per = atof(buff[planetDict['period']].split('(')[0])
            #Find the index of the host star, so we
            #can look up information from the star table
            for index in [starHDIndex,starHippIndex,starNameIndex]:
                try:
                    starIndex = index[star]
                except:
                    continue
            #
            planetNames.append(buff[planetDict['name']].strip())
            MPlanetList.append(m)
            rsmList.append(rsm)
            eccList.append(ecc)
            yearList.append(per)
            IPlanetList.append(IBoloList[starIndex])
            TstarPlanetList.append(TList[starIndex])
            #
            count += 1
    elif 'EVERYTHING BELOW HERE CONFORMS' in line:
        start = True
IPlanet = Numeric.array(IPlanetList)
rsm = Numeric.array(rsmList)
ecc = Numeric.array(eccList)
#
axrat = (1.+ecc)/(1.-ecc)
LstarPeri = IPlanet/(rsm*rsm*(1-ecc)**2)
LstarAp = IPlanet/(rsm*rsm*(1+ecc)**2)

c = Curve()
c.addCurve(MPlanetList,'mPlanet')
c.addCurve(rsmList,'rsm')
c.addCurve(eccList,'ecc')
c.addCurve(axrat,'PeriApRat')
c.addCurve(yearList,'yearLength')
c.addCurve(IPlanetList,'Luminosity')
c.addCurve(LstarPeri,'PeriSolarConstant')
c.addCurve(LstarAp,'ApSolarConstant')
c.addCurve(TstarPlanetList,'TStar')
c.dump('planetsOriginal.temp.txt')