#!/usr/bin/env python

import argparse
import os

import numpy as np

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(
    "--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
var_elem = 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)

igen = IGMGenerator(
    args.emin,
    args.emax,
    args.nbins,
    binscale=args.binscale,
    resonant_scattering=args.resonant_scattering,
    cxb_factor=args.cxb_factor,
    var_elem=var_elem,
)
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)
