#!/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 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+'/0862*')

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)

    search_str = f'{events_dir}/{obsid}/????_{obsid}_EPN_S???_ImagingEvts.ds'
    print(search_str)
    epfiles = glob.glob(search_str)
    if not (epfiles):
        print("*** run 01_init_events.py ***")
        #w('epproc', []).run()    
    
    eventfile = epfiles[0]
    print("Checking for EPIC-pn Event Files ..... \n")
    # Check if epproc has already run.
    if os.path.isfile(eventfile):
        print ("File "+eventfile+" exists. \n")
    else:
        print ("File "+eventfile+" does not exist, please check. \n")

    #############################################################
    # For display purposes only define the following event cuts #
    #############################################################
    pn_pattern   = 4        # pattern selection
    pn_pi_min    = 300.     # Low energy range eV
    pn_pi_max    = 12000.   # High energy range eV
    pn_flag      = 0        # FLAG

    plt.figure(figsize=(20,8))
    pl=1

    hdu_list = fits.open(eventfile, memmap=True)
    evt_data = Table(hdu_list[1].data)

    mask = ((evt_data['PATTERN'] <= pn_pattern) &
            (evt_data['FLAG'] == pn_flag) &
            (evt_data['PI'] >= pn_pi_min) &
            (evt_data['PI'] <= pn_pi_max))

    print("Events in event file" + " " + eventfile + ": " + str(len(evt_data)) + "\n")
    print("Events in filtered event file" + " " + eventfile + ": " + str(np.sum(mask)) + "\n")
    print("   Filter: PATTERN <= " + str(pn_pattern) +
          " : " + str(pn_pi_min) + " <= E(eV) <= " + str(pn_pi_max) +
          " : " + " FLAG == " + str(pn_flag)+ "\n")


    xmax=np.amax(evt_data['X'])
    xmin=np.amin(evt_data['X'])
    xmid=(xmax-xmin)/2.+xmin
    ymax=np.amax(evt_data['Y'])
    ymin=np.amin(evt_data['Y'])
    xbin_size=80
    ybin_size=80
    NBINS = (int((xmax-xmin)/xbin_size),int((ymax-ymin)/ybin_size))

    plt.subplot(1, 2, pl)
    
    img_zero_mpl = plt.hist2d(evt_data['X'], evt_data['Y'], NBINS, cmap='GnBu', norm=LogNorm())

    cbar = plt.colorbar(ticks=[10.,100.,1000.])
    cbar.ax.set_yticklabels(['10','100','1000'])

    plt.title(obsid)
    plt.xlabel('x')
    plt.ylabel('y')

    pl=pl+1

    # Create Filtered Events image
    
    xmax=np.amax(evt_data['X'][mask])
    xmin=np.amin(evt_data['X'][mask])
    xmid=(xmax-xmin)/2.+xmin
    ymax=np.amax(evt_data['Y'][mask])
    ymin=np.amin(evt_data['Y'][mask])
    xbin_size=80
    ybin_size=80
    NBINS = (int((xmax-xmin)/xbin_size),int((ymax-ymin)/ybin_size))

    plt.subplot(1, 2, pl)
    
    img_zero_mpl = plt.hist2d(evt_data['X'][mask], evt_data['Y'][mask], NBINS, cmap='GnBu', norm=LogNorm())

    cbar = plt.colorbar(ticks=[10.,100.,1000.])
    cbar.ax.set_yticklabels(['10','100','1000'])

    plt.title(obsid)
    plt.xlabel('x')
    plt.ylabel('y')

    txt=("PATTERN <= " + str(pn_pattern) +
         " : " + str(pn_pi_min) + " <= E(eV) <= " + str(pn_pi_max) +
         " : " + " FLAG == " + str(pn_flag))
    plt.text(xmid, ymin+0.1*(ymax-ymin), txt, ha='center')

    pl=pl+1
    
    hdu_list.close()
    plt.show()

    ##################################################################
    # Define a SAS filter expression to derive a background rate cut #
    ##################################################################
    pn_pattern   = 0        # pattern selection
    pn_pi_min    = 10000.   # Low energy range eV
    pn_pi_max    = 12000.   # High energy range eV
    pn_threshold = 0.75      # cts/sec (only used here for display purposes)

    out_LCFile   = work_dir+'/EPIC_PN_FlareBKGRate.fit'  # Name of the output BKG lightcurve

    # SAS Command
    cmd        = "evselect" # SAS task to be executed                  
    
    # Arguments of SAS Command
    expression = f'#XMMEA_EP&&(PI>={pn_pi_min}&&PI<={pn_pi_max})&&(PATTERN=={pn_pattern})'  # event filter expression
    inargs     = [f'table={eventfile}','withrateset=Y',f'rateset={out_LCFile}','maketimecolumn=Y','timebinsize=100','makeratecolumn=Y',f'expression={expression}']

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

    # Execute SAS task with parameters
    w(cmd, inargs).run()

    # Open event file

    hdu_list = fits.open(eventfile, memmap=True)
    evt_data = Table(hdu_list[1].data)
    prihdu   = hdu_list[1].header

    mask = ((evt_data['PATTERN'] <= pn_pattern) &
            (evt_data['FLAG'] == pn_flag) &
            (evt_data['PI'] >= pn_pi_min) &
            (evt_data['PI'] <= pn_pi_max))

    # Read some information from keywords to be used later on

    if ('INSTRUME' in prihdu):
        ins = prihdu['INSTRUME']
        print("Looking into instrument: "+ins+" \n")
    if ('EXPIDSTR' in prihdu):
        expid = prihdu['EXPIDSTR']
        print("Looking at exposure: "+expid+" \n")

    # Check number of event in initial event file

    print("Events in event file" + " " + eventfile + ": " + str(len(evt_data)) + "\n")
    print("Events in filtered event file" + " " + eventfile + ": " + str(np.sum(mask)) + "\n")
    print("   Filter: PATTERN <= " + str(pn_pattern) +
          " : " + str(pn_pi_min) + " <= E(eV) <= " + str(pn_pi_max) +
          " : " + " FLAG == " + str(pn_flag)+ "\n")

    # Create events image and background lightcurve

    plt.figure(figsize=(20,8))

    pl=1

    xmax=np.amax(evt_data['X'][mask])
    xmin=np.amin(evt_data['X'][mask])
    xmid=(xmax-xmin)/2.+xmin
    ymax=np.amax(evt_data['Y'][mask])
    ymin=np.amin(evt_data['Y'][mask])
    xbin_size=80
    ybin_size=80
    NBINS = (int((xmax-xmin)/xbin_size),int((ymax-ymin)/ybin_size))

    # Plot image
    
    plt.subplot(1, 2, pl)

    img_zero_mpl = plt.hist2d(evt_data['X'][mask], evt_data['Y'][mask], NBINS, cmap='GnBu', norm=LogNorm())

    cbar = plt.colorbar(ticks=[10.,100.,1000.])
    cbar.ax.set_yticklabels(['10','100','1000'])

    plt.title(obsid)
    plt.xlabel('x')
    plt.ylabel('y')

    plt.text(xmid, ymin+0.1*(ymax-ymin), expression, ha='center')

    pl=pl+1
    plt.subplot(1, 2, pl)

    # Plot BKG lightcurve

    plotLC(plt,pn_threshold,out_LCFile)

    pl=pl+1
    plt.show()
    hdu_list.close()

    ############################################
    # Define energy range to filter event file #
    ############################################
    pn_pattern   = 4        # pattern selection
    pn_pi_min    = 200.     # Low energy range eV
    pn_pi_max    = 10000.   # High energy range eV
    pn_threshold = 0.75      # Cut to be applied to filter event file (cts/sec)

    # Define the input and output file names

    in_LCFile         = work_dir+'/EPIC_PN_FlareBKGRate.fit'   # Name of the input BKG lightcurve
    out_gti_set       = work_dir+'/EPIC_PN_gti.fit'            # Name of the output file containing GTI intervals
    out_clean_evtFile = work_dir+'/EPIC_PN_gtiFilteredEvts.ds' # Name of the output Event file filtered by GTI

    # SAS Command
    cmd        = "tabgtigen"
    
    # Arguments of SAS Command
    expression = 'RATE<='+str(pn_threshold)  # event filter expression
    inargs     = [f'table={in_LCFile}',f'gtiset={out_gti_set}',f'expression={expression}']
    
    print("   Filter expression to use: "+expression+" \n")
    print("   SAS command to be executed: "+cmd+", with arguments; \n")
    # Execute SAS task with parameters
    w(cmd, inargs).run()

    # SAS Command
    cmd        = "evselect"
    
    # Arguments of SAS Command
    expression = ('#XMMEA_EP&&FLAG==0&&(PI>='+str(pn_pi_min)+'&&PI<='+str(pn_pi_max)+
                  ')&&(gti('+str(out_gti_set)+',TIME))')
    inargs     = [f'table={eventfile}','withfilteredset=Y',f'filteredset={out_clean_evtFile}',
                  'destruct=Y','keepfilteroutput=T',f'expression={expression}']
    
    print("   Filter expression to use: "+expression+" \n")
    print("   SAS command to be executed: "+cmd+", with arguments; \n")
    # Execute SAS task with parameters
    w(cmd, inargs).run()


    plt.figure(figsize=(20,8))

    pl=1

    hdu_list = fits.open(eventfile, memmap=True)
    evt_data = Table(hdu_list[1].data)
    prihdu   = hdu_list[1].header
    print("Events in event file" + " " + eventfile + ": " + str(len(evt_data)) + "\n")

    gti_hdu_list = fits.open(out_clean_evtFile, memmap=True)
    gti_evt_data = Table(gti_hdu_list[1].data)
    print("Events in GTI clean event file" + " " + out_clean_evtFile + ": " + str(len(gti_evt_data)) + "\n")

    # Create Events image
    
    xmax=np.amax(evt_data['X'])
    xmin=np.amin(evt_data['X'])
    xmid=(xmax-xmin)/2.+xmin
    ymax=np.amax(evt_data['Y'])
    ymin=np.amin(evt_data['Y'])
    xbin_size=80
    ybin_size=80
    NBINS = (int((xmax-xmin)/xbin_size),int((ymax-ymin)/ybin_size))

    plt.subplot(1, 2, pl)

    img_zero_mpl = plt.hist2d(evt_data['X'], evt_data['Y'], NBINS, cmap='GnBu', norm=LogNorm())

    cbar = plt.colorbar(ticks=[10.,100.,1000.])
    cbar.ax.set_yticklabels(['10','100','1000'])

    plt.title(obsid)
    plt.xlabel('x')
    plt.ylabel('y')

    pl=pl+1

    # Create Clean Events image

    xmax=np.amax(gti_evt_data['X'])
    xmin=np.amin(gti_evt_data['X'])
    xmid=(xmax-xmin)/2.+xmin
    ymax=np.amax(gti_evt_data['Y'])
    ymin=np.amin(gti_evt_data['Y'])
    xbin_size=80
    ybin_size=80
    NBINS = (int((xmax-xmin)/xbin_size),int((ymax-ymin)/ybin_size))

    plt.subplot(1, 2, pl)

    img_zero_mpl = plt.hist2d(gti_evt_data['X'], gti_evt_data['Y'], NBINS, cmap='GnBu', norm=LogNorm())

    cbar = plt.colorbar(ticks=[10.,100.,1000.])
    cbar.ax.set_yticklabels(['10','100','1000'])

    plt.title(out_clean_evtFile)
    plt.xlabel('x')
    plt.ylabel('y')

    plt.text(xmid, ymin-0.1*(ymax-ymin), expression, ha='center')

    pl=pl+1

    gti_hdu_list.close()
    hdu_list.close()
    plt.show()

    ###############################################################################
    # Define some parameters to produce the image and the name of the output file #
    ###############################################################################
    xbin=80     # xbin size
    ybin=80     # ybin size
    xcoord='X'  # coordinate system
    ycoord='Y'  # coordinate system

    out_IMFile   = work_dir+f'/EPIC_PN_Image_{obsid}.fit'  # Name of the output Image file 
    # SAS Command
    cmd        = "evselect" # SAS task to be executed                  

    # Arguments of SAS Command
    
    inargs     = [f'table={out_clean_evtFile}','imagebinning=binSize',f'imageset={out_IMFile}','withimageset=yes',f'xcolumn={xcoord}',f'ycolumn={ycoord}',f'ximagebinsize={xbin}',f'yimagebinsize={ybin}']

    print("   SAS command to be executed: "+cmd+", with arguments; \n")
    # Execute the SAS task with the parameters to produce an image
    w(cmd, inargs).run()
    
    # 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")