#!/usr/bin/env python

# Program to dump out spex mekal model at various temperature grid points
# The lines and continuum are then gathered up, and stored in an apec-format
# table model for xspec

# Jeremy Sanders 2005-2008

import os
import string

import numpy as N
import pyfits

#####################################################################
# Adjustable parameters

# executable to use to run spex
spexexecutable = 'spex_linux'

# use a lot of temperatures or a more normal set
large = True

if large:
    # temperatures for which to dump the data in the output file
    temperature_list = ( list(N.arange(0.05, 0.5, 0.01)) +
                         list(N.arange(0.5, 1.5, 0.02)) +
                         list(N.arange(1.5, 2.5, 0.05)) +
                         list(N.arange(2.5, 5.0, 0.1)) +
                         list(N.arange(5.0, 10.0, 0.2)) +
                         list(N.arange(10., 20., 1)) +
                         list(N.arange(20., 68., 2)) )
    # output root (using spex version) for apec format filenames
    # creates outroot_(line|coco).fits
    outroot = 'spex_2.00.11_large'
else:
    temperature_list = [0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6,
                        0.75, 0.8, 0.95, 1.0, 1.05, 1.1, 1.15, 1.2, 1.25, 1.3, 1.35,
                        1.4, 1.45, 1.5, 1.55, 1.6, 1.65, 1.7, 1.75, 1.8, 1.85, 1.9, 1.95,
                        2.0, 2.1, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.2,
                        4.4, 4.6, 4.8, 5.0, 5.4, 5.8, 6.2, 6.6, 7.0, 7.4, 7.8, 8.2, 8.6, 9.0,
                        9.4, 9.8, 10.2, 10.6, 11., 11.4, 12.0, 13.0, 14.0, 15.0, 16.0,
                        18., 20., 22., 24., 26., 30., 34., 38., 42., 46., 50.,
                        54., 58., 62., 66.]
    outroot = 'spex_2.00.11'

# energy range and stepping to sample continuua (log spacing used)
contminenergy = 0.05
contmaxenergy = 15.
contenergysteps = 300

###########################################################

# conversions
keV_K = 11.6048e6
keV_erg = 1.6022e-9

# convert from unit norm to cm3 in spex
norm_factor_cm3 = 1e58

# identify element names with element numbers in spex
elements = ('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')

# create dict to convert element names into numbers
element_nums = {}
for num, element in enumerate(elements):
    element_nums[element] = num+1

# these are the apec elements
apec_elements = ['He', 'C', 'N', 'O', 'Ne', 'Mg', 'Al', 'Si', 'S', 'Ar',
                 'Ca', 'Fe', 'Ni']

# roman numerals (bah)
roman_numerals = ('I', 'II', 'III', 'IV', 'V', 'VI', 'VII',
                  'VIII', 'IX', 'X', 'XI', 'XII', 'XIII', 'XIV',
                  'XV', 'XVI', 'XVII', 'XVIII', 'XIX', 'XX',
                  'XXI', 'XXII', 'XXIII', 'XXIV', 'XXV', 'XXVI', 'XXVII',
                  'XXVIII', 'XXIX', 'XXX')

# dict to convert numerals to numbers
roman_to_number = {}
for num, numeral in enumerate(roman_numerals):
    roman_to_number[numeral] = num+1

class SpexCommandStream:

    # abundance to use to get continuua of elements more exactly
    continuum_mult = 10000.

    temperatures = N.array(temperature_list)
    
    def __init__(self):
        """Builds up a command stream to send to spex, to generate
        the dumped output."""

        self.outfilename = 'dumped_spex_commands'
        self.file = open(self.outfilename, 'w')
        print >>self.file, 'egrid log %e %e %i' % (contminenergy, contmaxenergy,
                                                   contenergysteps)
        print >>self.file, 'var calc new'
        print >>self.file, 'var newmekal all true'
        print >>self.file, 'comp cie'

        # switch off all elements
        self.setElements(elements[1:], 0.)

    def setElements(self, elements, val):
        """Set elements (list) to abundance value."""
        for e in elements:
            print >>self.file, 'par %02i val %f' % (element_nums[e], val)

    def generateLines(self, T):
        """Generate lines for model with temperature given."""

        # set temperature
        print >>self.file, 'par t val %e' % T

        # switch on all apec elements
        self.setElements(apec_elements, 1.)

        # compute model
        print >>self.file, 'calc'

        # dump out lines
        outfile = 'dumped_lines_T%f' % T
        print >>self.file, 'ascdump file %s 1 1 line' % outfile

        # switch off all apec elements
        self.setElements(apec_elements, 0.)

    def generateContinuua(self, T):
        """Generate continuua for each element (H is always included)."""

        # set temperature
        print >>self.file, 'par t val %e' % T

        print >>self.file, 'calc'
        outfile = 'dumped_continuum_T%f_%s' % (T, 'H')
        print >>self.file, 'ascdump file %s 1 1 tcl' % outfile

        for el in apec_elements:
            self.setElements( (el,), self.continuum_mult )
            print >>self.file, 'calc'
            outfile = 'dumped_continuum_T%f_%s' % (T, el)
            print >>self.file, 'ascdump file %s 1 1 tcl' % outfile
            self.setElements( (el,), 0. )

    def generateScript(self):
        """Get spex to simulate the spectra"""

        for T in self.temperatures:
            self.generateLines(T)
            self.generateContinuua(T)

    def runScript(self):
        """Get spex to run generated script."""
        
        print >>self.file, 'quit'
        self.file.close()

        print "Starting spex"
        cmd = '%s < %s' % (spexexecutable, self.outfilename)
        os.system(cmd)

        os.unlink(self.outfilename)

    def interpretAllLines(self):
        """Interpret spex dumped spectra."""

        # construct line file
        hdus = []
        Nelement = []
        Nline = []
        for T in self.temperatures:
            hdu, numlines, numelements = self.interpretDumpedLines(T)
            Nline.append(numlines)
            Nelement.append(numelements)
            hdus.append(hdu)

        # construct hdu describing parameters
        col_kT = pyfits.Column(name='kT', format='1E', unit='keV',
                               array=self.temperatures)
        col_EDensity = pyfits.Column(name='EDensity', format='1E', unit='cm**-3',
                                     array=N.zeros(len(self.temperatures))+1)
        col_Nelement = pyfits.Column(name='Nelement', format='1J',
                                     array=Nelement)
        col_Nline = pyfits.Column(name='Nline', format='1J',
                                  array=Nline)
        tabhdu = pyfits.new_table([col_kT, col_EDensity, col_Nelement, col_Nline])
        tabhdu.name = 'PARAMETERS'

        # make output file containing all lines
        hdulist = pyfits.HDUList([pyfits.PrimaryHDU(), tabhdu] + hdus)
        hdulist.writeto(outroot + '_line.fits')

    def interpretDumpedLines(self, T):
        """Interpret dumped lines file.

        Returns new HDU, number of lines, and number of elements
        """

        totflux = 0.
        elements = {}
        lines = []
        outfile = 'dumped_lines_T%f.asc' % T
        for line in open(outfile):
            p = line.strip().split()
            # numerical line
            if p[0][0] in string.digits:
                # horribly, have to hard code in column numbers here
                element = element_nums[line[9:11].strip()]
                ion = roman_to_number[ line[12:17].strip() ]
                wavelength = float( line[88:94] )
                energy_keV = float( line[71:82] )

                # convert from total photon flux to normalised photon flux
                strength = float( line[103:114] ) / norm_factor_cm3

                # keep track of total flux in ergs
                totflux += energy_keV*keV_erg*strength

                lines.append( (element, ion, wavelength, strength) )
                elements[element] = True
        os.unlink(outfile)

        # sort lines by element and ion
        lines.sort()

        # construct up FITS table to APEC format
        col_lambda = pyfits.Column(name='Lambda', format='1E', unit='A',
                                   array=[i[2] for i in lines])
        col_lambda_err = pyfits.Column(name='Lambda_Err', format='1E', unit='A',
                                       array=N.zeros( (len(lines),) ) + N.nan)
        col_epsilon = pyfits.Column(name='Epsilon', format='1E',
                                    unit='photons cm^3 s^-1',
                                    array=[i[3] for i in lines])
        col_epsilon_err = pyfits.Column(name='Epsilon_Err', format='1E',
                                        unit='photons cm^3 s^-1',
                                        array=N.zeros((len(lines),)) + N.nan)
        col_element = pyfits.Column(name='Element', format='1J',
                                    array=[i[0] for i in lines])
        col_ion = pyfits.Column(name='Ion', format='1J',
                                array=[i[1] for i in lines])
        
        col_upperlev = pyfits.Column(name='UpperLev', format='1J',
                                     array=N.zeros((len(lines),)) + 2)
        col_lowerlev = pyfits.Column(name='LowerLev', format='1J',
                                     array=N.zeros((len(lines),)) + 1)

        tabhdu = pyfits.new_table([col_lambda, col_lambda_err, col_epsilon,
                                   col_epsilon_err, col_element, col_ion,
                                   col_upperlev, col_lowerlev])
        tabhdu.name = 'EMISSIVITY'
        h = tabhdu.header
        h.update('HIERARCH TEMPERATURE', T*keV_K, 'Electron temperature')
        h.update('XTEMP', T, 'e temperature (keV)')
       
        # fixme wrong below (erg not photon)
        h.update('TOT_LINE', totflux,
                 'Total Line Emission (erg cm^3 s^-1)')
        h.update('N_LINES', len(lines), 'Number of emission lines')
        return tabhdu, len(lines), len(elements)

    def _readContinuum(self, filename):
        """Take spex dumped model spectrum file, and extract continuum."""
        outenergy = []
        outval = []
        for line in open(filename):
            p = line.strip().split()
            if p[0][0] in string.digits:
                outenergy.append(float(p[1]))
                outval.append(float(p[2]) * 1e44 / norm_factor_cm3)
        return (N.array(outenergy), N.array(outval))

    def interpretDumpedContinuum(self, T):
        """Interpret dumped continuum file."""

        # read in continum from each file
        continuua = {}

        allelements = ['H']+apec_elements
        for element in allelements:
            filename = 'dumped_continuum_T%f_%s.asc' % (T, element)
            energy, vals = self._readContinuum(filename)
            continuua[element] = vals
            os.unlink(filename)

        # subtract H continuum from each element except hydrogen
        # also divide by abundance these were generated at
        for element in apec_elements:
            continuua[element] -= continuua['H']
            continuua[element] /= self.continuum_mult

        # construct table
        col_element = pyfits.Column(name='Z', format='1J',
                                    array=[element_nums[i] for i in allelements])
        col_rmJ = pyfits.Column(name='rmJ', format='1J', array=N.zeros(len(allelements)))
        col_N_Cont = pyfits.Column(name='N_Cont', format='1J',
                                   array= [len(energy)]*len(allelements))
        col_E_Cont = pyfits.Column(name='E_Cont', format='%iE' % len(energy),
                                   unit='keV', array=N.resize(energy, (len(allelements),
                                                                       len(energy))))
        col_Continuum = pyfits.Column(name='Continuum', format='%iE' % len(energy),
                                      unit='photons cm^3 s^-1 keV^-1',
                                      array=[continuua[i] for i in allelements])
        col_Cont_Err = pyfits.Column(name='Cont_Err', format='%iE' % len(energy),
                                      unit='photons cm^3 s^-1 keV^-1',
                                      array=N.zeros( (len(allelements), len(energy)) ))

        # we make all the pseudo continuua zero
        col_N_Pseudo = pyfits.Column(name='N_Pseudo', format='1J',
                                     array=N.zeros(len(allelements)) + 1)

        # can't get below to work, unless I set array size to 2 columns
        # doesn't seem to matter, though
        col_E_Pseudo = pyfits.Column(name='E_Pseudo', format='2E',
                                     array=N.zeros( (len(allelements), 2) ))
        col_Pseudo = pyfits.Column(name='Pseudo', format='2E',
                                   array=N.zeros( (len(allelements), 2) ))
        col_Pseudo_Err = pyfits.Column(name='Pseudo_Err', format='2E',
                                       array=N.zeros( (len(allelements), 2) ))
        
        tabhdu = pyfits.new_table([col_element, col_rmJ,
                                   col_N_Cont, col_E_Cont,
                                   col_Continuum, col_Cont_Err,                               
                                   col_N_Pseudo, col_E_Pseudo,
                                   col_Pseudo, col_Pseudo_Err])
        
        tabhdu.name = 'EMISSIVITY'
        h = tabhdu.header
        h.update('HIERARCH TEMPERATURE', T*keV_K, 'Electron temperature')
        h.update('XTEMP', T, 'e temperature (keV)')
        h.update('DENSITY', 1., 'Electron density')
        # sum flux
        totcoco = 0.
        for i in allelements:
            totcoco += continuua[i].sum()
        h.update('TOT_COCO', totcoco, 'Total Emission (erg cm^3 s^-1)')

        return (tabhdu, len(allelements), len(energy)*len(allelements), 0)
        
    def interpretAllContinuum(self):
        """Build up continuum output file."""
        
        hdus = []
        NElement = []
        NCont = []
        NPseudo = []
        for T in self.temperatures:
            hdu, numelem, numcont, numpseudo = self.interpretDumpedContinuum(T)
            hdus.append(hdu)
            NElement.append(numelem)
            NCont.append(numcont)
            NPseudo.append(numpseudo)

        # construct hdu describing parameters
        col_kT = pyfits.Column(name='kT', format='1E', unit='keV',
                               array=self.temperatures)
        col_EDensity = pyfits.Column(name='EDensity', format='1E', unit='cm**-3',
                                     array=N.zeros(len(self.temperatures))+1)
        col_NElement = pyfits.Column(name='NElement', format='1J',
                                     array=NElement)
        col_NCont = pyfits.Column(name='NCont', format='1J',
                                  array=NCont)
        col_NPseudo = pyfits.Column(name='NPseudo', format='1J',
                                    array=NPseudo)
        tabhdu = pyfits.new_table([col_kT, col_EDensity, col_NElement, col_NCont, col_NPseudo])
        tabhdu.name = 'PARAMETERS'

        # make output file containing all lines
        hdulist = pyfits.HDUList([pyfits.PrimaryHDU(), tabhdu] + hdus)
        hdulist.writeto(outroot + '_coco.fits')
     

c = SpexCommandStream()
c.generateScript()
c.runScript()
c.interpretAllLines()
c.interpretAllContinuum()