#!/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 not (obsid in good):
        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
    #

    attfile = get_first_file(events_dir+f'/{obsid}/*{obsid}_AttHk.ds')

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

        # first element contains all energies
        label=['',]
        pimin=[200,]
        pimax=[12000,]
        for idx,e1 in enumerate(emin):
            label.append(f"_{idx}")
            pimin.append(emin[idx])
            pimax.append(emax[idx])
            
        run_evqpb(key, work_dir=work_dir, evtfile=evtfile, attfile=attfile, label=label, pimin=pimin, pimax=pimax)
        
        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()
