#!/usr/bin/env python
from xspec import *

from loguru import logger as log
import matplotlib.pyplot as plt
import argparse
import os, os.path
import sys

# import numpy as np
# import astropy.units as u

inittemplate = """
# This filename should be the python load script output by nicerl3-spect 
exec(open("merged_cutmpu7_loadRPP.py").read())

# Here you choose the base model and initialize the parameters
# To freeze parameters, either init with the string "val,0" (including the quotes)
# or use m1(2).frozen = True after the setPars.
m1 = Model("tbabs* (bb + pow)")
m1.setPars(0.01, 0.044, 1.0, 1.9, 0.01)

# Define the parameters that can be frozen if they peg at 0.0
# These should be STRINGS that can be eval()'ed to give the parameter object
FREEZEPARAMS = [
    "m1(1)",  # nH
    "nxb(4)",  # trel_norm
]

# Now define the parameters that can be thawed if they are significant.
# These should be STRINGS that can be eval()'ed to give the parameter object
THAWPARAMS = [
        "sky(12)", # swcxovii_norm - OVII normalization 
]

"""

parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    description="""Do RPP catalog spectral analysis.

    First run with --init specinit.py to create a template called specinit.py;
    then edit specinit.py to configure spectral analysis;
    then run RPP-spec.py specinit.py to run the analysis with that configuration.""",
)
parser.add_argument(
    "initfile", help="Python init script (template generated by --init)"
)
parser.add_argument(
    "--init",
    action="store_true",
    default=False,
    help="Write template initialization script to file",
)
args = parser.parse_args()

if args.init:
    if os.path.exists(args.initfile):
        log.error(f"File {args.initfile} already exists, not overwriting.")
        raise FileExistsError
    with open(args.initfile, "w") as outf:
        print(inittemplate, file=outf)
    sys.exit(0)

Plot.splashPage = False
exec(open(args.initfile).read())
AllData.show()

sky = AllModels(1, modName="sky")
nxb = AllModels(1, modName="nxb")

# Thaw params in SCORPEON, as requested
for pname in THAWPARAMS:
    p = eval(pname)
    p.frozen = False

Fit.statMethod = "pgstat"
Fit.nIterations = 200
Fit.query = "no"
Fit.perform()

# Now check if parameters are pegged at zero and freeze them
# Each time fit is redone.
# This needs to be more sophisticated
for pname in FREEZEPARAMS:
    p = eval(pname)
    if p.sigma == -1:
        # Freeze norm at 0.0 and refit
        p.values = "0.0,0"
        Fit.perform()

AllModels.calcFlux("0.3 2.0 err")
fluxes = AllData(1).flux
# Eflux, Eflux,min, Eflux,max, Phflux, Phflux,min, Phflux,max, ... for each of the 3 sources
print(fluxes)

# Don't generate XSPEC plots, just use this to get the plot data out
Plot.device = "/null"
# Plot.device = "/xs" # For debugging
Plot.xAxis = "keV"
Plot.add = True
Plot.xLog = True
Plot.yLog = True
Plot.setRebin()
Plot.addCommand("res y 0.01 20")
Plot.addCommand("wenv output")
Plot("ldata ratio")
# Plot.delCommand(1)
en = Plot.x()
rates = Plot.y()
rates_err = Plot.yErr()
folded = Plot.model()


fig, ax = plt.subplots(1, 1)
ax.errorbar(en, rates, yerr=rates_err, label="Data")
ax.plot(en, folded, label="Full Model")
ax.plot(en, m1.folded(1), label="Source")
ax.set_xscale("log")
ax.set_yscale("log")
ax.legend()

plt.show()


# Now need to output several things:
# * XSPEC model file
# * YAML output file with all the useful parameters like band fluxes, spectral
#   parameters with errors, etc. Whatever is needed for Tables or further
#   analysis
# * Nice, fully labeled plot
