#!/usr/bin/env python
"""
НАЗВАНИЕ:

 01_init_obs.py
 
 
НАЗНАЧЕНИЕ:

   Подготавливает списки событий в разных энергетических диапазонах.
   Производит списки источников в наждом наблюдении и делает астрокоррекцию с помощью wcs_match/wcs_update
 
ВЫЗОВ:

   conda activate ciao-4.15
   esass
   ./01_init_obs.py


УПРАВЛЕНИЕ:

   Запуск отдельных команд управляется переменными, например: do_init = True
   Выбранные энергетические диапазоны управляется массивом eband_selected
 
ПАРАМЕТРЫ:
   eband_selected : Выбранные энергетические диапазоны


ВЫВОД:
 
   Выходные файлы записываются в директорию outfile_dir
  

ИСТОРИЯ:
   Роман Кривонос, ИКИ РАН, krivonos@cosmos.ru
   Март 2023
   
"""

from astropy.wcs import WCS
from astropy.io import fits
import sys, os, os.path, time, subprocess
from pathlib import Path
import numpy as np
import glob
from multiprocessing import Pool

from os.path import dirname
import inspect
import uds

from uds.utils import *
from uds.config import *


""" find UDS root dir """
root_path=dirname(dirname(dirname(inspect.getfile(uds))))
print("UDS root path: {}".format(root_path))

infile_dir=root_path+'/data/processed'
outfile_dir=root_path+'/products'
create_folder(outfile_dir)

run_Pool=False

do_init         = False
do_ermask       = False

do_erbox1       = False # local mode
do_erbackmap1   = False # 
do_erbox2       = False # map mode, with background map
do_erbackmap2   = False # 
do_erbox3       = False # map mode, with background map
do_erbackmap3   = False # 

do_ermldet      = False
do_catprep      = False
do_cross_match  = False

do_astro_corr   = False # search optimal shift
do_astro_update = True 

do_wcs_match    = False # Chandra task -- DEPRECATED
do_wcs_update   = False # Chandra task -- DEPRECATED

eband_selected=[5]

vign=True
vignetting = 'vign' if (vign==True) else 'novign'

    
def runme(datakey):
    """ runs datakey over energy bands """
    events=[]
    expmaps=[]
    outfile_boxlist1=[]
    outfile_boxlist2=[]
    outfile_boxlist3=[]
    outfile_backmap1=[]
    outfile_backmap2=[]
    outfile_backmap3=[]
    cheesemask=[]
    bkgimage=[]
    srcmaps=[]

    print(datakey)
    print('module name:', __name__)
    print('parent process:', os.getppid())
    print('process id:', os.getpid())
    
    for ii in range(len(eband_selected)):
        index=eband_selected[ii]
        print("\t>>> Energy band en{} -- {}-{} keV".format(eband[index],emin_kev[index],emax_kev[index]))
        outfile_evtool, outfile_expmap = init_events(key=datakey, eband_index=eband[index],
                                                     infile_dir=infile_dir,
                                                     outfile_dir=outfile_dir,
                                                     do_init=do_init,
                                                     do_obsmode=True,
                                                     do_center=True,
                                                     do_evtool=True,
                                                     do_expmap=True,
                                                     vign=vign,
                                                     ra_cen=ra_cen, de_cen=de_cen,
                                                     emin_kev=emin_kev[index],
                                                     emax_kev=emax_kev[index])
        expmaps.append(outfile_expmap)
        events.append(outfile_evtool)

        
        # After astrometry-corrected files (*.attcorr.fits) are obtained, one can take them as original, in order to check the full chain:
        #events.append(outfile_evtool.replace(".fits", ".attcorr.fits"))
        

    """ Detmask """
    
    detmask="{}_DetectionMask{}".format(os.path.join(outfile_dir,datakey), outfile_post)
    if(do_ermask==True):
        cmd=["ermask",
             "expimage=%s" %(expmaps[0]), # use the first exposure maps calculated for that skyfield, independent of the energy band
             "detmask=%s" %(detmask),
             "threshold1=0.01",
             "threshold2=10.0",
             "regionfile_flag=no"
             ]
        remove_file(detmask)
        print((" ").join(cmd))
        os.system((" ").join(cmd))

    for ii in range(len(eband_selected)):
        index=eband_selected[ii]
        print("\t>>> Energy band en{} -- {}-{} keV".format(eband[index],emin_kev[index],emax_kev[index]))

        """ erbox in local mode """
        outfile_boxlist1.append("{}_BoxList1_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        if(do_erbox1==True):
            """ erbox in local mode """
            cmd=["erbox",
                 "images=%s" %(events[ii]),
                 "boxlist=%s" %(outfile_boxlist1[ii]),
                 "expimages=%s" %(expmaps[ii]),
                 "detmasks=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "ecf=1.0",
                 "nruns=2",
                 "likemin=6.0",
                 "boxsize=4",
                 "compress_flag=N",
                 "bkgima_flag=N",
                 "expima_flag=Y",
                 "detmask_flag=Y"
                 ]
            
            remove_file(outfile_boxlist1[ii])

            print((" ").join(cmd))
            os.system((" ").join(cmd))
            save_ds9reg(outfile_boxlist1[ii])


        outfile_backmap1.append("{}_BackMap1_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        cheese_mask="{}_CheeseMask1_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)
        if(do_erbackmap1==True):
            """ back map 1 """
            cmd=["erbackmap",
                 "image=%s" %(events[ii]),
                 "expimage=%s" %(expmaps[ii]),
                 "boxlist=%s" %(outfile_boxlist1[ii]),
                 "detmask=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "bkgimage=%s" %(outfile_backmap1[ii]),
                 "cheesemask=%s" %(cheese_mask),
                 "idband=1",
                 "scut=0.001",
                 "mlmin=6",
                 "maxcut=0.5",
                 "fitmethod=smooth smoothval=15",
                 "snr=40.",
                 ]

            remove_file(cheese_mask)
            remove_file(outfile_backmap1[ii])

            os.system((" ").join(cmd))
            print((" ").join(cmd))

        outfile_boxlist2.append("{}_BoxList2_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        if(do_erbox2==True):
            """ erbox in background mode """
            cmd=["erbox",
                 "images=%s" %(events[ii]),
                 "boxlist=%s" %(outfile_boxlist2[ii]),
                 "expimages=%s" %(expmaps[ii]),
                 "detmasks=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "ecf=1.0",
                 "nruns=2",
                 "likemin=4.0",
                 "boxsize=4",
                 "compress_flag=N",
                 "bkgima_flag=Y", 
                 "bkgimages={}".format(outfile_backmap1[ii]),
                 "expima_flag=Y",
                 "detmask_flag=Y"
                 ]

            remove_file(outfile_boxlist2[ii])
            
            print((" ").join(cmd))
            os.system((" ").join(cmd))
            save_ds9reg(outfile_boxlist2[ii])


        outfile_backmap2.append("{}_BackMap2_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        cheese_mask="{}_CheeseMask2_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)
        if(do_erbackmap2==True):
            """ back map 2 """
            cmd=["erbackmap",
                 "image=%s" %(events[ii]),
                 "expimage=%s" %(expmaps[ii]),
                 "boxlist=%s" %(outfile_boxlist2[ii]),
                 "detmask=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "bkgimage=%s" %(outfile_backmap2[ii]),
                 "cheesemask=%s" %(cheese_mask),
                 "idband=1",
                 "scut=0.001",
                 "mlmin=6",
                 "maxcut=0.5",
                 "fitmethod=smooth smoothval=15",
                 "snr=40.",
                 ]

            remove_file(cheese_mask)
            remove_file(outfile_backmap2[ii])

            os.system((" ").join(cmd))
            print((" ").join(cmd))

        outfile_boxlist3.append("{}_BoxList3_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        if(do_erbox3==True):
            """ erbox in map mode FINAL """
            cmd=["erbox",
                 "images=%s" %(events[ii]),
                 "boxlist=%s" %(outfile_boxlist3[ii]),
                 "expimages=%s" %(expmaps[ii]),
                 "detmasks=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "ecf=1.0",
                 "nruns=2",
                 "likemin=4.0",
                 "boxsize=4",
                 "compress_flag=N",
                 "bkgima_flag=Y", 
                 "bkgimages={}".format(outfile_backmap2[ii]),
                 "expima_flag=Y",
                 "detmask_flag=Y"
                 ]
        
            remove_file(outfile_boxlist3[ii])
        
            print((" ").join(cmd))
            os.system((" ").join(cmd))
            save_ds9reg(outfile_boxlist3[ii])

        outfile_backmap3.append("{}_BackMap3_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post))
        cheese_mask="{}_CheeseMask3_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)
        if(do_erbackmap3==True):
            """ back map 3 FINAL """
            cmd=["erbackmap",
                 "image=%s" %(events[ii]),
                 "expimage=%s" %(expmaps[ii]),
                 "boxlist=%s" %(outfile_boxlist3[ii]),
                 "detmask=%s" %(detmask),
                 "emin=%s" %(emin_ev[index]),
                 "emax=%s" %(emax_ev[index]),
                 "bkgimage=%s" %(outfile_backmap3[ii]),
                 "cheesemask=%s" %(cheese_mask),
                 "idband=1",
                 "scut=0.001",
                 "mlmin=6",
                 "maxcut=0.5",
                 "fitmethod=smooth smoothval=15",
                 "snr=40.",
                 ]

            remove_file(cheese_mask)
            remove_file(outfile_backmap3[ii])

            os.system((" ").join(cmd))
            print((" ").join(cmd))

        mllist="{}_MaxLikSourceList_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)
        srcmap="{}_SourceMap_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)

        cmd=["ermldet",
             "mllist=%s" %(mllist),
             "boxlist=%s" %(outfile_boxlist3[ii]),
             "images=%s" %(events[ii]),
             "expimages=%s" %(expmaps[ii]),
             "detmasks=%s" %(detmask),
             "bkgimages=%s" %(outfile_backmap3[ii]),
             "emin=%s" %(emin_ev[index]),
             "emax=%s" %(emax_ev[index]),
             "hrdef=",
             "ecf={}".format(ecf[index]),
             "likemin=5.",
             "extlikemin=6.",
             "compress_flag=N",
             "cutrad=15.",
             "multrad=20.",
             "extmin=2.0",
             "extmax=15.0",
             #"bkgima_flag=Y", looks outdated
             "expima_flag=Y",
             "detmask_flag=Y",
             "extentmodel=beta",
             "thres_flag=N",
             "thres_col=like",
             "thres_val=30.",
             "nmaxfit=4",
             "nmulsou=2",
             "fitext_flag=yes",
             "srcima_flag=yes",
             "srcimages=%s" %(srcmap)
             
         ]

        if(do_ermldet==True):
            test_exe('ermldet')
            remove_file(mllist)
            remove_file(srcmap)

            os.system((" ").join(cmd))
            print((" ").join(cmd))
            save_ermldet_ds9reg(mllist,scale=60*60)

        catprep="{}_SourceCatalog_en{}{}".format(os.path.join(outfile_dir,datakey), eband[index], outfile_post)
        catprep_en0="{}_SourceCatalog_en{}{}".format(os.path.join(outfile_dir,datakey), eband[0], outfile_post)
        if(do_catprep==True):
            cmd=["catprep",
                 "infile={}".format(mllist),
                 "outfile={}".format(catprep),]
            remove_file(catprep)
            os.system((" ").join(cmd))
            print((" ").join(cmd))
            save_catprep_ds9reg(catprep,scale=60*60)

        if(do_cross_match==True):
            crossmatch_shu2019(catprep,dlmin=10,refimage=events[ii],crval=wcslist[datakey],
                               catalog=root_path+"/data/Gaia_unWISE/Gaia_unWISE_UDS.fits.catalog",errlim=5.0)

        if(do_astro_corr==True and eband[index]=='0'):
            """ run astro_corr for 0.3-2.3 keV only """
            wcs_astro_corr(catprep)
            #wcs_match_ciao(catprep, method='rst',radius=12,residlim=0,residtype=0,residfac=1)

        if(do_astro_update==True):
            """ run astro_corr for 0.3-2.3 keV only """
            attcorr=wcs_update_shift(events[ii],flog=catprep_en0.replace(".fits", ".shift.log"))
            do_evtool_esass(evfile=attcorr,outfile=attcorr,rmlock=False, do_center=True, ra_cen=ra_cen, de_cen=de_cen)


        if(do_wcs_match==True and eband[index]=='0'):
            """ run wcs_match for 0.3-2.3 keV only """
            wcs_match_ciao(catprep, method='trans',radius=12,residlim=5)
            #wcs_match_ciao(catprep, method='rst',radius=12,residlim=0,residtype=0,residfac=1)

        if(do_wcs_update==True):
            """ use 0.3-2.3 keV transform matrix for all other bands """ 
            attcorr=wcs_update_ciao(events[ii],crval=wcslist[datakey],transformfile=catprep_en0.replace(".fits", ".xfm"),clean=False)
            do_evtool_esass(evfile=attcorr,outfile=attcorr,rmlock=False, do_center=True, ra_cen=ra_cen, de_cen=de_cen)


"""
# individual run, testing
runme("tm7_obs_1")
runme("tm5_obs_1")
runme("tm6_scan_1")
"""



if(run_Pool==True):
    # parallel run
    items=[]
    for tmkey in keylist_tm.keys():
        for datakey in keylist_tm[tmkey]:
            items.append(datakey)     
        
    with Pool() as pool:
        pool.map(runme, items)

else:
    # conventional run
    for tmkey in keylist_tm.keys():
        for datakey in keylist_tm[tmkey]:
            print("--> {}".format(datakey))
            runme(datakey)
            #sys.exit()
