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

mdir=root_path+'/products/mosaic_2021-2024'
create_folder(mdir)

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):
    #if not (idx==1):
    #    continue
    
    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 not (obsid in good):
            continue

        work_dir = f'{products_dir}/{obsid}'
    
        for ikey,key in enumerate(['MOS1','MOS2','PN']):
            shortkey = key.replace("OS","")
            #expkey=f'00{ikey+1}'
            #expkey=f'00{ikey+1}'
            # sci and qpb image are produced by 06_sas_QPB.py
            sci_image=work_dir+f'/EPIC_{key}_sci_image{label}.fits'
            if not (test_file(sci_image, alive=True)):
                continue
            qpb_image=work_dir+f'/EPIC_{key}_qpb_image{label}.fits'
            if not (test_file(qpb_image, alive=True)):
                continue
            # exposure image is produced by eimageget (07_sas_eimageget.py)
            # example: P0862470801PNS003_ima_4exp.fits
            exp_image = get_first_file(work_dir+f'/P{obsid}{shortkey}?00?_ima{label}exp.fits', alive=True)
            if not (exp_image):
                continue
            
            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)

    print("MOS1 ",len(m1_ima))
    print(m1_ima)
    print("MOS2 ",len(m2_ima))
    print(m2_ima)
    print("MOS1+2 ",len(mos_ima))
    print(mos_ima)
    print("PN ",len(pn_ima))
    print(pn_ima)
    print("All ",len(ima))
    print(ima)

    run_mosaic(out_cts=mdir+f'/m1_cts{label}.fits', out_exp=mdir+f'/m1_exp{label}.fits', out_qpb=mdir+f'/m1_qpb{label}.fits', out_sub=mdir+f'/m1_sub{label}.fits', out_flx=mdir+f'/m1_flx{label}.fits', out_pix=mdir+f'/m1_pix{label}.fits', cts=m1_ima,qpb=m1_qpb,exp=m1_exp)
    
    run_mosaic(out_cts=mdir+f'/m2_cts{label}.fits', out_exp=mdir+f'/m2_exp{label}.fits', out_qpb=mdir+f'/m2_qpb{label}.fits', out_sub=mdir+f'/m2_sub{label}.fits', out_flx=mdir+f'/m2_flx{label}.fits', cts=m2_ima,qpb=m2_qpb,exp=m2_exp)
    
    run_mosaic(out_cts=mdir+f'/mos_cts{label}.fits', out_exp=mdir+f'/mos_exp{label}.fits', out_qpb=mdir+f'/mos_qpb{label}.fits', out_sub=mdir+f'/mos_sub{label}.fits', out_flx=mdir+f'/mos_flx{label}.fits', cts=mos_ima,qpb=mos_qpb,exp=mos_exp)
    
    run_mosaic(out_cts=mdir+f'/pn_cts{label}.fits', out_exp=mdir+f'/pn_exp{label}.fits', out_qpb=mdir+f'/pn_qpb{label}.fits', out_sub=mdir+f'/pn_sub{label}.fits', out_flx=mdir+f'/pn_flx{label}.fits', cts=pn_ima,qpb=pn_qpb,exp=pn_exp)

    
    run_mosaic(out_cts=mdir+f'/cts{label}.fits', out_exp=mdir+f'/exp{label}.fits', out_qpb=mdir+f'/qpb{label}.fits', out_sub=mdir+f'/sub{label}.fits', out_flx=mdir+f'/flx{label}.fits', cts=ima,qpb=qpb,exp=exp)
    
    
