#!/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'
ds9reg_dir=root_path+'/data/ds9reg'

imos=2

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)

    work_dir = init_work_dir(obsid)

    os.chdir(work_dir)

    eventfile=f'{products_dir}/{obsid}/EPIC_MOS{imos}_gtiFilteredEvts.ds'
    if(os.path.isfile(eventfile)==False):
        print("*** run 02_filter_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_filter_flares_mos.py ***")
        sys.exit()
        
    out_IMFile   = emfiles[0]

    # Visualize the image with ds9

    d = pyds9.DS9()
    d.set("file "+out_IMFile)
    d.set('cmap bb')
    d.set('scale log')
    #d.set(f"region {ds9reg_dir}/arches.reg")
    src_fn=f'{ds9reg_dir}/{obsid}/src.reg'
    if(os.path.isfile(src_fn)==True):
        d.set(f"region {src_fn}")
    bkg_fn=f'{ds9reg_dir}/{obsid}/bkg.reg'
    if(os.path.isfile(bkg_fn)==True):
        d.set(f"region {bkg_fn}")
        
    reply = input(f"{obsid}: Proceed to make spectrum/lightcurve? [y/[n]] ") 
    
    if reply!='y':
        continue
        
    print(d.get("regions"))
    
    # Extract the relevant information from the ds9 regions.
    
    region1=(re.split("circle|annulus",d.get("regions").partition("physical")[2]))[1].replace('(','').replace(')','').replace('#',',')
    region2=(re.split("circle|annulus",d.get("regions").partition("physical")[2]))[2].replace('(','').replace(')','').replace('#',',')
    print("Identified first region: ", region1)
    print("Identified second region: ", region2)
    
    #print("region1 "+region1.partition("color")[2].replace('=','').replace('\n',''))
    #print("region2 "+region2.partition("color")[2].replace('=','').replace('\n',''))
    
    # Identify source and background regions using the 'white' color.
    c1=region1.partition("color")[2].replace('=','').replace('\n','')
    print("Region1 color: "+c1)
    if(c1=='white'):
        regionsrc = region1
        regionbkg = region2
    else:
        regionsrc = region2
        regionbkg = region1
        
    # Save and print selected region coordinates.
    x_source = regionsrc.split(",")[0].replace('\n','')

    y_source = regionsrc.split(",")[1].replace('\n','')
    r_source = regionsrc.split(",")[2].replace('\n','')
    print("The coordinates of the selected source region are: \n")
    print("  x_source = ", x_source, "(physical)")
    print("  y_source = ", y_source, "(physical)")
    print("  r_source = ", r_source, "(physical) \n")

    x_bkg = regionbkg.split(",")[0].replace('\n','')
    y_bkg = regionbkg.split(",")[1].replace('\n','')
    r_bkg = regionbkg.split(",")[2].replace('\n','')
    print("The coordinates of the selected background region are: \n")
    print("  x_bkg = ", x_bkg, "(physical)")
    print("  y_bkg = ", y_bkg, "(physical)")
    print("  r_bkg = ", r_bkg, "(physical) \n")

    # If the background is an annulus, save and print R2.

    if "annulus" in str(d.get("regions")):
        r2_bkg = regionbkg.split(",")[3].replace('\n','')
        print("  r2_bkg = ", r2_bkg, "(physical)")

    #########################################
    # 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

    # 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 [{mos_pi_min}:{mos_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,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

    # 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}))'  # 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")
    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+f'/EPIC_MOS{imos}_background_spectrum.fits'   # Name of the output background 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_bkg},{y_bkg},{r_bkg}))'  # event filter expression
    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")
    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()

    in_RESPFile = work_dir+f'/EPIC_MOS{imos}.rmf'   # Name of the output redistribution
    w("rmfgen", [f'spectrumset={in_SPSRCFile}',f'rmfset={in_RESPFile}']).run()

    in_ARFFile = work_dir+f'/EPIC_MOS{imos}.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+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()