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

 05_srctool.py
 
 
НАЗНАЧЕНИЕ:

   Запускает scrtool для самого широкого канала 0.2-10 кэВ, чтобы спектры имели самое полное покрытие по энергиям. Список источников берется из 0.3-2.3 кэВ.

ВЫЗОВ:

   esass
   ./05_srctool.py


УПРАВЛЕНИЕ:

   Требуется запуск предыдущего скрипта 04_mosaics.py
 
ПАРАМЕТРЫ:

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


ВЫВОД:
 
   Выходные файлы записываются в директорию outfile_dir/srctool_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 uds
from scipy.stats import norm
import matplotlib.pyplot as plt

import pandas as pd

from uds.utils import *
from uds.config import *
from uds.sherpa 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)

srctool_dir="{}/{}".format(outfile_dir,"srctool-products")
create_folder(srctool_dir)

outkey="tm0"

outfile_srctool="{}_SrcTool_".format(outkey)

do_flux_distr         = False
do_sens_curve         = False
do_4xmm_ratio         = False
do_euds_radec_err     = False
do_euds_dr12_diff     = False
do_euds_dr12_stat     = True

do_print_ecf          = False

index=0

catalog = "{}_SourceCatalog_en{}.main.selected.csv".format(os.path.join(outfile_dir,outkey), eband[0])

if not (os.path.isfile(catalog)==True):
    print("{} not found, run 05_srctool.py?".format(catalog))
    sys.exit()

if(do_flux_distr==True):

    
    data, logbins, mean, median = get_log_distr(infile=catalog, field='ml_rate', minval=1e-3, maxval=2)
    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.hist(data, bins=logbins, histtype='step', color='blue', linewidth=1, linestyle='solid')
    plt.xlabel('Count rate (counts s$^{-1}$)',fontsize=14, fontweight='normal')
    plt.ylabel('Number',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('log')
    plt.yscale('log')
    plt.savefig(catalog.replace("main.selected.csv", "ml_rate.png"), bbox_inches='tight')
    plt.close(fig)
    
    
    data, logbins, mean, median = get_log_distr(infile=catalog, field='ml_flux', minval=1e-15, maxval=2e-12)
    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.hist(data, bins=logbins, histtype='step', color='blue', linewidth=1, linestyle='solid')
    plt.xlabel('Energy flux (erg s$^{-1}$ cm$^{-2}$)',fontsize=14, fontweight='normal')
    plt.ylabel('Number',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('log')
    plt.yscale('log')
    plt.savefig(catalog.replace("main.selected.csv", "ml_flux.png"), bbox_inches='tight')
    plt.close(fig)

if(do_sens_curve==True):
    coeff = 2.4336e-13 / 3.4012e-13 # see below

    areatab="{}_AreaTable_dl10_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
    hdul = fits.open(areatab)
    tbdata = hdul[1].data
    """ convert limflux from 0.3-2.3 keV to 0.5-2 keV """
    limflux_dl10 = tbdata['LIMFLUX']*coeff
    area_dl10 = tbdata['SKY_AREA']
    hdul.close()

    areatab="{}_AreaTable_dl6_en{}{}".format(os.path.join(outfile_dir,outkey), eband[index], outfile_post)
    hdul = fits.open(areatab)
    tbdata = hdul[1].data
    """ convert limflux from 0.3-2.3 keV to 0.5-2 keV """
    limflux_dl6 = tbdata['LIMFLUX']*coeff
    area_dl6 = tbdata['SKY_AREA']
    hdul.close()
    
    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    ax.set_xlim([3e-16,1e-13])
    ax.set_ylim([0.3,160])

    plt.plot(limflux_dl10,area_dl10, color="black", linestyle='solid',label="eUDS (DL10)")
    plt.plot(limflux_dl6,area_dl6, color="black", linestyle='dashed', label="eUDS (DL6)")

    df = pandas.read_csv("../data/surveys/eFEDS.dat",header=None,names=['limflux','area'])
    plt.plot(df['limflux'],df['area'], color="brown", linestyle='solid',label="eFEDS (DL6)")


    df = pandas.read_csv("../data/surveys/cosmos-legacy.dat",header=None,names=['limflux','area'])
    plt.plot(df['limflux'],df['area'], color="blue", linestyle='solid',label="COSMOS Legacy")

    df = pandas.read_csv("../data/surveys/CDWFS.dat",header=None,names=['limflux','area'])
    plt.plot(df['limflux'],df['area'], color="green", linestyle='solid', label="CDWFS")

    df = pandas.read_csv("../data/surveys/XMM-RM.dat",header=None,names=['limflux','area'])
    plt.plot(df['limflux'],df['area'], color="magenta", linestyle='solid',label="XMM-RM")

    df = pandas.read_csv("../data/surveys/XMM-XXL-N.dat",header=None,names=['limflux','area'])
    plt.plot(df['limflux'],df['area'], color="red", linestyle='solid',label="XMM-XXL-N")

 
    ax.legend()

    #plt.hist(data, bins=logbins, histtype='step', color='blue', linewidth=1, linestyle='solid')
    plt.xlabel('Limiting flux (0.5-2 keV, erg s$^{-1}$ cm$^{-2}$)',fontsize=14, fontweight='normal')
    plt.ylabel('Area (deg$^{2}$)',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('log')
    plt.yscale('log')
    png="{}_AreaTable_en{}.png".format(os.path.join(outfile_dir,outkey), eband[index])
    plt.savefig(png, bbox_inches='tight')
    plt.close(fig)

    """
    ========================================================================
    Model TBabs<1>*powerlaw<2> Source No.: 1   Active/On
    Model Model Component  Parameter  Unit     Value
    par  comp
    1    1   TBabs      nH         10^22    2.00000E-02  frozen
    2    2   powerlaw   PhoIndex            2.00000      frozen
    3    2   powerlaw   norm                1.14851E-04  +/-  3.83988E-06  
    ________________________________________________________________________
    
    Model Flux 0.00028334 photons (3.4012e-13 ergs/cm^2/s) range (0.30000 - 2.3000 keV)
    Model Flux 0.0001619 photons (2.4336e-13 ergs/cm^2/s) range (0.50000 - 2.0000 keV)
    """
if(do_4xmm_ratio==True):
    filename="../products/eUDS_4XMM-DR12.flux.csv"

    data, logbins, mean, median = get_log_distr(infile=filename, field='ratio', minval=8e-2, maxval=60, nbin=60)
    print("Median {}".format(median))
    print("  Mean {}".format(mean))
    print("Ntotal {}".format(len(data)))
    fig, ax = plt.subplots()
    #plt.figure(figsize=(5,5))
    #plt.figure().set_figheight(3.6)
    #plt.rcParams['figure.figsize'] = [4, 4]
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.hist(data, bins=logbins, histtype='step', color='blue', linewidth=1, linestyle='solid')
    plt.axvline(x = median, color = 'black', label = 'Median value', linestyle='dashed')
    plt.xlabel('Energy flux ratio',fontsize=14, fontweight='normal')
    plt.ylabel('Number',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('log')
    plt.yscale('log')
    plt.savefig(filename.replace(".csv", ".distr.png"), bbox_inches='tight')
    plt.close(fig)

    df = pandas.read_csv(filename)
    fig, ax = plt.subplots()
    #plt.rcParams['figure.figsize'] = [4, 4]
    #plt.figure(figsize=(5,5))
    #plt.figure().set_figheight(4)
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.plot(np.array([1e-17,1e-11]), np.array([1e-17,1e-11]), linestyle = 'solid', color='black')
    plt.plot(np.array([1e-17,1e-11]), 10*np.array([1e-17,1e-11]), linestyle = 'dashed', color='black')
    plt.plot(np.array([1e-17,1e-11]), 0.1*np.array([1e-17,1e-11]), linestyle = 'dotted', color='black')
    #plt.plot(np.array([1e-17,1e-11]), median*np.array([1e-17,1e-11]), linestyle = 'dashed', color='blue')
    plt.errorbar(df['dr12_flux'],df['euds_flux'], yerr=df['euds_flux_err'], xerr=df['dr12_flux_err'], linestyle='None', color='black', label='Errors')
    plt.plot(df['dr12_flux'],df['euds_flux'], marker="o",  linewidth=1, linestyle='None', markerfacecolor='Gold',markeredgecolor="black",)
    #plt.axvline(x = median, color = 'black', label = 'Median value', linestyle='dashed')
    plt.xlabel('4XMM-DR12 Energy flux (erg s$^{-1}$ cm$^{-2}$)',fontsize=14, fontweight='normal')
    plt.ylabel('eUDS Energy flux (erg s$^{-1}$ cm$^{-2}$)',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xlim(2e-16,1e-12)
    plt.ylim(1e-15,2e-12)
    plt.xscale('log')
    plt.yscale('log')
    plt.savefig(filename.replace(".csv", ".png"), bbox_inches='tight')
    plt.close(fig)
    
if(do_euds_radec_err==True):
    filename="../products/eUDS.fits"

    ecat={}
    hdul = fits.open(filename)
    etab = hdul[1].data
    hdul.close()

    x=[]
    y=[]
    for s in etab:
        if ("XMM" in s['DR12_IAU_NAME']):
            print("Skip ",s['DR12_IAU_NAME'])
            continue
        if (s['EXT_LIKE']>0.0):
            print("Skip extended ",s['ID_SRC'])
            continue
        x.append(s['DET_LIKE'])
        y.append(s['RADEC_ERR'])
        
            
    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    ax.set_xlim([10,10000])
    ax.set_ylim([0.1,40])

    plt.plot(x,y,".", color="gold", markersize=12,  markeredgewidth=1.5, markeredgecolor="black")

    plt.xlabel('DET_LIKE',fontsize=14, fontweight='normal')
    plt.ylabel('RADEC_ERR (arcsec)',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('log')
    plt.yscale('log')
    png="../products/en0_radec_err.png"
    plt.savefig(png, bbox_inches='tight')
    plt.close(fig)

    stat={}
    hdul = fits.open("../products/eUDS.dr12.cross.fits")
    tab = hdul[1].data
    hdul.close()
    
    print(len(np.unique(tab['ID'])))
    for rec in tab:
        sid=rec['ID']
        if sid in stat.keys():
            stat[sid]=stat[sid]+1
        else:
            stat[sid]=1

    unique=0
    double=0
    triple=0
    for key in stat.keys():
        if(stat[key]==1):
            unique=unique+1
        if(stat[key]==2):
            double=double+1
        if(stat[key]==3):
            triple=triple+1
    print("unique: {}, double: {}. triple: {}".format(unique, double, triple))
    

if(do_euds_dr12_diff==True):
    # read forced first
    # fieldnames = ['euds_det_like', 'euds_flux', 'euds_flux_err', 'dr12_flux', 'dr12_flux_err', 'ratio']
    filename="../products/eUDS_4XMM-DR12.flux.csv"
    df = pandas.read_csv(filename)
    # calculate difference similar to Bruner

    threshold=5e-14

    # forced catalog contains CONF flag, take it to remove from histogram
    hdul = fits.open("../products/eUDS_4XMM-DR12.fits")
    forced = hdul[1].data
    fcat={}
    for rec in forced:
        key=rec['DR12_SRCID']
        fcat[key]={'conf':rec['CONF'],}
            
    # read cross-match results (not forced)
    hdul = fits.open('../products/eUDS.dr12.cross.fits')
    tb = hdul[1].data
    cross_diff=[]
    cross_ratio=[]
    mean=0.0
    count=0
    for ind,s in enumerate(tb):
        key=rec['DR12_SRCID']
        #if not (tb[ind]['dr12_flux']>threshold and tb[ind]['src_flux']>threshold):
        if not (tb[ind]['dr12_flux']>threshold):
            continue
        ape_flux=(tb[ind]['ape_cts']-tb[ind]['ape_bkg'])/tb[ind]['ape_exp']/ecf[index]
        
        cross_dr12_flux=(tb[ind]['dr12_flux'])
        cross_dr12_flux_error=(tb[ind]['dr12_flux_error'])
        cross_euds_flux=(tb[ind]['src_flux'])
        #cross_euds_flux=ape_flux
        cross_euds_flux_error=(tb[ind]['src_flux_error'])

        print(cross_euds_flux,cross_euds_flux_error,tb[ind]['ape_cts'],tb[ind]['det_like'])
        
        d=(cross_dr12_flux - cross_euds_flux) / np.sqrt(cross_dr12_flux_error**2 + cross_euds_flux_error**2)
        #if(d < -15):
        #    continue
        r=cross_euds_flux/cross_dr12_flux
        cross_diff.append(d)
        cross_ratio.append(r)

        #if(d < -15.0):
        #print("GGG",d,r,cross_euds_flux,cross_dr12_flux)
        #print(tb[ind]['DR12_NAME'],tb[ind]['DR12_RA'],tb[ind]['DR12_DEC'])
        mean=mean + (cross_euds_flux/cross_dr12_flux)
        count=count+1
        
    
   
    print("min={:.2f}, max={:.2f}, mean={:.2f}, median={:.2f}, std={:.2f} N={}".format(min(cross_diff),
                                                                                       max(cross_diff),
                                                                                       np.mean(cross_diff),
                                                                                       np.median(cross_diff),
                                                                                       np.std(cross_diff),
                                                                                       len(cross_diff)))
    print("ratio mean",mean/count)
    print("ratio median",np.median(cross_ratio))


    
    nbin=50
    bins=np.linspace(-30,30, nbin)

    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.hist(cross_diff, bins=bins,histtype='step', color='green', linewidth=1, linestyle='solid', density=True)

    #x = np.arange(-10, 10, 0.001)
    #plot normal distribution with mean 0 and standard deviation 1
    #plt.plot(x, norm.pdf(x, 0, 1), color='red', linewidth=2)
    
    plt.ylabel('Relative fraction',fontsize=14, fontweight='normal')
    plt.xlabel('(F$_{XMM}$-F$_{eUDS}$)/$\sqrt{\Delta F_{XMM}^{2}+\Delta F_{eUDS}^{2}}$',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('linear')
    plt.yscale('linear')
    ax.set_xlim([-31, 31])
    ax.set_ylim([0, 0.2])
    plt.savefig("../products/cross-match_diff.png", bbox_inches='tight')
    plt.close(fig)

    
    nbin=200
    bins=np.linspace(-1,1.0, nbin)

    fig, ax = plt.subplots()
    for axis in ['top','bottom','left','right']:
        ax.spines[axis].set_linewidth(1)
    ax.tick_params(axis="both", width=1, labelsize=14)
    plt.hist(cross_ratio, bins=bins, histtype='step', color='green', linewidth=1, linestyle='solid', density=False)

    #x = np.arange(-10, 10, 0.001)
    #plot normal distribution with mean 0 and standard deviation 1
    #plt.plot(x, norm.pdf(x, 0, 1), color='red', linewidth=2)
    
    plt.xlabel('eUDS and 4XMM-DR12 flux ratio',fontsize=14, fontweight='normal')
    plt.ylabel('Relative fraction',fontsize=14, fontweight='normal')
    plt.grid(visible=True)
    plt.xscale('linear')
    plt.yscale('linear')
    plt.savefig("../products/cross-match_ratio.png", bbox_inches='tight')
    plt.close(fig)
    

if(do_print_ecf==True):
    filename='../data/ECF/ecf_tbabspow_g2nh0.02.pkl'
    with open(filename, 'rb') as f:
        ecf_table = pickle.load(f)
    """    
    for key in table.keys():
        print("{} --> {}".format(key,table[key]))
    """
    
    print(ecf_table[(0.3,2.3)])
    print(ecf_table[(0.3,0.6)])
    print(ecf_table[(0.6,2.3)])
    print(ecf_table[(2.3,5.0)])
    print(ecf_table[(5.0,8.0)])
    print()
    print(ecf_table[(0.5,1.0)]) # 4XMM-DR12 EP2 band
    print(ecf_table[(1.0,2.0)]) # 4XMM-DR12 EP3 band

if(do_euds_dr12_stat==True):
    dr12_stat(infile='../products/eUDS_4XMM-DR12.fits',
              xmmslim='../data/4XMM-DR12/4XMM_DR12cat_slim_v1.0_UDS.fits.catalog',
              xmmlim=None,
              outfile_reg='../products/dr12_fluxlim.reg')
