#!python

import os
import numpy as np
import argparse
from soxs.thermal_spectra import IGMGenerator

parser = argparse.ArgumentParser(description='Create a thermal spectrum using the SOXS IGM model '
                                             'and write it to a file. The abundances of individual '
                                             'elements can be set by supplying optional arguments '
                                             'in the form of --O=0.5, --Mg=0.6, etc.')
parser.add_argument("kT", help='The temperature in keV.')
parser.add_argument("nH", help='The hydrogen number density in cm**-3.')
parser.add_argument("abund", type=float, help="The metal abundance in solar units.")
parser.add_argument("redshift", type=float, help="The redshift of the source.")
parser.add_argument("norm", type=float, 
                    help='The normalization of the model, in the standard Xspec units of '+
                         '1.0e-14*EM/(4*pi*(1+z)**2*D_A**2).')
parser.add_argument("specfile", type=str, help="The filename to write the spectrum to.")
parser.add_argument("emin", help='The minimum energy in keV.')
parser.add_argument("emax", help='The maximum energy in keV.')
parser.add_argument("nbins", type=int, help='The number of bins in the spectrum.')
parser.add_argument("--binscale", type=str, default='linear',
                    help='The scale of the energy binning: '
                         '"linear" or "log". Default: "linear"')
parser.add_argument("--resonant_scattering", action='store_true',
                    help="Whether or not to include the effects of resonant scattering "
                         "from CXB photons. Default: False")
parser.add_argument("--cxb_factor", type=float, default=0.5,
                    help="The fraction of the CXB photons that are resonant scattered "
                         "to enhance the lines. Default: 0.5")
parser.add_argument("--var_elem_option", type=int, 
                    help="An integer to choose between options for variable elements, "
                         "which are: 1: specify abundances of O, Ne, and Fe separately "
                         "from other metals, 2: specify abundances of O, Ne, Mg, Si, S, "
                         "and Fe separately from other metals. 3: specify abundances of "
                         "C, N, O, Ne, Mg, Si, S, Ca, and Fe separately from other metals. "
                         "Default: None, which means no metal abundances can be specified "
                         "separately.")
parser.add_argument("--absorb_model", type=str,
                    help='Model for applying foreground Galactic absorption.')
parser.add_argument("--nH_abs", 
                    help='The hydrogen column in units of 10**22 atoms/cm**2.')
parser.add_argument("--overwrite", action='store_true', 
                    help='Overwrite an existing file with the same name.')

args, unknown = parser.parse_known_args()

elem_abund = None

if len(unknown) > 0:
    elem_abund = {}
    for uarg in unknown:
        key, value = uarg[2:].split("=")
        elem_abund[key] = float(value)

igen = IGMGenerator(args.emin, args.emax, args.nbins, binscale=args.binscale,
                    resonant_scattering=args.resonant_scattering, cxb_factor=args.cxb_factor,
                    var_elem_option=args.var_elem_option)
spec = igen.get_spectrum(args.kT, args.nH, args.abund, args.redshift, args.norm, 
                         elem_abund=elem_abund)

if args.absorb_model is not None:
    if args.nH_abs is None:
        raise RuntimeError("Must specify a value for --nH_abs if including"
                           "foreground absorption!")
    spec.apply_foreground_absorption(args.nH_abs, model=args.absorb_model,
                                     abund_table="feld")

spec.write_file(args.specfile, overwrite=args.overwrite)