#!/usr/bin/env python3
#
# Copyright (c) 2017 Weitian LI <liweitianux@live.com>
# MIT license

"""
Calculate the cooling function table with respect to the specified
temperature range, using the XSPEC model 'wabs*apec' with the
provided arguments.

Later, the cooling function profile w.r.t. a temperature profile
can be quickly derived by interpolating this cooling function table.


Description
-----------
Emission measure:
    EM = \int n_e n_H dV ~= (n_e^2 / ratio_eH) V  [ cm^-3 ]
where 'ratio_eH' is the ratio of electron density to proton density (n_H).

APEC normalization returned by XSPEC is simply the *emission measure* of
the gas scaled by the distance:
    eta = (\int n_e n_H dV) / (4 pi (D_A (1+z))^2)

The flux calculated with the XSPEC `flux` command has dimension:
    Flux: [ photon s^-1 cm^-2 ] or [ erg s^-2 cm^-2 ]

If we let EM=1 and then set the APEC's normalization,
the cooling function is therefore derived by calculating the flux
using the XSPEC `flux` command, and the cooling function has
dimension of [ FLUX / EM ].
"""

import argparse
import subprocess
import sys
import os
from datetime import datetime

from _context import acispy
from acispy.cosmo import Calculator


def make_xspec_script(outfile, data):
    """
    Generate the XSPEC script for cooling function table calculation.

    Parameters
    ----------
    outfile: str
        Filename of the output XSPEC script
    data: dict
        Data used to format the template XSPEC script
    """
    data["prog_name"] = os.path.basename(sys.argv[0])
    data["date"] = datetime.now().isoformat()

    xspec_script = """\
# Calculate the cooling function table w.r.t the temperature range.
#
# Generated by: %(prog_name)s
# Date: %(date)s

# debug (off)
chatter 0

set xs_return_results 1
set xs_echo_script 0
# set tcl_precision 12

query yes
abund %(abund_table)s
dummyrsp 0.01 100.0 4096 linear
model wabs*apec & %(nh)s & 1.0 & %(abundance)s & %(redshift)s & %(norm)s & /*

# flux unit and value index
set unit "%(unit)s"
if {$unit eq "erg"} {
    set fidx 0
} else {
    set fidx 3
}

# output cooling function table
set cf_fn "%(outfile)s"
set cf_fd [open $cf_fn w]

set elow  %(elow)s
set ehigh %(ehigh)s
set tmin  %(tmin)s
set tmax  %(tmax)s
set tstep %(tstep)s

puts $cf_fd "# temperature(keV)    flux($elow-$ehigh)($unit/cm^2/s)"

# temperature sampling points
for {set t $tmin} {$t <= $tmax} {set t [expr {$t + $tstep}]} {
    newpar 2 $t
    flux $elow $ehigh
    set cf [lindex [tcloutr flux 1] $fidx]
    puts $cf_fd "$t    $cf"
}

close $cf_fd
tclexit
""" % data
    with open(outfile, "w") as f:
        f.write(xspec_script)


def main():
    parser = argparse.ArgumentParser(
        description="Calculate the cooling function table for " +
                    "specified temperature range")
    parser.add_argument("-d", "--debug", dest="debug", action="store_true",
                        help="debug; only make XSPEC script")
    parser.add_argument("-C", "--clobber", dest="clobber", action="store_true",
                        help="overwrite existing files")
    parser.add_argument("-A", "--abund-table", dest="abund_table",
                        default="grsa",
                        help="abundance table (default: grsa)")
    parser.add_argument("-L", "--elow", dest="elow",
                        type=float, default=0.7,
                        help="lower energy limit [keV] (default: 0.7 keV)")
    parser.add_argument("-H", "--ehigh", dest="ehigh",
                        type=float, default=7.0,
                        help="upper energy limit [keV] (default: 7.0 keV)")
    parser.add_argument("-t", "--tmin", dest="tmin",
                        type=float, default=0.1,
                        help="lower temperature limit [keV] (default: 0.1)")
    parser.add_argument("-T", "--tmax", dest="tmax",
                        type=float, default=15.0,
                        help="upper temperature limit [keV] (default: 15.0)")
    parser.add_argument("-s", "--tstep", dest="tstep",
                        type=float, default=0.02,
                        help="temperature step size [keV] (default: 0.02)")
    parser.add_argument("-u", "--unit", dest="unit", required=True,
                        choices=["erg", "photon"],
                        help="use flux values of unit [erg/cm^2/s] or " +
                        "[photon/cm^2/s]")
    parser.add_argument("-z", "--redshift", dest="redshift",
                        type=float, required=True,
                        help="redshift")
    parser.add_argument("-n", "--nh", dest="nh", type=float, required=True,
                        help="HI column density (unit: 1e22)")
    parser.add_argument("-Z", "--abundance", dest="abundance",
                        type=float, required=True,
                        help="average abundance (unit: solar)")
    parser.add_argument("-o", "--outfile", dest="outfile", required=True,
                        help="filename of the output cooling function table")
    args = parser.parse_args()

    xspec_script = os.path.splitext(args.outfile)[0] + ".xcm"
    if not args.clobber:
        if os.path.exists(args.outfile):
            raise OSError("Output file already exists: %s" % args.outfile)
        if os.path.exists(xspec_script):
            raise OSError("Output XSPEC script already exists: %s" %
                          xspec_script)

    cosmocalc = Calculator()
    xspec_data = vars(args)
    xspec_data["norm"] = cosmocalc.norm_apec(z=args.redshift)
    print("outfile: %s" % args.outfile, file=sys.stderr)
    print("xspec_script: %s" % xspec_script, file=sys.stderr)
    print("xspec_data:", xspec_data, sep="\n", file=sys.stderr)

    make_xspec_script(outfile=xspec_script, data=xspec_data)
    if not args.debug:
        print("Invoke XSPEC to calculate the cooling function table ...",
              file=sys.stderr)
        subprocess.check_call(["xspec", "-", xspec_script],
                              stdout=subprocess.DEVNULL)


if __name__ == "__main__":
    main()