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


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)
    if(obsid in skip):
        continue

    work_dir = init_work_dir(obsid, products_dir=products_dir)
    os.chdir(work_dir)

    #
    # Combine MOS1+2
    #
    
    spec=[]
    bkg=[]
    resp=[]
    arf=[]
    for imos in [1,2]:
        in_SPSRCFile = work_dir+f'/EPIC_MOS{imos}_source_spectrum.fits'   # Name of the output source spectrum
        in_SPBKGFile = work_dir+f'/EPIC_MOS{imos}_background_spectrum.fits'   # Name of the output background spectrum
        in_RESPFile = work_dir+f'/EPIC_MOS{imos}.rmf'   # Name of the output redistribution
        in_ARFFile = work_dir+f'/EPIC_MOS{imos}.arf'   # Name of the output ancillary
        spec.append(in_SPSRCFile)
        bkg.append(in_SPBKGFile)
        resp.append(in_RESPFile)
        arf.append(in_ARFFile)

    in_spec=" ".join(spec)
    in_bkg=" ".join(bkg)
    in_resp=" ".join(resp)
    in_arf=" ".join(arf)

    out_SPSRCFile = work_dir+f'/EPIC_MOS_merged_source_spectrum.fits'   # Name of the output source spectrum
    out_SPBKGFile = work_dir+f'/EPIC_MOS_merged_background_spectrum.fits'   # Name of the output background spectrum
    out_RESPFile = work_dir+f'/EPIC_MOS_merged.rmf'   # Name of the output redistribution
    out_ARFFile = work_dir+f'/EPIC_MOS_merged.arf'   # Name of the output ancillary

    inargs     = [f'pha=\"{in_spec}\"',f'bkg={in_bkg}',
                  f'rmf=\"{in_resp}\"',f'arf=\"{in_arf}\"',
                  f'allowHEdiff=yes',
                  f'filepha={out_SPSRCFile}',
                  f'filebkg={out_SPBKGFile}',
                  f'filersp={out_RESPFile}']
   
    w("epicspeccombine", inargs).run()

    # rebin the spectra and link associated files
    in_GRPFile = work_dir+f'/EPIC_MOS_merged_spectrum_grp.fits'   # Name of the output specgruop
    inargs     = [f'spectrumset={out_SPSRCFile}','mincounts=30','oversample=3',
                  f'rmfset={out_RESPFile}',
                  f'backgndset={out_SPBKGFile}',f'groupedset={in_GRPFile}']
    w("specgroup", inargs).run()

    #
    # Combine MOS1,2+PN
    #


    in_SPSRCFile = work_dir+f'/EPIC_PN_source_spectrum.fits'   # Name of the output source spectrum
    in_SPBKGFile = work_dir+f'/EPIC_PN_background_spectrum.fits'   # Name of the output background spectrum
    in_RESPFile = work_dir+f'/EPIC_PN.rmf'   # Name of the output redistribution
    in_ARFFile = work_dir+f'/EPIC_PN.arf'   # Name of the output ancillary

    spec.append(work_dir+f'/EPIC_PN_source_spectrum.fits')
    bkg.append(work_dir+f'/EPIC_PN_background_spectrum.fits')
    resp.append(work_dir+f'/EPIC_PN.rmf')
    arf.append(work_dir+f'/EPIC_PN.arf')

    in_spec=" ".join(spec)
    in_bkg=" ".join(bkg)
    in_resp=" ".join(resp)
    in_arf=" ".join(arf)

    out_SPSRCFile = work_dir+f'/EPIC_merged_source_spectrum.fits'   # Name of the output source spectrum
    out_SPBKGFile = work_dir+f'/EPIC_merged_background_spectrum.fits'   # Name of the output background spectrum
    out_RESPFile = work_dir+f'/EPIC_merged.rmf'   # Name of the output redistribution
    out_ARFFile = work_dir+f'/EPIC_merged.arf'   # Name of the output ancillary

    inargs     = [f'pha=\"{in_spec}\"',f'bkg={in_bkg}',
                  f'rmf=\"{in_resp}\"',f'arf=\"{in_arf}\"',
                  f'allowHEdiff=yes',
                  f'filepha={out_SPSRCFile}',
                  f'filebkg={out_SPBKGFile}',
                  f'filersp={out_RESPFile}']
    w("epicspeccombine", inargs).run()

    # rebin the spectra and link associated files
    in_GRPFile = work_dir+f'/EPIC_merged_spectrum_grp.fits'   # Name of the output specgruop
    inargs     = [f'spectrumset={out_SPSRCFile}','mincounts=30','oversample=3',
                  f'rmfset={out_RESPFile}',
                  f'backgndset={out_SPBKGFile}',f'groupedset={in_GRPFile}']
    w("specgroup", inargs).run()