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

from astropy.convolution import interpolate_replace_nans, convolve, convolve_fft
from astropy.convolution import Gaussian2DKernel

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'
create_folder(products_dir)
mdir=root_path+'/products/mosaic'
create_folder(mdir)

ds9reg_dir=root_path+'/data/ds9reg'


inargs = ['--version']
t = w('sasver', inargs)
t.run()

# run over energies
for idx,e in enumerate(emin):
    label=f"_{idx}"
    e1=emin[idx]
    e2=emax[idx]

    for ikey,key in enumerate(['','m1_','m2_','mos_','pn_']):
        in_cts=f'{mdir}/{key}cts{label}.fits'
        if not (test_file(in_cts,alive=True)):
            continue

        hdul0 = fits.open(in_cts)
        c_data = hdul0[0].data
        c_header = hdul0[0].header
        c_wcs = WCS(c_header)
        hdul0.close()
        
        in_exp=f'{mdir}/{key}exp{label}.fits'
        if not (test_file(in_exp,alive=True)):
            continue

        hdul0 = fits.open(in_exp)
        e_data = hdul0[0].data
        e_header = hdul0[0].header
        e_wcs = WCS(e_header)
        hdul0.close()
        
        in_qpb=f'{mdir}/{key}qpb{label}.fits'
        if not (test_file(in_qpb,alive=True)):
            continue

        hdul0 = fits.open(in_qpb)
        q_data = hdul0[0].data
        q_header = hdul0[0].header
        q_wcs = WCS(q_header)
        hdul0.close()
        
        gauss2d = Gaussian2DKernel(x_stddev=0.530785)

        convol_cts = convolve(c_data, gauss2d)
        out_cts = f'{mdir}/{key}cts{label}_gauss.fits'
        hdul0[0].data=convol_cts
        hdul0.writeto(out_cts,overwrite=True)

        convol_qpb = convolve(q_data, gauss2d)
        out_qpb = f'{mdir}/{key}qpb{label}_gauss.fits'
        hdul0[0].data=convol_qpb
        hdul0.writeto(out_qpb,overwrite=True)

        (nx,ny) = c_data.shape
        map_flx = np.zeros(c_data.shape,dtype=np.float64)
        index = e_data > 0
        map_flx[index]=(convol_cts[index]-convol_qpb[index])/e_data[index]
        out_flx = f'{mdir}/{key}flx{label}_gauss.fits'
        hdul0[0].data=map_flx
        hdul0.writeto(out_flx,overwrite=True)