#!/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'

if len(sys.argv) > 1:
    print(f"The first command-line argument is: {sys.argv[1]}")
    imos=sys.argv[1]
else:
    print("No command-line arguments provided. Use MOS1.")
    imos=1

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_MOS{imos}_gtiFilteredEvts.ds'
    if(os.path.isfile(eventfile)==False):
        print("*** run 02_sas_flares_mos.py ***")
        sys.exit()

    # EPIC_pn_gtiFilteredEvts.ds
    search_str = f'{products_dir}/{obsid}/EPIC_MOS{imos}_Image_{obsid}.fit'
    print(search_str)
    emfiles = glob.glob(search_str)
    if not (emfiles):
        print("*** run 02_sas_flares_mos.py ***")
        sys.exit()
        
    out_IMFile   = emfiles[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_EM" # Quality flag for EPIC pn
    n_pattern    = 12           # Pattern selection
    mos_pi_min    = 200.        # Low energy range eV
    mos_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+f'/EPIC_MOS{imos}_source_lightcurve.lc'   # Name of the output source light curve

    # Arguments of SAS Command
    expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_source},{y_source},{r_source}))&&(PI in [{mos_pi_min}:{mos_pi_max}])'
    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")
    # Execute the SAS task with the parameters to produce the source region light curve
    w("evselect", inargs).run()

    # Inspect light curve
    #plt.figure(figsize=(20,8))  # Size of figure
    #plotLC(plt,mos_threshold[imos-1],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_EM" # Quality flag for EPIC pn
    n_pattern    = 12           # Pattern selection

    # Define the output ligthcurve file name
    in_SPSRCFile = work_dir+f'/EPIC_MOS{imos}_source_spectrum.fits'   # Name of the output source spectrum

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

    print("   Filter expression to use: "+expression+" \n")
    w("evselect", inargs).run()

    # Extract the background region spectrum with the same criteria

    # Define the output ligthcurve file name    
    in_SPBKGFile = work_dir+f'/EPIC_MOS{imos}_background_spectrum.fits'   # Name of the output background spectrum

    # 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','withspecranges=yes','specchannelmin=0',
                  'specchannelmax=11999',f'expression={expression}']

    print("   Filter expression to use: "+expression+" \n")
    w("evselect", 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=12
    NBINS=1190
    in_RESPFile = work_dir+f'/EPIC_MOS{imos}.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+f'/EPIC_MOS{imos}.arf'   # Name of the output ancillary

    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("arfgen", inargs).run()

    # rebin the spectra and link associated files
    in_GRPFile = work_dir+f'/EPIC_MOS{imos}_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()