#!python

import os
import numpy as np
import argparse
from soxs.thermal_spectra import CIEGenerator, MekalGenerator, \
    CloudyCIEGenerator

parser = argparse.ArgumentParser(description='Create a thermal CIE 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  Only available for 'apec' and 'spex' models.")
parser.add_argument("--model_vers", type=str,
                    help='The version of the CIE tables to use. Default is '
                         'to use the version currently included with this '
                         'version of SOXS.')
parser.add_argument("--binscale", type=str, default='linear',
                    help='The scale of the energy binning: '
                         '"linear" or "log". Default: "linear"')
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.')
parser.add_argument("--nolines", action='store_true',
                    help="Make a spectrum without lines. Only available for 'apec' and "
                         "'spex' models.")
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'.")
parser.add_argument("--var_elem_option", type=int,
                    help="Only for the 'cloudy' model. 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("--model", type=str, default="apec",
                    help="The CIE model to use when generating the spectrum, either "
                         "'apec', 'spex', 'mekal', or 'cloudy'. Default: 'apec'")
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. "
                                 "Only available for 'apec' and 'spex' models.")
feature_parser.add_argument("--no_broadening", dest='broadening', action='store_false',
                            help="Turn thermal and velocity broadening off. On by default. "
                                 "Only available for 'apec' and 'spex' models.")
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

if args.model in ['apec', 'spex']:
    cgen = CIEGenerator(args.model, args.emin, args.emax, args.nbins, binscale=args.binscale,
                        model_vers=args.model_vers, broadening=args.broadening, 
                        nolines=args.nolines, var_elem=var_elem, abund_table=abund_table)
    spec = cgen.get_spectrum(args.kT, args.abund, args.redshift, args.norm, 
                             velocity=args.velocity, elem_abund=elem_abund)
elif args.model == 'mekal':
    mgen = MekalGenerator(args.emin, args.emax, args.nbins, binscale=args.binscale,
                          var_elem=var_elem, abund_table=abund_table)
    spec = mgen.get_spectrum(args.kT, args.abund, args.redshift, args.norm,
                             elem_abund=elem_abund)
elif args.model == "cloudy":
    cgen = CloudyCIEGenerator(args.emin, args.emax, args.nbins, binscale=args.binscale,
                              var_elem_option=args.var_elem_option)
    spec = cgen.get_spectrum(args.kT, args.abund, args.redshift, args.norm,
                             elem_abund=elem_abund)
else:
    raise ValueError(f"Unknown model '{args.model}'!")

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!")
    if args.model == "cloudy":
        atable = "feld"
    else:
        atable = abund_table
    spec.apply_foreground_absorption(args.nH_abs, model=args.absorb_model,
                                     abund_table=atable)

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