#!/usr/bin/env python

import os
import numpy as np
import argparse
from soxs import ApecGenerator

parser = argparse.ArgumentParser(description='Create a thermal spectrum 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("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("--velocity", default=0.0,
                    help='The velocity broadening parameter, in units of km/s. Default: 0.0')
parser.add_argument("--apec_vers", type=str,
                    help='The version of the AtomDB tables to use. Default is'
                         'to use the version currently included with this'
                         'version of SOXS.')
parser.add_argument("--absorb_model", type=str,
                    help='Model for applying foreground Galactic absorption.')
parser.add_argument("--nh", default=0.02, 
                    help='The hydrogen column in units of 10**22 atoms/cm**2. Default: 0.02')
parser.add_argument("--overwrite", action='store_true', 
                    help='Overwrite an existing file with the same name.')
parser.add_argument("--nolines", action='store_true',
                    help="Make a spectrum without lines.")
parser.add_argument("--abund_table", type=str, default="angr",
                    help="The abundance table to be used for solar abundances. "
                         "Either a string corresponding to a built-in table or "
                         "an ASCII file containing a column of 30 floats "
                         "corresponding to the abundances of each element relative "
                         "to the abundance of H. Default is set in the SOXS "
                         "configuration file, the default for which is 'angr'.")
feature_parser = parser.add_mutually_exclusive_group(required=False)
feature_parser.add_argument("--broadening", dest='broadening', action='store_true',
                            help='Turn thermal and velocity broadening on. On by default.')
feature_parser.add_argument("--no_broadening", dest='broadening', action='store_false',
                            help='Turn thermal and velocity broadening off. On by default.')
parser.set_defaults(broadening=True)

args, unknown = parser.parse_known_args()

var_elem = None
elem_abund = None

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

if os.path.exists(args.abund_table):
    abund_table = np.loadtxt(args.abund_table)
else:
    abund_table = args.abund_table
agen = ApecGenerator(args.emin, args.emax, args.nbins, apec_vers=args.apec_vers, 
                     broadening=args.broadening, nolines=args.nolines, 
                     var_elem=var_elem, abund_table=abund_table)
spec = agen.get_spectrum(args.kT, args.abund, args.redshift, args.norm, 
                         velocity=args.velocity, elem_abund=elem_abund)

if args.absorb_model is not None:
    spec.apply_foreground_absorption(args.nh, model=args.absorb_model)

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