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

ds9reg_dir=root_path+'/data/ds9reg'


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

files = glob.glob(archive_dir+'/*')


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

    # all MOS1
    m1_ima=[]
    m1_qpb=[]
    m1_exp=[]

    # all MOS2
    m2_ima=[]
    m2_qpb=[]
    m2_exp=[]

    # all MOS
    mos_ima=[]
    mos_qpb=[]
    mos_exp=[]

    # all PNs
    pn_ima=[]
    pn_qpb=[]
    pn_exp=[]

    # all MOS1,2,PN
    ima=[]
    qpb=[]
    exp=[]
    
    for obsid in files:
        obsid = os.path.basename(obsid)
        if(obsid in skip):
            continue

        work_dir = f'{products_dir}/{obsid}'
    
        for ikey,key in enumerate(['MOS1','MOS2','PN']):
            shortkey = key.replace("OS","")
            expkey=f'S00{ikey+1}'
            # sci and qpb image are produced by 06_sas_QPB.py
            sci_image=work_dir+f'/EPIC_{key}_sci_image{label}.fits'
            test_file(sci_image)
            qpb_image=work_dir+f'/EPIC_{key}_qpb_image{label}.fits'
            test_file(qpb_image)
            # exposure image is produced by eimageget (07_sas_eimageget.py)
            # example: P0862470801PNS003_ima_4exp.fits
            exp_image=work_dir+f'/P{obsid}{shortkey}{expkey}_ima{label}exp.fits'
            test_file(exp_image)

            ima.append(sci_image)
            qpb.append(qpb_image)
            exp.append(exp_image)
            
            if(key=='PN'):
                pn_ima.append(sci_image)
                pn_qpb.append(qpb_image)
                pn_exp.append(exp_image)
            if(key=='MOS1'):
                m1_ima.append(sci_image)
                m1_qpb.append(qpb_image)
                m1_exp.append(exp_image)
            if(key=='MOS2'):
                m2_ima.append(sci_image)
                m2_qpb.append(qpb_image)
                m2_exp.append(exp_image)
            if(key=='MOS1' or key=='MOS2'):
                mos_ima.append(sci_image)
                mos_qpb.append(qpb_image)
                mos_exp.append(exp_image)

    
    run_mosaic(outfile_cts=mosaic_dir+f'/m1_cts{label}.fits',
               outfile_exp=mosaic_dir+f'/m1_exp{label}.fits',
               outfile_qpb=mosaic_dir+f'/m1_qpb{label}.fits',
               outfile_sub=mosaic_dir+f'/m1_sub{label}.fits',
               outfile_flx=mosaic_dir+f'/m1_flx{label}.fits',
               outfile_pix=mosaic_dir+f'/m1_pix{label}.fits',
               cts=m1_ima,qpb=m1_qpb,exp=m1_exp)
    run_mosaic(outfile_cts=mosaic_dir+f'/m2_cts{label}.fits',
               outfile_exp=mosaic_dir+f'/m2_exp{label}.fits',
               outfile_qpb=mosaic_dir+f'/m2_qpb{label}.fits',
               outfile_sub=mosaic_dir+f'/m2_sub{label}.fits',
               outfile_flx=mosaic_dir+f'/m2_flx{label}.fits',
               cts=m2_ima,qpb=m2_qpb,exp=m2_exp)
    run_mosaic(outfile_cts=mosaic_dir+f'/mos_cts{label}.fits',
               outfile_exp=mosaic_dir+f'/mos_exp{label}.fits',
               outfile_qpb=mosaic_dir+f'/mos_qpb{label}.fits',
               outfile_sub=mosaic_dir+f'/mos_sub{label}.fits',
               outfile_flx=mosaic_dir+f'/mos_flx{label}.fits',
               cts=mos_ima,qpb=mos_qpb,exp=mos_exp)
    run_mosaic(outfile_cts=mosaic_dir+f'/pn_cts{label}.fits',
               outfile_exp=mosaic_dir+f'/pn_exp{label}.fits',
               outfile_qpb=mosaic_dir+f'/pn_qpb{label}.fits',
               outfile_sub=mosaic_dir+f'/pn_sub{label}.fits',
               outfile_flx=mosaic_dir+f'/pn_flx{label}.fits',
               cts=pn_ima,qpb=pn_qpb,exp=pn_exp)
    run_mosaic(outfile_cts=mosaic_dir+f'/cts{label}.fits',
               outfile_exp=mosaic_dir+f'/exp{label}.fits',
               outfile_qpb=mosaic_dir+f'/qpb{label}.fits',
               outfile_sub=mosaic_dir+f'/sub{label}.fits',
               outfile_flx=mosaic_dir+f'/flx{label}.fits',
               cts=ima,qpb=qpb,exp=exp)