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

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

parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    description="""Do RPP catalog spectral analysis from a full src+bkg model saved from 
    xspec as an *.xcm file.

    First use Xspec to fit your src+bck model and use 'save all mymodel.xcm'.
    then run RPP-spec-from-xcm.py mymodel.xcm.""",
)
parser.add_argument(
    "xcmfile", help="xcm file  init script (template generated by --init)"
)
parser.add_argument(
    "psrname", help="PSR name for output figures and yaml"
)
parser.add_argument("--outroot", help="rootname of output files",type=str, default=None)
parser.add_argument("--ylim", help="tuple for y-limits. Ex: --ylim 1e-6 1e-2",
                    nargs="+",type=float, default=None)

args = parser.parse_args()

if args.outroot is not None:
    figfile = "{}.png".format(args.outroot)
    ymlfile = "{}.yml".format(args.outroot)
else:
    figfile = "{}.png".format(os.path.splitext(args.xcmfile)[0])
    ymlfile = "{}.yml".format(os.path.splitext(args.xcmfile)[0])

Plot.splashPage = False
XspecSettings.restore(args.xcmfile, args.xcmfile)
AllData.show()

# Plot setup
Plot.device = "/null"
Plot.xAxis = 'keV'
Plot.add = True
Plot.xLog = True
Plot.yLog = True

# Re-fit (it should find the same fit as in the XCM file
Fit.statMethod = "pgstat"
Fit.nIterations = 1000
Fit.query = "no"
Fit.perform()

# TODO: Check for pegged parameters
# TODO: Check for "bad" nullhyp
# TODO: Check for standards...

# Get the fluxes
# Eflux, Eflux,min, Eflux,max, Phflux, Phflux,min, Phflux,max, ... for each of the 3 sources
AllModels.calcFlux("0.3 2.0 err")
fluxes = AllData(1).flux
fluxarray = [fluxes[0], fluxes[1], fluxes[2]]

# Get all three source models independently
src = AllModels(1)
sky = AllModels(1, modName="sky")
nxb = AllModels(1, modName="nxb")

# Fill dictionnary with parameters values and errors
param_errors = dict()
for i in range(1,src.nParameters+1):
    if not src(i).frozen:
        Fit.error("{}".format(i))
        param_errors[src(i).name] = [src(i).values[0],
                                     src(i).error[0],
                                     src(i).error[1],
                                    ]

Plot("ldata")
en = Plot.x()
en_err = Plot.xErr()
rates = Plot.y()
rates_err = Plot.yErr()
folded = Plot.model()
Plot('ratio')
en = Plot.x()
en_err = Plot.xErr()
ratio = Plot.y()
ratio_err = Plot.yErr()

fig, ax = plt.subplots(2, 1, height_ratios=[2,1],sharex='all')
fig.subplots_adjust(hspace=0.0)

# Top panel
ax[0].errorbar(en, rates, xerr=en_err, yerr=rates_err, label="Data", alpha=1.0,linestyle='')
ax[0].step(en, folded, label="Full Model")
ax[0].step(en, sky.folded(1),  label="Sky", alpha=0.3)
ax[0].step(en, nxb.folded(1), label="NXB", alpha=0.3)
ax[0].step(en, src.folded(1), label="PSR")
ax[0].set_xscale("log")
ax[0].set_yscale("log")
ax[0].set_ylabel("Photon flux\n(counts/s/keV)")
ax[0].legend()
if args.ylim is not None:
    ax[0].set_ylim(tuple(args.ylim))

# Residuals panel
ax[1].set_xlim(0.22,15.0)
ax[1].set_ylim(0.8,1.2)
ax[1].errorbar(en, ratio, xerr=en_err, yerr=ratio_err, alpha=0.5,linestyle='')
ax[1].hlines(y=1.0, xmin=0.1, xmax=15.0, alpha=0.7, linewidth=0.5, color='k')
ax[1].set_xlabel("Energy (keV)")
ax[1].set_ylabel("Ratio (data/model)")
ax[1].tick_params('x', top=True, which='both', labeltop=False)
plt.savefig(figfile)


outdict = {}
objname = args.psrname
outdict[objname] = {}
outdict[objname]["name"] = objname
outdict[objname]["psr model"] = src.expression
outdict[objname]["parameter (bestfit, -90%, +90%)"] = param_errors
outdict[objname]["fit pgstat"] = Fit.statistic
outdict[objname]["fit NHP (chi2)"] = Fit.nullhyp
outdict[objname]["flux (central, -90%, +90%)"] = fluxarray

with open(ymlfile, "w") as outfile:
    yaml.dump(outdict, outfile, default_flow_style=False)


# Now need to output several things:
# * 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
