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

    #
    # Create QPB
    # https://www.cosmos.esa.int/web/xmm-newton/sas-thread-background
    #
    
    attfiles = glob.glob(events_dir+f'/{obsid}/*{obsid}_AttHk.ds')
    if(attfiles):
        print(f"*** {attfiles} ***")
    else:
        print("AttHk?")
        sys.exit()

    attfile = attfiles[0]
    if not os.path.isfile(attfile):
        print ("File "+attfile+" does not exist, please check. \n")
        sys.exit()

    spec=[]
    bkg=[]
    resp=[]
    arf=[]
    #for imos in [1,2]:
    for key in ['MOS1','MOS2','PN']:
        evtfile = get_first_file(events_dir+f'/{obsid}/*{obsid}_E{key}*_ImagingEvts.ds')

        # Run the SAS task evqpb over each one of these event files.
        evqpb_outset=work_dir+f'/EPIC_{key}_QPB.fits'   # Name of the output file
        inargs     = [f'table={evtfile}', f'attfile={attfile}',
                      f'outset={evqpb_outset}',
                      f'exposurefactor=2.0',]
        w("evqpb", inargs).run()

        gti=f'EPIC_{key}_gti.fit'
        expression=f'gti({gti},TIME)'
        # filter the EPIC event files to create cleaned and filtered for
        # flaring particle background event files for your observation
        sci_clean=work_dir+f'/EPIC_{key}_filtered.fits'   # Name of the output file
        inargs     = [f'table={evtfile}', 'withfilteredset=yes',
                      f'filteredset={sci_clean}',
                      'destruct=yes','keepfilteroutput=true',
                      f'expression={expression}']
        w("evselect", inargs).run()

        # Use the same expression to filter the FWC event files:
        qpb_clean=work_dir+f'/EPIC_{key}_QPB_clean.fits'   # Name of the output file
        inargs     = [f'table={evqpb_outset}', 'withfilteredset=yes',
                      f'filteredset={qpb_clean}',
                      'destruct=yes','keepfilteroutput=true',
                      f'expression={expression}']
        w("evselect", inargs).run()
        
        if(key=='PN'):
            expression='(#XMMEA_EP)&&(PATTERN<=4)&&(PI>=200)&&(PI<=12000)'
        else:
            expression='(#XMMEA_EM)&&(PATTERN<=12)&&(PI>=200)&&(PI<=12000)'

        sci_image=work_dir+f'/EPIC_{key}_sci_image.fits' # Name of the output file
        # Extract an image for the science exposure
        inargs     = [f'table={sci_clean}', f'expression={expression}',
                      'imagebinning=binSize', f'imageset={sci_image}',
                      'withimageset=yes','xcolumn=X', 'ycolumn=Y',
                      'ximagebinsize=80', 'yimagebinsize=80',
                      ]
        w("evselect", inargs).run()

        qpb_image=work_dir+f'/EPIC_{key}_qpb_image.fits' # Name of the output file
        # Extract an image  for the FWC exposure
        inargs     = [f'table={qpb_clean}', f'expression={expression}',
                      'imagebinning=binSize', f'imageset={qpb_image}',
                      'withimageset=yes','xcolumn=X', 'ycolumn=Y',
                      'ximagebinsize=80', 'yimagebinsize=80','zcolumn=EWEIGHT',
                      ]
        w("evselect", inargs).run()

        # Subtract background
        cor_image=work_dir+f'/EPIC_{key}_cor_image.fits' # Name of the output file
        cmd     = ['farith', f'{sci_image}', f'{qpb_image}', f'{cor_image}', 'SUB',
                   'copyprime=yes',
                   'clobber=yes']
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)

        continue
        
        #
        # How to use the FWC event files for spectra generation
        #
        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(sci_image, src_fn, bkg_fn)

        ######################################
        # Extract the source region spectrum #
        ######################################
        
        # Define some parameters for filtering the event file
        q_flag       = "#XMMEA_EP" if (key=='PN') else "#XMMEA_EM" # Quality flag for EPIC pn
        n_pattern    = 4  if (key=='PN') else 12         # Pattern selection
        
        # for the source region of the science exposure
        sci_spec_file = work_dir+f'/EPIC_{key}_source_spectrum.fits' # Name of the output source spectrum
        expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_source},{y_source},{r_source}))'
        inargs     = [f'table={sci_clean}','withspectrumset=yes',f'spectrumset={sci_spec_file}',
                      'energycolumn=PI','spectralbinsize=5','withspecranges=yes','specchannelmin=0',
                      'specchannelmax=11999',f'expression={expression}']
        w("evselect", inargs).run()
        
        # for the background region of the science exposure
        sci_bkg_file = work_dir+f'/EPIC_{key}_background_spectrum.fits'
        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={sci_clean}','withspectrumset=yes',f'spectrumset={sci_bkg_file}',
                      'energycolumn=PI','spectralbinsize=5','withspecranges=yes','specchannelmin=0',
                      'specchannelmax=11999',f'expression={expression}']
        w("evselect", inargs).run()

        # for the source region of the science exposure
        qpb_spec_file = work_dir+f'/EPIC_{key}_qpb_source_spectrum.fits' # Name of the output source spectrum
        expression = f'{q_flag}&&(PATTERN<={n_pattern})&&((X,Y) IN circle({x_source},{y_source},{r_source}))'
        inargs     = [f'table={qpb_clean}','withspectrumset=yes',f'spectrumset={qpb_spec_file}',
                      'energycolumn=PI','spectralbinsize=5','withspecranges=yes','specchannelmin=0',
                      'specchannelmax=11999',f'expression={expression}', 'zcolumn=EWEIGHT',]
        w("evselect", inargs).run()

        # for the source region of the science exposure
        qpb_bkg_file = work_dir+f'/EPIC_{key}_qpb_background_spectrum.fits' # Name of the output source spectrum
        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={qpb_clean}','withspectrumset=yes',f'spectrumset={qpb_bkg_file}',
                      'energycolumn=PI','spectralbinsize=5','withspecranges=yes','specchannelmin=0',
                      'specchannelmax=11999',f'expression={expression}', 'zcolumn=EWEIGHT',]
        w("evselect", inargs).run()
        """
        By using the zcolumn parameter in the generation of the spectra for the FWC exposure,
        we ensure that the produced FWC spectra are correctly scaled for the ratio of exposure times
        between science and FWC exposures. By doing so, the FWC spectra are ready to be
        subtracted from the science spectra. 
        """
        #sys.exit()