#!/usr/bin/env python
# -*- coding:  utf-8  -*-
"""
НАЗВАНИЕ:

 04_mosaics.py
 
 
НАЗНАЧЕНИЕ:

   Собирает мозайки в разных энергетических диапазонах.
 
ВЫЗОВ:

   esass
   ./04_mosaics.py


УПРАВЛЕНИЕ:

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

   index : Выбранный энергетический диапазон


ВЫВОД:
 
   Выходные файлы записываются в директорию 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 os.path import dirname
import inspect
import pickle



import uds

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


""" find UDS root dir """
#root_path=dirname(dirname(dirname(inspect.getfile(uds))))
"""
ftools does not like long file path names,
for this reason, we use relative path here
"""
root_path='..'
print("UDS root path: {}".format(root_path))

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

local_run = False

outkey="tm0"

do_init            = True
do_merge           = True
do_detmask         = True
do_expmap          = True
do_erbox1          = True # local mode
do_erbackmap1      = True # 
do_erbox2          = True # map mode, with background map
do_erbackmap2      = True # 
do_erbox3          = True # map mode, with background map
do_erbackmap3      = True # 
do_ersensmap       = False
do_ermldet         = True
do_fixcat          = False  # only for index=0
do_fixxmm          = False # prepare forced photometry, only for index=0
do_apetool         = False
do_catprep         = True
do_filter_catalog  = False

do_cross_match     = False

index=3

forced=False
""" If forced=True, take input catalog from energy range en0 """

comm='' # for 4XMM-DR12 forced photometry use '-xmm'


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

attcorr=True

events=[]
expmaps=[]
bkgmaps=[]
for tmkey in keylist_tm.keys():
    print("TM{} in work... init events".format(tmkey))
    for datakey in keylist_tm[tmkey]:
        print("--> {}".format(datakey))
        """ Подготавливаем списки событий индивидуальных наблюдений """
        outfile_evtool,outfile_expmap=init_events(key=datakey,attcorr=attcorr,
                                                  eband_index=eband[index],
                                                  infile_dir=infile_dir,
                                                  outfile_dir=outfile_dir,
                                                  do_init=do_init,
                                                  do_obsmode=False,
                                                  do_center=False,
                                                  do_evtool=True,
                                                  do_expmap=False,
                                                  vign=vign,
                                                  ra_cen=ra_cen, de_cen=de_cen,
                                                  emin_kev=emin_kev[index],
                                                  emax_kev=emax_kev[index])
        events.append(outfile_evtool)
        expmaps.append(outfile_expmap)

   
""" Собираем общий список событий """    
outfile_evtool="{}_EventList_en{}.fits".format(os.path.join(outfile_dir,outkey),
                                               eband[index])

if(do_merge==True):
    do_evtool_esass(events=events, outfile=outfile_evtool)


    """ makes detmask from TM exposures """
detmask="{}/{}_DetectorMask_en{}{}".format(outfile_dir,
                                           outkey,
                                           eband[index],
                                           outfile_post)
if(do_detmask==True):
   create_detmask_merged(expmaps,detmask,minval=100)

"""
Собираем общую карту экспозиции, обратите внимание на коэффициент 7.
Экспозиция рассчитывается на 7 телескопов.
outfile_expmap="{}_ExposureMap_en{}.{}.fits".format(os.path.join(outfile_dir,outkey), eband[index], vignetting)
outfile_bkgmap="{}_BackMap_en{}.{}.fits".format(os.path.join(outfile_dir,outkey), eband[index], vignetting)
"""    
outfile_expmap="{}_ExposureMap_en{}.{}{}".format(os.path.join(outfile_dir,outkey),
                                              eband[index],vignetting,
                                              outfile_post)
if(do_expmap==True):
   create_expmap_merged(expmaps,outfile_expmap,scale=7.0)

outfile_boxlist1="{}/{}_BoxList1_en{}{}".format(outfile_dir,outkey, eband[index], outfile_post)
if(do_erbox1==True):
   cmd=["erbox",
        "images=\'{}\'".format(outfile_evtool),
        "boxlist=%s" %(outfile_boxlist1),
        "expimages=\'{}\'".format(outfile_expmap),
        "detmasks=\'{}\'".format(detmask),
        "emin=\'{}\'".format(emin_ev[index]),
        "emax=\'{}\'".format(emax_ev[index]),
        "ecf=\'{}\'".format(ecf[index]),
        "nruns=2",
        "likemin=6.0",
        "boxsize=4",
        "compress_flag=N",
        "bkgima_flag=N",
        "expima_flag=Y",
        "detmask_flag=Y"
        ]
   remove_file(outfile_boxlist1)
   print((" ").join(cmd))
   os.system((" ").join(cmd))
   save_ds9reg(outfile_boxlist1)

""" Background map 1 """
outfile_backmap1="{}_BackMap1_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
cheese_mask="{}_CheeseMask1_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
if(do_erbackmap1==True):
   do_erbackmap_esass(outfile_evtool,outfile_expmap,outfile_boxlist1,detmask,emin_ev[index],emax_ev[index],
                      outfile_backmap1,cheese_mask)

outfile_boxlist2="{}/{}_BoxList2_en{}{}".format(outfile_dir,outkey, eband[index], outfile_post)
if(do_erbox2==True):
   cmd=["erbox",
        "images=\'{}\'".format(outfile_evtool),
        "boxlist=%s" %(outfile_boxlist2),
        "expimages=\'{}\'".format(outfile_expmap),
        "detmasks=\'{}\'".format(detmask),
        "emin=\'{}\'".format(emin_ev[index]),
        "emax=\'{}\'".format(emax_ev[index]),
        "ecf=\'{}\'".format(ecf[index]),
        "nruns=2",
        "likemin=6.0",
        "boxsize=4",
        "compress_flag=N",
        "bkgima_flag=Y",
        "bkgimages={}".format(outfile_backmap1),
        "expima_flag=Y",
        "detmask_flag=Y"
        ]
   remove_file(outfile_boxlist2)
   print((" ").join(cmd))
   os.system((" ").join(cmd))
   save_ds9reg(outfile_boxlist2)

""" Background map 2 """
outfile_backmap2="{}_BackMap2_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
cheese_mask="{}_CheeseMask2_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
if(do_erbackmap2==True):
   do_erbackmap_esass(outfile_evtool,outfile_expmap,outfile_boxlist2,detmask,emin_ev[index],emax_ev[index],
                      outfile_backmap2,cheese_mask)

outfile_boxlist3="{}/{}_BoxList3_en{}{}".format(outfile_dir,outkey, eband[index], outfile_post)
if(do_erbox3==True):
   cmd=["erbox",
        "images=\'{}\'".format(outfile_evtool),
        "boxlist=%s" %(outfile_boxlist3),
        "expimages=\'{}\'".format(outfile_expmap),
        "detmasks=\'{}\'".format(detmask),
        "emin=\'{}\'".format(emin_ev[index]),
        "emax=\'{}\'".format(emax_ev[index]),
        "ecf=\'{}\'".format(ecf[index]),
        "nruns=2",
        "likemin=6.0",
        "boxsize=4",
        "compress_flag=N",
        "bkgima_flag=Y",
        "bkgimages={}".format(outfile_backmap2),
        "expima_flag=Y",
        "detmask_flag=Y"
        ]
   remove_file(outfile_boxlist3)
   print((" ").join(cmd))
   os.system((" ").join(cmd))
   save_ds9reg(outfile_boxlist3)

""" Background map 3 """
outfile_backmap3="{}_BackMap3_en{}.{}{}".format(os.path.join(outfile_dir,outkey), eband[index], vignetting, outfile_post)
cheese_mask="{}_CheeseMask3_en{}.{}{}".format(os.path.join(outfile_dir,outkey), eband[index], vignetting, outfile_post)
if(do_erbackmap3==True):
    boxlist3 = outfile_boxlist3 if(forced == False) else "{}/{}_BoxList3_en{}{}".format(outfile_dir,outkey, eband[0], outfile_post)
    do_erbackmap_esass(outfile_evtool,outfile_expmap,boxlist3,detmask,emin_ev[index],emax_ev[index],
                       outfile_backmap3,cheese_mask)

if(forced==True):
    mllist="{}_MaxLikSourceList_en{}.forced{}{}".format(os.path.join(outfile_dir,outkey), eband[index], comm, outfile_post)
    srcmap="{}_SourceMap_en{}.forced{}{}".format(os.path.join(outfile_dir,outkey), eband[index], comm, outfile_post)

    """ for en1,2,3,6 give mllist from en0 as input """
    #boxlist3="{}_MaxLikSourceList_en{}.forced{}{}".format(os.path.join(outfile_dir,outkey), eband[0], comm, outfile_post)
    #if(index==0):
    boxlist3="{}_MaxLikSourceList_en{}.fixed{}{}".format(os.path.join(outfile_dir,outkey), eband[0], comm, outfile_post)
    if not (os.path.exists(boxlist3)):
        print("{} not found. Run do_fixcat=True, index=0, forced=False".format(boxlist3))
        sys.exit()
            
    add_specific_columns(boxlist3)
    fitpos_flag="fitpos_flag=no"
    fitext_flag="fitext_flag=no"
    nmulsou = "nmulsou=1"
    nmaxfit="nmaxfit=10"
    multrad="multrad=15." 
    cutrad="cutrad=15."  
    if(index == 3 or index == 6):
        """ for hard band take unvignetted background """
        outfile_backmap3="{}_BackMap3_en{}.{}{}".format(os.path.join(outfile_dir,outkey), eband[index], "novign", outfile_post)

else:
    mllist="{}_MaxLikSourceList_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
    srcmap="{}_SourceMap_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
    boxlist3 = outfile_boxlist3
    fitpos_flag="fitpos_flag=yes"
    fitext_flag="fitext_flag=yes"
    nmulsou = "nmulsou=2"
    nmaxfit="nmaxfit=4"
    multrad="multrad=20." 
    cutrad="cutrad=20." 
    """ allow ermpldet to split sources (no more than two) """ 

cmd=["ermldet",
     "mllist={}".format(mllist),
     "boxlist=%s" %(boxlist3),
     "images=\'{}\'".format(outfile_evtool),
     "expimages=\'{}\'".format(outfile_expmap),
     "detmasks=\'{}\'".format(detmask),
     "bkgimages=\'{}\'".format(outfile_backmap3),
     "emin=\'{}\'".format(emin_ev[index]),
     "emax=\'{}\'".format(emax_ev[index]),
     "ecf=\'{}\'".format(ecf[index]),
     "hrdef=",
     "likemin=0.",
     "extlikemin=5.",
     "compress_flag=N",
     cutrad, 
     multrad, 
     "extmin=2.0",
     "extmax=35.0",
     #"bkgima_flag=Y", looks outdated
     "expima_flag=Y",
     "detmask_flag=Y",
     "shapelet_flag=yes",
     "photon_flag=yes",
     "extentmodel=beta",
     "thres_flag=N",
     "thres_col=like",
     "thres_val=30.",
     nmaxfit,
     nmulsou,
     fitpos_flag,
     fitext_flag,
     "srcima_flag=yes",
     "srcimages=\'{}\'".format(srcmap)
     ]

if(do_ersensmap==True):
   detlike=10
   create_sensmap(sensmap="{}_SensitivityMap_dl{}_en{}{}".format(os.path.join(outfile_dir,outkey), detlike,
                                                                 eband[index], outfile_post),
                  areatab="{}_AreaTable_dl{}_en{}{}".format(os.path.join(outfile_dir,outkey),
                                                            detlike, eband[index], outfile_post),
                  expmap=outfile_expmap, backmap=outfile_backmap3,detlike=detlike,
                  detmask=detmask, emin=emin_ev[index], emax=emax_ev[index],ecf=ecf[index])
   """
   detlike=6
   create_sensmap(sensmap="{}_SensitivityMap_dl{}_en{}{}".format(os.path.join(outfile_dir,outkey), detlike,
                                                                 eband[index], outfile_post),
                  areatab="{}_AreaTable_dl{}_en{}{}".format(os.path.join(outfile_dir,outkey),
                                                            detlike, eband[index], outfile_post),
                  expmap=outfile_expmap, backmap=outfile_backmap3,detlike=detlike,
                  detmask=detmask, emin=emin_ev[index], emax=emax_ev[index],ecf=ecf[index])
   """

if(do_ermldet==True):
    test_exe('ermldet')
    if(vign==False):
        print('Run ermldet with vignetted exposure!')
        sys.exit()
    remove_file(mllist)
    remove_file(srcmap)
    print(cmd)
    runme(cmd, local_run=local_run)
    print(cmd)
    save_ermldet_ds9reg(mllist,scale=60*60,label='det_like')
    save_ermldet_ds9reg(mllist,scale=60*60,label='id_src')

    correct_fluxerr_ermldet_forced(mllist)
    if(forced==True):
        result = check_ermldet_forced(mllist)
        # for a some reason, for an arbitrary energy band, ermldet break order of sources. Do this forced correction. 
        if(result == False):
             correct_srcid_ermldet_forced(mllist)

if(do_fixcat==True):
    if not index == 0:
        print("ERROR: You can fix only reference catalog for en0.")
        sys.exit()
    if forced == True:
        print("ERROR: You can fix only non-forced catalog for en0.")
        sys.exit()

    srcs_remove=[341,446,346,96] 
    srcs_add = {'4XMM J021738.8-051257':[34.4117002, -5.2159135, 0.624],# 341
                '4XMM J021733.8-051311':[34.3910215,-5.2199877,2.247],# 341
                '4XMM J021929.4-051220':[34.8725460,-5.2056849,1.074],#446
                '4XMM J021930.7-051225':[34.8782267,-5.2072112,0.624],# 446
                '4XMM J021945.2-045331':[34.9383593,-4.8919843,1.538],#346
                '4XMM J021929.4-043224':[34.8728586,-4.5400022,0.659555],#96
                #'4XMM J021929.4-043224':[34.8728586,-4.5400022, 0.660], 
                #'4XMM J021831.8-050059':[34.6328841,-5.0163909,1.529], 
                #'4XMM J022131.1-050027':[35.3797879,-5.0075498,0.941], 
                #'4XMM J022129.5-045914':[35.3732136,-4.9874025,0.332], 
                #'4XMM J022026.3-050251':[35.1098619,-5.0476199,0.551],
                #'4XMM J021925.4-042647':[34.8559099,-4.4465007,1.366],
                #'4XMM J021910.9-045108':[34.7954311,-4.8522901,0.898],
                #'4XMM J021945.2-045331':[34.9383593,-4.8919843,1.538],
                #'4XMM J021733.8-051311':[34.3910215,-5.2199877,2.247],
                }
    fix_catalog(mllist=mllist,refimage=outfile_evtool, srcs_remove=srcs_remove, srcs_add=srcs_add)
    """
    Note that fix_catalog added ID_SRC to each XMM source.
    Next, we save forced XMM sources (with new ID_SRC!) for later catalog compilation
    """
    with open(mllist.replace(".fits", ".xmm.pickle"), 'wb') as f:
        pickle.dump(srcs_add, f)

if(do_fixxmm==True):
    if not index == 0:
        print("ERROR: You can fix only reference catalog for en0.")
        sys.exit()
    if forced == True:
        print("ERROR: You can fix only non-forced catalog for en0.")
        sys.exit()

    fix_xmm_sources(mllist=mllist,refimage=outfile_evtool, xmm_catalog='../data/4XMM-DR12/4XMM_DR12cat_slim_v1.0_UDS.fits.catalog')

        
if(do_apetool==True):
    psfmap="{}_PsfMap{}".format(os.path.join(outfile_dir,outkey), outfile_post)
    #remove_file(psfmap)
    #cmd=["apetool",
    #     "images=\'{}\'".format(outfile_evtool),
    #     "psfmaps=\'{}\'".format(psfmap),
    #     "psfmapflag=yes",]
    #runme(cmd, local_run=local_run)
    
    cmd=["apetool",
         "mllist={}".format(mllist),
         "apelistout={}".format(mllist), # give the same file
         "images=\'{}\'".format(outfile_evtool),
         "expimages=\'{}\'".format(outfile_expmap),
         "detmasks=\'{}\'".format(detmask),
         "bkgimages=\'{}\'".format(outfile_backmap3),
         "emin=\'{}\'".format(emin_ev[index]),
         "emax=\'{}\'".format(emax_ev[index]),
         "srcimages=\'{}\'".format(srcmap),
         "psfmaps={}".format(psfmap),
         "psfmapflag=no",
         "stackflag=no",
         "apexflag=yes",
         "apesenseflag=no",
         "eefextract=0.65",
         "cutrad=15",
         "eindex=1",]
    runme(cmd, local_run=local_run)

if(forced==True):
    catprep="{}_SourceCatalog_en{}.forced{}{}".format(os.path.join(outfile_dir,outkey), eband[index], comm, outfile_post)
else:
    catprep="{}_SourceCatalog_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)

if(do_catprep==True):
   cmd=["catprep",
        "infile={}".format(mllist),
        "outfile={}".format(catprep),]
   remove_file(catprep)
   runme(cmd, local_run=local_run)


if(do_filter_catalog==True):
    #filter_mllist(mllist,expcut=5000.0,dlcut=10.0,dlmin=10,dlmax=10000)
    """ works the same """
    filter_catprep(catprep,expcut=5000.0,dlmin=10,dlmax=10000,outkey='bright')
    #filter_catprep(catprep,expcut=5000.0,dlmin=6,dlmax=10,outkey='faint')

if(do_cross_match==True):
    crossmatch_shu2019(catprep,dlmin=10,refimage=outfile_evtool,crval=[ra_cen, de_cen],
                       catalog=root_path+"/data/Gaia_unWISE/Gaia_unWISE_UDS.fits.catalog")
