#!/usr/bin/env python

from pysas.wrapper import Wrapper as w
import os, sys
from os.path import dirname
import inspect
import glob

import os.path
from os import path
import subprocess
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.table import Table
from matplotlib.colors import LogNorm

import re
import pyds9

import arches
from arches.utils import *
from arches.config import *

root_path=dirname(dirname(dirname(inspect.getfile(arches))))
print("Arches root path: {}".format(root_path))

archive_dir=root_path+'/data/archive'
events_dir=root_path+'/data/processed'
products_dir=root_path+'/products/sas'
ds9reg_dir=root_path+'/data/ds9reg'

create_folder(products_dir)

inargs = ['--version']
t = w('sasver', inargs)
t.run()

files = glob.glob(archive_dir+'/*')

for obsid in files:
    obsid = os.path.basename(obsid)
    if(obsid in skip):
        continue

    work_dir = init_work_dir(obsid, products_dir=products_dir)

    os.chdir(work_dir)

    eventfile=f'{products_dir}/{obsid}/EPIC_pn_gtiFilteredEvts.ds'
    if(os.path.isfile(eventfile)==False):
        print("*** run 02_filter_flares_pn.py ***")
        sys.exit()

    # EPIC_pn_gtiFilteredEvts.ds
    search_str = f'{products_dir}/{obsid}/EPIC_PN_Image_{obsid}.fit'
    print(search_str)
    epfiles = glob.glob(search_str)
    if not (epfiles):
        print("*** run 02_filter_flares_pn.py ***")
        sys.exit()
        
    out_IMFile   = epfiles[0]

    # Visualize the image with ds9
    src_fn=f'{ds9reg_dir}/{obsid}/src.reg'
    bkg_fn=f'{ds9reg_dir}/{obsid}/bkg.reg'
    x_source,y_source,r_source,x_bkg,y_bkg,r_bkg,r2_bkg = get_ds9_regions(out_IMFile, src_fn, bkg_fn)

    #########################################
    # Extract the source region light curve #
    #########################################
    
    # Define some parameters for filtering the event file and define the light curve binning

    q_flag       = "#XMMEA_EP" # Quality flag for EPIC pn
    n_pattern    = 4           # Pattern selection
    pn_pi_min    = 200.        # Low energy range eV
    pn_pi_max    = 10000.      # High energy range eV
    lc_bin       = 100         # Light curve bin in secs

    # Define the output ligthcurve file name
    in_LCSRCFile = work_dir+'/EPIC_PN_source_lightcurve.lc'   # Name of the output source light curve

    # SAS Command
    cmd        = "evselect" # SAS task to be executed                  

    # Arguments of SAS Command

    expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_source},{y_source},{r_source}))&&(PI in [{pn_pi_min}:{pn_pi_max}])'  # event filter expression
    inargs     = [f'table={eventfile}','energycolumn=PI','withrateset=yes',f'rateset={in_LCSRCFile}',
                  f'timebinsize={lc_bin}','maketimecolumn=yes','makeratecolumn=yes',f'expression={expression}']
    
    print("   Filter expression to use: "+expression+" \n")
    print("   SAS command to be executed: "+cmd+", with arguments; \n")
    # Execute the SAS task with the parameters to produce the source region light curve
    w(cmd, inargs).run()

    # Inspect light curve
    #plt.figure(figsize=(20,8))  # Size of figure
    #plotLC(plt,pn_threshold,in_LCSRCFile) # Plot source region light curve
    #plt.legend()
    #plt.show()
    
    ######################################
    # Extract the source region spectrum #
    ######################################

    # Define some parameters for filtering the event file
    q_flag       = "#XMMEA_EP" # Quality flag for EPIC pn
    n_pattern    = 4           # Pattern selection

    # Define the output ligthcurve file name
    in_SPSRCFile = work_dir+'/EPIC_PN_source_spectrum.fits'   # Name of the output source spectrum

    # SAS Command
    cmd        = "evselect" # SAS task to be executed                  

    # Arguments of SAS Command
    expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_source},{y_source},{r_source}))'
    inargs     = [f'table={eventfile}','withspectrumset=yes',f'spectrumset={in_SPSRCFile}',
                  'energycolumn=PI','spectralbinsize=5','withspecranges=yes','specchannelmin=0',
                  'specchannelmax=20479',f'expression={expression}']

    print("   Filter expression to use: "+expression+" \n")
    print("   SAS command to be executed: "+cmd+", with arguments; \n")

    w(cmd, inargs).run()

    # Extract the background region spectrum with the same criteria

    # Define the output ligthcurve file name    
    in_SPBKGFile = work_dir+'/EPIC_PN_background_spectrum.fits'   # Name of the output background spectrum

    # SAS Command
    cmd        = "evselect" # SAS task to be executed                  

    # Arguments of SAS Command
    if(r2_bkg==None):
        expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_bkg},{y_bkg},{r_bkg}))'
    else:
        expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN annulus({x_bkg},{y_bkg},{r_bkg},{r2_bkg}))'
            
    inargs     = [f'table={eventfile}','withspectrumset=yes',f'spectrumset={in_SPBKGFile}',
                  'energycolumn=PI','spectralbinsize=5','specchannelmin=0',#'withspecranges=yes',
                  'specchannelmax=20479',f'expression={expression}']

    print("   Filter expression to use: "+expression+" \n")
    print("   SAS command to be executed: "+cmd+", with arguments; \n")

    w(cmd, inargs).run()

    w("backscale", [f'spectrumset={in_SPSRCFile}',f'badpixlocation={eventfile}']).run()
    w("backscale", [f'spectrumset={in_SPBKGFile}',f'badpixlocation={eventfile}']).run()

    # https://www.cosmos.esa.int/web/xmm-newton/sas-thread-epic-merging
    MAXENERGY=15
    NBINS=1490
    in_RESPFile = work_dir+'/EPIC_PN.rmf'   # Name of the output redistribution
    w("rmfgen", [f'spectrumset={in_SPSRCFile}',f'rmfset={in_RESPFile}',
                 'withenergybins=yes', 'energymin=0.1',
                 f'energymax={MAXENERGY}', f'nenergybins={NBINS}',]).run()

    in_ARFFile = work_dir+'/EPIC_PN.arf'   # Name of the output ancillary
    cmd        = "arfgen" # SAS task to be executed                  

    print("   Checking for Response File ..... \n")
    # Check if RESP file is available.
    if os.path.isfile(in_RESPFile):
        print ("File "+in_RESPFile+" exists. \n")
    else:
        print ("File "+in_RESPFile+" does not exist, please check. \n")
        sys.exit()

    # Arguments of SAS Command
    inargs     = [f'spectrumset={in_SPSRCFile}',f'arfset={in_ARFFile}',
                  'withrmfset=yes',f'rmfset={in_RESPFile}',f'badpixlocation={eventfile}','detmaptype=psf']
    w(cmd, inargs).run()

    # rebin the spectra and link associated files
    in_GRPFile = work_dir+'/EPIC_PN_spectrum_grp.fits'   # Name of the output specgruop
    inargs     = [f'spectrumset={in_SPSRCFile}','mincounts=30','oversample=3',
                  f'rmfset={in_RESPFile}',f'arfset={in_ARFFile}',
                  f'backgndset={in_SPBKGFile}',f'groupedset={in_GRPFile}']
    w("specgroup", inargs).run()
