#!/usr/bin/env python3

import numpy as np
import tensorflow.keras as keras

from tensorflow.keras.layers import Input

## Basic blocks ##

def skip_connection(x, xskip):
    """
    A long skip connection that concatenates output from an encoding layer to a
    layer in the decoder along the channel dimension.

    Parameters
    ----------
    x :
        Input tensor from the decoder phase.
    xskip :
        Tensor from the encoder phase.
    """

    con = keras.layers.Concatenate(axis=3)([x, xskip])
    return con

def convolution_block(x, filters, kernel_size, padding, strides, act, dropval, name=None, first_block=False, upsample=False):
    """
    Fundamental convolution block.

    This block has three possible configurations.
    1) conv -> activation -> batchnorm
    2) dropout -> conv -> activation -> batchnorm
    3) upsample -> dropout -> conv -> activation -> batchnorm

    This convolution block is used in both the encoder and decoder of the network.
    Upsampling in the decoder is done with a dedicated layer. Downsampling is done
    with a convolutional layer.

    Parameters
    ----------
    x :
        Input tensor.
    filters : int
        Number of activation maps in the convolutiional layer.
    kernel_size : int
        Dimension of filter in convolutional layer. Only a single integer is accepted.
    padding : str
        Convolutional layer padding.
    strides : int
        Stride number for convolutional layer.
    act :
        Activation function
    dropval :
        Drop rate for droupout layer.
    name : str
        Name that will be applied to only the batch normalization layer.
    first_block : boolean
        If True, dropout layer will be omitted from the block.
    upsample : boolean
        If True, a 2D upsampling layer will be placed before the dropout layer.
        This enables configuration (3).

    Returns
    -------
    con :
        Output tensor.
    """

    if first_block:
        con = keras.layers.Conv2D(filters, kernel_size=kernel_size, padding=padding, strides=strides)(x)
    elif upsample:
        con = keras.layers.UpSampling2D(strides)(x)
        con = keras.layers.Dropout(dropval)(con)
        con = keras.layers.Conv2D(filters, kernel_size=kernel_size, padding=padding, strides=1)(con)
    else:
        con = keras.layers.Dropout(dropval)(x)
        con = keras.layers.Conv2D(filters, kernel_size=kernel_size, padding=padding, strides=strides)(con)
    con = keras.layers.Activation(act)(con)
    con = keras.layers.BatchNormalization(name=name)(con)
    return con

def residual_block(x, filters, kernel_size, padding, strides, act, dropval, name=[None, None], first_block=False, upsample=[False, False, False], filters_changed=False, mid_skip=None):
    """
    Parameters
    ----------
    x :
        Input tensor.
    filters : list
        A list with two elements. The items in the list, filters[0] and filters[1],
        are the number of filters in the first and second convolution block,
        respectively.
    kernel_size : int
        Dimension of filter in convolutional layer. Only a single integer is accepted.
        The kernel_size provided is applied to both convolutional blocks.
    padding : str
        Convolutional layer padding. Given string is applied to both convolution
        blocks.
    strides : list
        A list with three elements. Two items in the list, strides[0] and strides[1],
        are the stride values for the first and second convolutional block, respectively.
        The third item, stride[2], is the stride that will be applied to the convolutional
        layer or upsampling layer in the residual connection.
    act : list
        A list with two elements. The items in the list, act[0] and act[1], are the
        activation functions for the first and second convolutional block, respectively.
    dropval :
        Drop rate for droupout layer. Given value is applied to both convolution
        blocks.
    name : list
        A list with two elements. The items in the list, name[0] and name[1], are the
        names that will be applied to only the batch normalization layer of the first
        and second convolution block, respectively.
    first_block : boolean
        If True, dropout layer will be omitted from the first convolution block.
    upsample : list
        A list with three elements. Two items in the list, upsample[0] and upsample[1],
        indicate if an 2D upsampling layer should be added before the dropout layer in
        the first and second convolution block, respectively. The third item, upsample[2],
        indicates if a 2D upsampling layer should be used. All elements in the list
        must be boolean.
    filters_changed : boolean
        Set to True if the input to the residual block and the output have a different
        number of channels (filters). If True, a convolutional layer, if not already present
        due to automatic triggers, will be placed in the residual connection to fix the channel
        dimension.
    mid_skip :
        A keras tensor. This keras tensor will be concatenated with the output of the
        first convolution block.

    Returns
    -------
    con0 :
        Output tensor of the first convolution block.
    con1 :
        Output tensor of the second convolution block. This is before the residual
        connection is applied.
    resoutput :
        Output tensor from running the entire residual block. Residual connection is applied.
    """

    con0 = convolution_block(x, filters[0], kernel_size, padding, strides[0], act[0], dropval, name=name[0], first_block=first_block, upsample=upsample[0])
    if mid_skip is not None:
        con0 = skip_connection(con0, mid_skip)
    con1 = convolution_block(con0, filters[1], kernel_size, padding, strides[1], act[1], dropval, name=name[1], upsample=upsample[1])

    if strides[2] != 1 and upsample[2] == True:
        x = keras.layers.UpSampling2D(strides[2])(x)
        if filters_changed == True:
            x = keras.layers.Conv2D(filters[1], kernel_size=1, padding=padding, strides=1)(x)
    elif (strides[2] != 1 and upsample[2] == False):
        x = keras.layers.Conv2D(filters[1], kernel_size=1, padding=padding, strides=strides[2])(x)
    elif filters_changed == True:
        x = keras.layers.Conv2D(filters[1], kernel_size=1, padding=padding, strides=1)(x)
    rescon = keras.layers.BatchNormalization()(x)

    resoutput = keras.layers.Add()([rescon, con1])
    return con0, con1, resoutput

def ResUNet_CMB(params):
    """
    ResUNet-CMB Network

    Network used in "Reconstructing Patchy Reionization with Deep Learning."

    Parameters
    ----------
    params:
        params is a container for the variables defined in the configuration file.
        An instance of class resunet.utils.Params.

    Returns
    -------
    model :
        model object.
    """

    input_img1 = Input(shape=(params.imagesize, params.imagesize, 1), dtype=np.float32, name="qlen")
    input_img2 = Input(shape=(params.imagesize, params.imagesize, 1), dtype=np.float32, name="ulen")

    # encoder
    enc_0 = keras.layers.Concatenate(axis=3)([input_img1, input_img2])
    enc_1 = residual_block(enc_0, [64,64], 5, "same", [1,1,1], ["selu","selu"], params.dropval, first_block=True, filters_changed=True)[2]
    enc_2 = residual_block(enc_1, [64,128], 5, "same", [1,2,2], ["selu","selu"], params.dropval, filters_changed=True)
    enc_3 = residual_block(enc_2[2], [128,128], 5, "same", [1,1,1], ["selu","selu"], params.dropval)[2]

    # bridge between encoder and decoder
    enc_dec = residual_block(enc_3, [256,128], 5, "same", [2,2,1], ["selu","selu"], params.dropval, upsample=[False, True, False])[2]

    lskip1 = skip_connection(enc_dec, enc_3)
    dec_1 = residual_block(lskip1, [128,128], 5, "same", [1,1,1], ["selu","selu"], params.dropval, filters_changed=True)[2]
    dec_2 = residual_block(dec_1, [64, 64], 5, "same", [2,1,2], ["selu","selu"], params.dropval, upsample=[True, False, True], mid_skip=enc_2[0], filters_changed=True)[2]

    # Block all branches use for final residual connection
    dec_3 = convolution_block(dec_2, 64, 5, "same", 1, "selu", params.dropval)

    # kappa branch
    kappa_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    kappa_res_1 = keras.layers.BatchNormalization()(dec_2)
    kappa_1 = keras.layers.Add()([kappa_res_1, kappa_1])
    kappa_2 = convolution_block(kappa_1, 64, 5, "same", 1, "selu", params.dropval)
    kappa_3 = convolution_block(kappa_2, 1, 5, "same", 1, "linear", params.dropval, name="kappa")

    # primordial E branch
    unle_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    unle_res_1 = keras.layers.BatchNormalization()(dec_2)
    unle_1 = keras.layers.Add()([unle_res_1, unle_1])
    unle_2 = convolution_block(unle_1, 64, 5, "same", 1, "selu", params.dropval)
    unle_3 = convolution_block(unle_2, 1, 5, "same", 1, "linear", params.dropval, name="unle")

    # tau branch
    tau_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    tau_res_1 = keras.layers.BatchNormalization()(dec_2)
    tau_1 = keras.layers.Add()([tau_res_1, tau_1])
    tau_2 = convolution_block(tau_1, 64, 5, "same", 1, "selu", params.dropval)
    tau_3 = convolution_block(tau_2, 1, 5, "same", 1, "linear", params.dropval, name="tau")

    model = keras.models.Model(inputs=[input_img1, input_img2], outputs=[tau_3, unle_3, kappa_3])
    return model

def ResUNet_CMB_4out(params):
    """
    ResUNet-CMB 4-output Network

    Network used in "Reconstructing Cosmic Polarization Rotation with ResUNet-CMB"

    Parameters
    ----------
    params:
        params is a container for the variables defined in the configuration file.
        An instance of class resunet.utils.Params.

    Returns
    -------
    model :
        model object.
    """

    input_img1 = Input(shape=(params.imagesize, params.imagesize, 1), dtype=np.float32, name="qlen")
    input_img2 = Input(shape=(params.imagesize, params.imagesize, 1), dtype=np.float32, name="ulen")

    # encoder
    enc_0 = keras.layers.Concatenate(axis=3)([input_img1, input_img2])
    enc_1 = residual_block(enc_0, [64,64], 5, "same", [1,1,1], ["selu","selu"], params.dropval, first_block=True, filters_changed=True)[2]
    enc_2 = residual_block(enc_1, [64,128], 5, "same", [1,2,2], ["selu","selu"], params.dropval, filters_changed=True)
    enc_3 = residual_block(enc_2[2], [128,128], 5, "same", [1,1,1], ["selu","selu"], params.dropval)[2]

    # bridge between encoder and decoder
    enc_dec = residual_block(enc_3, [256,128], 5, "same", [2,2,1], ["selu","selu"], params.dropval, upsample=[False, True, False])[2]

    lskip1 = skip_connection(enc_dec, enc_3)
    dec_1 = residual_block(lskip1, [128,128], 5, "same", [1,1,1], ["selu","selu"], params.dropval, filters_changed=True)[2]
    dec_2 = residual_block(dec_1, [64, 64], 5, "same", [2,1,2], ["selu","selu"], params.dropval, upsample=[True, False, True], mid_skip=enc_2[0], filters_changed=True)[2]

    # Block all branches use for final residual connection
    dec_3 = convolution_block(dec_2, 64, 5, "same", 1, "selu", params.dropval)

    # kappa branch
    kappa_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    kappa_res_1 = keras.layers.BatchNormalization()(dec_2)
    kappa_1 = keras.layers.Add()([kappa_res_1, kappa_1])
    kappa_2 = convolution_block(kappa_1, 64, 5, "same", 1, "selu", params.dropval)
    kappa_3 = convolution_block(kappa_2, 1, 5, "same", 1, "linear", params.dropval, name="kappa")

    # primordial E branch
    unle_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    unle_res_1 = keras.layers.BatchNormalization()(dec_2)
    unle_1 = keras.layers.Add()([unle_res_1, unle_1])
    unle_2 = convolution_block(unle_1, 64, 5, "same", 1, "selu", params.dropval)
    unle_3 = convolution_block(unle_2, 1, 5, "same", 1, "linear", params.dropval, name="unle")

    # tau branch
    tau_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    tau_res_1 = keras.layers.BatchNormalization()(dec_2)
    tau_1 = keras.layers.Add()([tau_res_1, tau_1])
    tau_2 = convolution_block(tau_1, 64, 5, "same", 1, "selu", params.dropval)
    tau_3 = convolution_block(tau_2, 1, 5, "same", 1, "linear", params.dropval, name="tau")

    # alpha branch
    cbf_1 = convolution_block(dec_3, 64, 5, "same", 1, "selu", params.dropval)
    cbf_res_1 = keras.layers.BatchNormalization()(dec_2)
    cbf_1 = keras.layers.Add()([cbf_res_1, cbf_1])
    cbf_2 = convolution_block(cbf_1, 64, 5, "same", 1, "selu", params.dropval)
    cbf_3 = convolution_block(cbf_2, 1, 5, "same", 1, "linear", params.dropval, name="cbf")

    model = keras.models.Model(inputs=[input_img1, input_img2], outputs=[tau_3, unle_3, kappa_3, cbf_3])
    return model


#   ZEROPOINT CALCULATION
#   19.03.09    MERRY CHRISTMAS!
#   MADE BY GREGORY S.H. PAEK
#	UPDATE : 20.01.16
#=========================================================================#
# import os, glob, subprocess
import numpy as np
from imsng import tool_tbd
# from astropy.table import Table
from astropy.table import vstack
import matplotlib.pyplot as plt
# from numba import jit
def which_obs(obs, path_and_file):
	"""
	=====================================================================
	GIVE CCD INFORMATION
	---------------------------------------------------------------------
	INPUT   :   observatory name, and path+table_name(ascii format)
				i.e.) path_and_file='/home/sonic/Research/table/obs.txt'
	OUTPUT  :   GAIN        []
				Pixel Scale ["/pix]
	---------------------------------------------------------------------
	obs         ccd         gain    RDnoise     dark    pixelscale
	SOAO        FLI         1.43    13.60       0.0     0.44540
	LOAO        E2V         2.68    4.84        0.0     0.794
	LOAO_FLI    FLI         3       15.0        0.0     1.28
	DOAO_sophia PIsophia    0.928   7.63        0.0     0.3855
	DOAO        FLI         1.27    14.0        0.005   0.464
	oldBOAO     UNKNOWN     8.9     20.0        0.0     1.7
	BOAO        UNKNOWN     5.0     2.0         0.0     0.2145
	BOAO        UNKNOWN     5.0     5.0         0.0     0.2145
	SAO         SBIG        1.35    9.0         0.0     0.31
	30inch      UNKNOWN     1.60    5.87        0.0     1.3553
	CCA250      MLI16803    1.46    10.35       0.0     2.06
	MAIDANAK    SNU4kCAM    1.45    4.7         0.0     0.266
	MAIDANAK    UNKNOWN     5.40    9.0         0.0     0.266
	LSGT        SNUCAMII    1.15    6.0         0.2     0.92
	UKIRT       WFCAM       99.0    0.0         0.0     0.4
	=====================================================================
	"""
	import numpy as np
	from astropy.io import ascii
	obsinfo     = ascii.read(path_and_file)
	indx_obs    = np.where( obs == obsinfo['obs'] )
	gain        = obsinfo['gain'][indx_obs]
	pixscale    = obsinfo['pixelscale'][indx_obs]
	return gain, pixscale

#-------------------------------------------------------------------------#
def image_list(imlist):
	"""
	INPUT   :   imlist = glob.glob('Calib-*.fits')
	OUTPUT  :   observatory list
				object list
				fillter list
	"""
	obslist = []
	objlist = []
	fillist = []
	for img in imlist:
		sp  = img.split('-')
		obslist.append(sp[1])
		objlist.append(sp[2])
		fillist.append(sp[5])
	obslist = list(set(obslist))
	objlist = list(set(objlist))
	fillist = list(set(fillist))
	return obslist, objlist, fillist

#-------------------------------------------------------------------------#
def secom(inim, gain, pixscale, zp=0, seeing=3, det_sigma=3, backsize=str(64), backfiltersize=str(3), psf=False, dual=False, detect='detection.fits', check=False):
	"""
	SourceEXtractor
	APERTURE    3", 5", 7",
				1.0seeing, 1.2seeing ,1.5seeing ,1.7seeing ,2.0seeeing
	INPUT   :   (image).fits
				aperture    []
				seeing_fwhm [pixel]
	OUTPUT  :   no return
				.cat
	"""
	import numpy as np
	import os
	from astropy.io import ascii
	#   FILE CHECK
	#	CONFIG FILES (USER BASE PATH)
	configfile      = '/home/sonic/Research/yourpy/config/targetphot.sex'
	paramfile       = '/home/sonic/Research/yourpy/config/targetphot.param'
	nnwfile		    = '/home/sonic/Research/yourpy/config/targetphot.nnw'
	convfile	    = '/home/sonic/Research/yourpy/config/targetphot.conv'
	try:
		comment = 'SourceEXtractor START\n' \
				+ 'IMAGE\t\t: '+inim+'\n' \
				+ 'GAIN\t\t: '+str(gain)+'\n' \
				+ 'PIXSCALE\t: '+str(pixscale)+'\n' \
				+ 'DETECTION SIGMA\t: '+str(det_sigma)+'\n' \
				+ 'PARAM\t\t: '+paramfile+'\n' \
				+ 'BACKSIZE\t: '+backsize+'\n' \
				+ 'BACKFILTER\t: '+backfiltersize+'\n' \
				+ 'CONFIG\t\t: '+configfile+'\n' \
				+ 'NNW\t\t: '+nnwfile+'\n' \
				+ 'CONVOLVE\t: '+convfile
		print(comment)
	except:
		comment = 'CHECK configfile/paramfile/nnewfile/convfile or others.'
		print(comment)
	#   FILE NAME
	cat     = inim[:-5]+'.cat'
	seg     = inim[:-5]+'.seg.fits'
	bkg     = inim[:-5]+'.bkg.fits'
	sub     = inim[:-5]+'.sub.fits'
	impsf   = inim[:-5]+'.psf'
	aper    = inim[:-5]+'.aper.fits'

	#   BASIC INFO.
	det_area        = 5
	det_thresh      = det_sigma/np.sqrt(det_area)
	detecminarea    = str(det_area)
	detectthresh    = str(det_thresh)
	#   OPTION
	aperture        = '%.2f'%(3./pixscale)+','+'%.2f'%(5./pixscale)+','+'%.2f'%(7./pixscale)+','+'%.2f'%(seeing)+','+'%.2f'%(1.2*seeing)+','+'%.2f'%(1.5*seeing)+','+'%.2f'%(1.7*seeing)+','+'%.2f'%(2.0*seeing)
	option0	= ' -MAG_ZEROPOINT {} '.format(zp)
	option1 = ' -CATALOG_NAME '+cat+' -PARAMETERS_NAME '+paramfile
	option2 = ' -DETECT_MINAREA '+detecminarea+' -DETECT_THRESH '+detectthresh \
			+' -FILTER Y '+'-FILTER_NAME '+convfile
	option3 = ' -PHOT_APERTURES '+aperture \
			+' -GAIN '+'%.1f'%(gain)+' -PIXEL_SCALE '+'%.1f'%(pixscale)
	#option4 = ' -SEEING_FWHM '+'%.3f'%(seeing)+' -STARNNW_NAME '+nnwfile
	option4 = ' -SEEING_FWHM '+'%.3f'%(seeing)+' -STARNNW_NAME '+nnwfile
	option5 = ' -BACK_SIZE '+ backsize \
			+ ' -BACK_FILTERSIZE '+ backfiltersize+' -BACKPHOTO_TYPE LOCAL'
	option6 = ' -CHECKIMAGE_TYPE SEGMENTATION,APERTURES,BACKGROUND,-BACKGROUND'
	option7 = ' -CHECKIMAGE_NAME '+seg+','+aper+','+','+bkg+','+sub
	if psf	!= False:
		option8 = ' -PSF_NAME '+impsf
	else:
		option8 = ''
	#   COMMAND
	#   detect = detection.fits is fine image show target significantly and have good seeing and many stars i
	dualcom ='sex -c '+configfile+' '+detect+' , '+inim+' -CATALOG_NAME dual'+cat+' -PARAMETERS_NAME '+paramfile+ ' '+option0+' '+option2+' '+option3+' '+option4+' '+option5+' '+option6+' '+option7+' '+option8
	sglcom  ='sex -c '+configfile+' '+inim+' '+option0+' '+option1+' '+option2+' '+option3+' '+option4+' '+option5+' '+option6+' '+option7+' '+option8
	clearcom='sex -c '+configfile+' '+inim+' '+option0+' '+option1+' '+option2+' '+option3+' '+option4+' '+option5+' '+option8
	if dual == False    :
		if check == False:
			os.system(clearcom)
		else:
			os.system(sglcom)
	else                : os.system(dualcom)

	secat   = ascii.read(cat)
	return secat, cat
#-------------------------------------------------------------------------#
def psfex(inim, pixscale):
	"""
	PSfextractor
	INPUT   :   (image).fits
	OUTPUT  :   FWHM    [pixel]
				FWHM    [arcsec]
	"""
	import os

	#   FILE CHECK
	#	CONFIG FILES (USER BASE PATH)
	psfexconf_prese_conf    = '/home/sonic/Research/yourpy/config/prepsfex.sex'
	psfexconf_prese_param   = '/home/sonic/Research/yourpy/config/prepsfex.param'
	psfexconf_psfex_conf    = '/home/sonic/Research/yourpy/config/default.psfex'
	psfexconf_psfex_conv    = '/home/sonic/Research/yourpy/config/default.conv'
	try:
		comment = '\nPSFex START\n' \
				+ 'IMAGE\t\t: '+inim+'\n' \
				+ 'PRE_CONFIG\t: '+psfexconf_prese_conf+'\n' \
				+ 'PRE_PARAM\t: '+psfexconf_prese_param+'\n' \
				+ 'CONFIG\t\t: '+psfexconf_psfex_conf+'\n' \
				+ 'CONV\t\t: '+psfexconf_psfex_conv
		print(comment)
	except:
		comment = 'CHECK psfexconf_prese/psfexconf_prese_param/psfexconf_psfex_conf/psfexconf_psfex_conv OR OTHERS.'
		print(comment)

	#   FILE NAME
	cat     = inim[:-5]+'.cat'
	xml     = inim[:-5]+'.xml'
	snap    = 'snap_'+inim+'[100:125,100:125]'
	psf     = 'psf-'+inim
	#   OPTION
	presecom1   = psfexconf_prese_conf+" "+inim
	presecom2   = " -CATALOG_NAME "+cat
	presecom3   = " -FILTER_NAME " + psfexconf_psfex_conv + " -PARAMETERS_NAME " + psfexconf_prese_param
	#   COMMAND
	presecom    = "sex -c "+presecom1+presecom2+presecom3
	psfexcom    = "psfex -c "+psfexconf_psfex_conf+" "+cat
	os.system(presecom)
	os.system(psfexcom)
	os.system('cp psfex.xml '+xml)
	#   SNAP IMAGE
	imcopycom   = 'imcopy '+snap+' '+psf
	print(imcopycom);   os.system(imcopycom)
	#   FWHM [pixel], FWHM [arcsec]
	fwhm_pix    = psfexxml(xml)
	fwhm_arcsec = round(fwhm_pix*pixscale, 3)
	comment     = '\n' \
				+ 'FILE NAME'+'\t'+': '+inim+'\n' \
				+ 'FWHM value'+'\t'+': '+str(fwhm_pix)+'\t'+'[pixel]'+'\n' \
				+ '\t'+'\t'+': '+str(fwhm_arcsec)+'\t'+'[arcsec]'+'\n'
	print(comment)
	return fwhm_pix, fwhm_arcsec
#------------------------------------------------------------
def sexcom(inim, param_insex, dualmode=False):
	'''

	'''
	param_sex = dict(	CONF_NAME = 'default.sex',
						#------------------------------
						#	CATALOG
						#------------------------------
						CATALOG_NAME = 'test.cat',
						CATALOG_TYPE = 'ASCII_HEAD',
						PARAMETERS_NAME = 'default.param',
						#------------------------------
						#	EXTRACTION
						#------------------------------
						DETECT_TYPE = 'CCD',
						DETECT_MINAREA = '5',
						DETECT_MAXAREA = '0',
						DETECT_THRESH = '1.5',
						# ANALYSIS_THRESH = 'RELATIVE',
						ANALYSIS_THRESH = '1.5',
						FILTER = 'Y',
						FILTER_NAME = 'default.conv',
						DEBLEND_NTHRESH = '64',
						DEBLEND_MINCONT = '0.0001',
						CLEAN = 'Y',
						CLEAN_PARAM = '1.0',
						MASK_TYPE = 'CORRECT',
						#------------------------------
						#	PHOTOMETRY
						#------------------------------
						# PHOT_APERTURES = '3',
						PHOT_AUTOPARAMS = '2.5,3.5',
						PHOT_PETROPARAMS = '2.0,3.5',
						SATUR_LEVEL  = '50000.0',
						SATUR_KEY = 'SQTURATE',
						MAG_ZEROPOINT = '0.0',
						MAG_GAMMA = '4.0',
						GAIN = '1.0',
						GAIN_KEY = 'GAIN',
						PIXEL_SCALE = '1.0',
						#------------------------------
						#	STAR/GALAXY SEPARATION
						#------------------------------
						SEEING_FWHM = '3.0',
						STARNNW_NAME = 'default.nnw',
						#------------------------------
						#	BACKGROUND
						#------------------------------
						BACK_SIZE = '128',
						BACK_FILTERSIZE = '10',
						BACKPHOTO_TYPE = 'LOCAL',
						#------------------------------
						#	CHECK IMAGE
						#------------------------------
						CHECKIMAGE_TYPE = 'NONE',
						CHECKIMAGE_NAME = 'check.fits',
						#==============================
						#	MEMORY & MISCELLANEOUS
						#==============================
						MEMORY_OBJSTACK = '3000',
						MEMORY_PIXSTACK = '300000',
						MEMORY_BUFSIZE = '1024',
						VERBOSE_TYPE = 'NORMAL',
						HEADER_SUFFIX = '.head',
						WRITE_XML = 'N',
						XML_NAME = 'sex.xml')

	for key in param_insex.keys():
		param_sex[key] = param_insex[key]

	sexcom_normal = 'sex -c {} {} '.format(param_sex['CONF_NAME'], inim)
	sexcom_dual = 'sex -c {} {} '.format(param_sex['CONF_NAME'], inim)
	for key in param_sex.keys():
		if key != 'CONF_NAME':
			sexcom_normal += '-{} {} '.format(key, param_sex[key])

	# print(sexcom_normal)
	# os.system(sexcom_normal)
	return sexcom_normal

def psfexxml(xmlfile):
	"""
	INPUT   :   .xml
	OUTPUT  :   FWHM    [pixel]
	"""
	from astropy.io.votable import parse
	from astropy.table import Table, Column, MaskedColumn
	votable     = parse(xmlfile)
	table       = votable.get_first_table()
	data        = table.array
	#   EXTRACT FWHM [pixel]
	fwhm        = data['FWHM_Mean'][0]
	fwhm        = round(fwhm, 3)
	return fwhm
#-------------------------------------------------------------------------#
#def matching(incat, refcat, sep=2.0):
def matching(intbl, reftbl, inra, indec, refra, refdec, sep=2.0):
	"""
	MATCHING TWO CATALOG WITH RA, Dec COORD. WITH python
	INPUT   :   SE catalog, SDSS catalog file name, sepertation [arcsec]
	OUTPUT  :   MATCED CATALOG FILE & TABLE
	"""
	import numpy as np
	import astropy.units as u
	# from astropy.table import Table, Column
	from astropy.coordinates import SkyCoord
	from astropy.io import ascii

	incoord		= SkyCoord(inra, indec, unit=(u.deg, u.deg))
	refcoord	= SkyCoord(refra, refdec, unit=(u.deg, u.deg))

	#   INDEX FOR REF.TABLE
	indx, d2d, d3d  = incoord.match_to_catalog_sky(refcoord)
	mreftbl			= reftbl[indx]
	mreftbl['sep']	= d2d
	mergetbl		= intbl
	for col in mreftbl.colnames:
		mergetbl[col]	= mreftbl[col]
	indx_sep		= np.where(mergetbl['sep']*3600.<sep)
	mtbl			= mergetbl[indx_sep]
	#mtbl.write(mergename, format='ascii', overwrite=True)
	return mtbl
#-------------------------------------------------------------------------#
def star4zp(intbl, inmagerkey, refmagkey, refmagerkey,
	refmaglower=14., refmagupper=17., refmagerupper=0.05,
	inmagerupper=0.1, flagcut=0, verbose=False, plot=False, plotout='zphist.png'):
	"""
	SELECT STARS FOR USING ZEROPOINT CALCULATION
	INPUT   :   TABLE, IMAGE MAG.ERR KEYWORD, REF.MAG. KEYWORD, REF.MAG.ERR KEYWORD
	OUTPUT  :   NEW TABLE
	"""
	# import numpy as np

	indx_flag = np.where( intbl['FLAGS'] <= flagcut )
	indx_refmag = np.where(	(intbl[refmagkey] <= refmagupper) &
							(intbl[refmagkey] >= refmaglower))
	indx_refmager = np.where( intbl[refmagerkey] <= refmagerupper )
	indx_mager = np.where( intbl[inmagerkey] <= inmagerupper )

	indx_all = np.where((intbl['FLAGS'] <= flagcut) &
						(intbl[refmagkey] < refmagupper) &
						(intbl[refmagkey] > refmaglower) &
						(intbl[refmagerkey] < refmagerupper) &
						(intbl[inmagerkey] < inmagerupper) )

	newtbl  = intbl[indx_all]
	comment = '-'*50+'\n' \
			+ 'ALL\t\t\t\t\t: {}\n'.format(len(intbl)) \
			+ '-'*50+'\n' \
			+ 'FLAG(<={})\t\t\t\t: {}\n'.format(flagcut, len(indx_flag[0])) \
			+ '{} REF. MAGCUT ({}-{})\t\t: {}\n'.format(refmagkey, refmaglower, refmagupper, len(indx_refmag[0])) \
			+ '{} REF. MAGERR CUT < {}\t\t: {}\n'.format(refmagerkey, refmagerupper, len(indx_refmager[0])) \
			+ '{} CUT < {}\t\t: {}\n'.format(inmagerkey, inmagerupper, len(indx_mager[0])) \
			+ '='*50+'\n' \
			+ 'TOTAL #\t\t\t\t\t: {}\n'.format(len(newtbl)) \
			+ '='*50
	if verbose != False:
		print(comment)
	'''
	if plot!=False:
		# import matplotlib.pyplot as plt
		plt.close('all')
		plt.rc('font', family='serif')
		fig = plt.figure()
		x = 1080 / fig.dpi
		y = 1080 / fig.dpi
		fig.set_figwidth(x)
		fig.set_figheight(y)
		#-------------------------------------------------------------
		plt.subplot(221)
		plt.title('FLAGS')

		bins = np.arange(np.min(intbl['FLAGS']), np.max(intbl['FLAGS'])+0.2, 0.2)
		plt.hist(intbl['FLAGS'], bins=bins, color='tomato', alpha=0.5,)
				# label='{}/{}'.format(len(intbl[indx_flag]), len(intbl)))
		plt.hist(intbl['FLAGS'][indx_flag], bins=bins, color='dodgerblue', alpha=0.75,
				label='{}/{}'.format(len(intbl[indx_flag]), len(intbl)))
		plt.minorticks_on()
		plt.grid(color='grey', alpha=0.5, linestyle='--')
		plt.legend()
		#-------------------------------------------------------------
		plt.subplot(222)
		plt.title('{} MAG'.format(refmagkey))
		bins = np.arange(np.min(intbl[refmagkey]), np.max(intbl[refmagkey])+0.1, 0.1)
		plt.hist(intbl[refmagkey], bins=bins, color='tomato', alpha=0.5)

		plt.hist(intbl[refmagkey][indx_refmag], bins=bins, color='dodgerblue', alpha=0.75,
				label='{}/{}'.format(len(intbl[indx_refmag]), len(intbl)))
		plt.axvline(x=refmaglower, linestyle='--', color='k',)
		plt.axvline(x=refmagupper, linestyle='--', color='k',)

		plt.minorticks_on()
		plt.grid(color='grey', alpha=0.5, linestyle='--')
		plt.legend()
		#-------------------------------------------------------------
		plt.subplot(223)
		plt.title('{}'.format(refmagerkey))
		bins = np.arange(np.min(intbl[refmagerkey]), np.max(intbl[refmagerkey])+0.001, 0.001)
		plt.hist(intbl[refmagerkey], bins=bins, color='tomato', alpha=0.5)

		plt.hist(intbl[refmagerkey][indx_refmager], bins=bins, color='dodgerblue', alpha=0.75,
				label='{}/{}'.format(len(intbl[indx_refmager]), len(intbl)))
		plt.axvline(x=refmagerupper, linestyle='--', color='k',)
		plt.minorticks_on()
		plt.grid(color='grey', alpha=0.5, linestyle='--')
		plt.legend()
		#-------------------------------------------------------------
		plt.subplot(224)
		plt.title('{}'.format(inmagerkey))
		bins = np.arange(np.min(intbl[inmagerkey]), np.max(intbl[inmagerkey])+0.01, 0.01)
		plt.hist(intbl[inmagerkey], bins=bins, color='tomato', alpha=0.5)
		plt.hist(intbl[inmagerkey][indx_mager], bins=bins, color='dodgerblue', alpha=0.75,
				label='{}/{}'.format(len(intbl[indx_mager]), len(intbl)))
		plt.axvline(x=inmagerupper, linestyle='--', color='k',)
		# plt.xlim([0.0, 1.0])
		plt.minorticks_on()
		plt.grid(color='grey', alpha=0.5, linestyle='--')
		plt.legend()
		#-------------------------------------------------------------
		plt.tight_layout()
		# plt.savefig(plotout, dpi=300)
		plt.savefig(plotout)
	'''
	return newtbl
#-------------------------------------------------------------------------#
# @jit
def zpcal(intbl, inmagkey, inmagerkey, refmagkey, refmagerkey, sigma=2.0, method='default'):
	"""
	ZERO POINT CALCULATION
	3 SIGMA CLIPPING (MEDIAN)

	import matplotlib.pyplot as plt
	from numpy import median
	import numpy as np
	"""
	from astropy.stats import sigma_clip
	#	REMOVE BLANK ROW (=99)
	indx_avail      = np.where( (intbl[inmagkey] != 99) & (intbl[refmagkey] != 99) )
	intbl           = intbl[indx_avail]
	zplist          = np.copy(intbl[refmagkey] - intbl[inmagkey])
	intbl['zp']		= zplist
	#	SIGMA CLIPPING
	zplist_clip     = sigma_clip(zplist, sigma=sigma, maxiters=None, cenfunc=np.median, copy=False)
	indx_alive      = np.where( zplist_clip.mask == False )
	indx_exile      = np.where( zplist_clip.mask == True )
	#	RE-DEF. ZP LIST AND INDEXING CLIPPED & NON-CLIPPED
	intbl_alive     = intbl[indx_alive]
	intbl_exile     = intbl[indx_exile]
	#	ZP & ZP ERR. CALC.
	if method == 'default':
		zp              = np.median(np.copy(intbl_alive['zp']))
		zper			= np.std(np.copy(intbl_alive['zp']))
	elif method == 'weightedmean':
		print(method)
		mager = sqsum(intbl_alive[inmagerkey], intbl_alive[refmagerkey])
		w0 = 1/mager
		w = w0/np.sum(w0)
		zp = np.sum(w*intbl_alive['zp'])/np.sum(w)
		zper = 1/np.sqrt(np.sum(w))
	return zp, zper, intbl_alive, intbl_exile
#-------------------------------------------------------------------------#
# def zpplot(outname, alltbl, otbl, xtbl, inmagkey, inmagerkey, refmagkey, refmagerkey, zp, zper, refmaglower, refmagupper):
def zpplot(outname, otbl, xtbl, inmagkey, inmagerkey, refmagkey, refmagerkey, zp, zper, refmaglower, refmagupper):
	# import numpy as np
	# import matplotlib.pyplot as plt
	#   FILE NAME
	plt.close('all')
	plt.rc('font', family='serif')
	fig = plt.figure()
	x = 1920 / 2 / fig.dpi
	y = 1080 / 2 / fig.dpi
	# x = 1920 / 4 / fig.dpi
	# y = 1080 / 4 / fig.dpi
	fig.set_figwidth(x)
	fig.set_figheight(y)

	alltbl = vstack([otbl, xtbl])
	# plt.rcParams.update({'font.size': 16})
	plt.axhline(
				zp,
				linewidth=2, linestyle='-', color='gray',
				label= r'ZP={}$\pm${}'.format(round(zp, 3), round(zper, 3))
				)
	plt.fill_between(
						# [np.min(otbl[refmagkey])-0.05, np.max(otbl[refmagkey])+0.05],
						[np.min(alltbl[refmagkey]), np.max(alltbl[refmagkey])],
						zp-zper, zp+zper,
						color='silver', alpha=0.3
					)
	# plt.errorbar(
	# 				alltbl[refmagkey], alltbl['zp'],
	# 				yerr=tool_tbd.sqsum(alltbl[inmagerkey], alltbl[refmagerkey]),
	# 				c='silver', ms=6, marker='.', ls='',
	# 				capsize=2.5, capthick=1,
	# 				label='All stars ({})'.format(len(alltbl)), alpha=0.25,
	# 			)
	plt.errorbar(
					otbl[refmagkey], otbl['zp'],
					yerr=tool_tbd.sqsum(otbl[inmagerkey], otbl[refmagerkey]),
					c='dodgerblue', ms=6, marker='o', ls='',
					capsize=5, capthick=1,
					label='Used stars ({})'.format(len(otbl)), alpha=0.5,
				)
	# plt.scatter( xtbl[refmagkey], xtbl['zp'], color='tomato', s=50, marker='x', linewidth=1, alpha=1.0, label='CLIPPED ({})'.format(len(xtbl)) )
	plt.errorbar(
					xtbl[refmagkey], xtbl['zp'],
					yerr=tool_tbd.sqsum(xtbl[inmagerkey], xtbl[refmagerkey]),
					c='tomato', ms=6, marker='x', ls='',
					capsize=5, capthick=1,
					label='Clipped stars ({})'.format(len(xtbl)), alpha=0.5,
				)
	plt.axhline(
				zp+zper,
				linewidth=2, linestyle='-.', color='gray', alpha=0.5,
				label='Clip upper ({})'.format(len(xtbl[xtbl['zp']>zp+zper])),
				)
	plt.axhline(
				zp-zper,
				linewidth=2, linestyle='--', color='gray', alpha=0.5,
				label='Clip lower ({})'.format(len(xtbl[xtbl['zp']<zp-zper])),
				)
	plt.axvline(x=refmaglower, linestyle='--', color='k', alpha=0.5)
	plt.axvline(x=refmagupper, linestyle='--', color='k', alpha=0.5)
	# plt.xlim(np.min(otbl[refmagkey])-0.05, np.max(otbl[refmagkey])+0.05)
	# plt.ylim(zp-0.5, zp+0.5)
	#	SETTING
	plt.title(outname, {'fontsize': 12})
	plt.gca().invert_yaxis()
	plt.xlabel('REF.MAG.', {'color': 'black', 'fontsize': 20})
	plt.ylabel('ZP [AB]', {'color': 'black', 'fontsize': 20})
	plt.legend(loc='best', prop={'size': 14}, edgecolor=None)
	plt.tight_layout()
	plt.minorticks_on()
	plt.savefig(outname)
	#	PRINT
	print('MAG TYP     : '+inmagkey)
	print('ZP          = '+str(round(zp, 3)))
	print('ZP ERR      = '+str(round(zper, 3)))
	print('STD.NUMB    = '+str(int(len(otbl))))
	print('REJ.NUMB    = '+str(int(len(xtbl))))
	print('CLIP UPPER  = {}'.format(len(xtbl[xtbl['zp']>zp+zper])))
	print('CTIP LOWER  = {}'.format(len(xtbl[xtbl['zp']<zp-zper])))
#-------------------------------------------------------------------------#
def bkgest_mask(inim):
	'''
	'''
	import numpy as np
	from astropy.io import fits
	from photutils import make_source_mask
	from numpy import mean,median
	from astropy.stats import sigma_clipped_stats

	data    =fits.getdata(inim)
	# mask    = make_source_mask(data, snr=3, npixels=5, dilate_size=11)
	mask    = make_source_mask(data, nsigma=3, npixels=5, dilate_size=11)
	mean, median, std = sigma_clipped_stats(data, sigma=3.0, mask=mask)
	return mean, median, std
#-------------------------------------------------------------------------#
def limitmag(N, zp, aper, skysigma):			# 3? 5?, zp, diameter [pixel], skysigma
	import numpy as np
	R           = float(aper)/2.				# to radius
	braket      = N*skysigma*np.sqrt(np.pi*(R**2))
	upperlimit  = float(zp)-2.5*np.log10(braket)
	return round(upperlimit, 3)
#-------------------------------------------------------------------------#
'''
def targetfind(ra1, de1, ra2, de2, sep):

	import numpy as np
	dist	= np.sqrt( (ra1-ra2)**2. + (de1-de2)**2. )
	indx	= np.where( (dist == np.min(dist)) &
						(dist < sep/3600.) )
	return indx
'''
#-------------------------------------------------------------------------#
def plotshow(inim, numb_list, xim_list, yim_list, outname='default', add=None, numb_addlist=None, xim_addlist=None, yim_addlist=None, invert=False):
	'''
	PLOT IMAGE AND SHOW DESINATED OBJECTS
	'''
	import numpy as np
	import matplotlib.pyplot as plt
	from astropy.io import fits
	from matplotlib.colors import LogNorm
	from matplotlib.patches import Circle
	from astropy.visualization import (MinMaxInterval, SqrtStretch, ImageNormalize)
	from astropy.visualization import ZScaleInterval, LinearStretch
	from astropy.wcs import WCS
	if outname == 'default':
		outname		= inim[:-5]+'.png'
	else:
		pass
	data, hdr	= fits.getdata(inim, header=True)
	if invert == True:
		data = -1*data
	#fig, ax		= plt.subplots(1)

	plt.close('all')
	plt.rc('font', family='serif')
	fig = plt.figure()
	x = 1080 / 2 / fig.dpi
	y = 1080 / 2 / fig.dpi
	fig.set_figwidth(x)
	fig.set_figheight(y)

	wcs			= WCS(hdr)
	norm_zscale	= ImageNormalize(data, interval=ZScaleInterval(), stretch=LinearStretch())
	# fig			= plt.figure()
	ax			= plt.subplot(projection=wcs)
	im			= ax.imshow(data, cmap='gray', origin='lower', norm=norm_zscale)

	for xx, yy in zip(xim_list, yim_list):
		circ = Circle((xx, yy), 15, color='gold', fill=None, linewidth='0.3')
		ax.add_patch(circ)
	for i, txt in enumerate(numb_list):
		xim		= xim_list[i]
		yim		= yim_list[i]
		ax.text(xim+7.5, yim+7.5, str(txt), color='gold', fontsize=5)
	if add != None:
		for xx, yy in zip(xim_addlist, yim_addlist):
			'''
			circ = Circle((xx, yy), 15, color='tomato', fill=None, linewidth='0.5')
			ax.add_patch(circ)
			'''
			ax.scatter(xx, yy, color='tomato', marker='o', alpha=0.1, s=10)
		for i, txt in enumerate(numb_addlist):
			xim		= xim_addlist[i]
			yim		= yim_addlist[i]
			ax.text(xim+7.5, yim+7.5, str(txt), color='tomato', fontsize=5)
	else:
		pass
	ax.grid('both', linestyle='--', color='silver', alpha=0.5)
	ax.set_xlabel('R.A.', fontsize=20)
	ax.set_ylabel('Dec.', fontsize=20)
	# plt.tight_layout()
	# plt.minorticks_on()
	fig.savefig(outname, facecolor='w', edgecolor='w',
		orientation='portrait', papertype=None, format=None,
		transparent=False, bbox_inches=None, pad_inches=0.1,
		metadata=None)
#-------------------------------------------------------------------------#
def sedualcom(inim, gain, pixscale, seeing, det_sigma=1.5, backsize=str(64), backfiltersize=str(3), detect='detection.fits'):
	"""
	SourceEXtractor
	APERTURE    3", 5", 7",
				1.0seeing, 1.2seeing ,1.5seeing ,1.7seeing ,2.0seeeing
	INPUT   :   (image).fits
				aperture    []
				seeing_fwhm [pixel]
	OUTPUT  :   no return
				.cat
	"""
	import numpy as np
	import os
	from astropy.io import ascii
	#   FILE CHECK
	#	CONFIG FILES (USER BASE PATH)
	sharepath       = '/home/sonic/Research/yourpy/config'
	configfile      = sharepath+'/targetphot.sex'
	paramfile       = sharepath+'/targetphot.param'
	nnwfile		    = sharepath+'/targetphot.nnw'
	convfile	    = sharepath+'/targetphot.conv'
	try:
		comment = '\nSourceEXtractor (DUAL MODE) START\n' \
				+ 'IMAGE\t\t: '+inim+'\n' \
				+ 'GAIN\t\t: '+str(gain)+'\n' \
				+ 'PIXSCALE\t: '+str(pixscale)+'\n' \
				+ 'DETECTION SIGMA\t: '+str(det_sigma)+'\n' \
				+ 'PARAM\t\t: '+paramfile+'\n' \
				+ 'BACKSIZE\t: '+backsize+'\n' \
				+ 'BACKFILTER\t: '+backfiltersize+'\n' \
				+ 'CONFIG\t\t: '+configfile+'\n' \
				+ 'NNW\t\t: '+nnwfile+'\n' \
				+ 'CONVOLVE\t: '+convfile
		print(comment)
	except:
		comment = 'CHECK configfile/paramfile/nnewfile/convfile or others.'
		print(comment)
	#   FILE NAME
	oriim   = inim[2:]
	cat     = inim[:-5]+'.dual.cat'
	seg     = inim[:-5]+'.seg.fits'
	bkg     = inim[:-5]+'.bkg.fits'
	sub     = inim[:-5]+'.sub.fits'
	# psf     = inim[2:-5]+'.psf'
	psf		= inim[:-5]+'.psf'
	aper    = inim[:-5]+'.aper.fits'

	#   BASIC INFO.
	det_area        = 5
	#det_thresh      = det_sigma/np.sqrt(det_area)
	det_thresh      = det_sigma
	detecminarea    = str(det_area)
	detectthresh    = str(det_thresh)
	# seeing, fwhm_arcsec = psfex(oriim, pixscale)
	# seeing, fwhm_arcsec = psfex(inim, pixscale)
	#pixscale        = pixscalecalc(inim)
	#   OPTION
	aperture        = '%.2f'%(3./pixscale)+','+'%.2f'%(5./pixscale)+','+'%.2f'%(7./pixscale)+','+'%.2f'%(seeing)+','+'%.2f'%(1.2*seeing)+','+'%.2f'%(1.5*seeing)+','+'%.2f'%(1.7*seeing)+','+'%.2f'%(2.0*seeing)

	option1 = ' -CATALOG_NAME '+cat+' -PARAMETERS_NAME '+paramfile
	option2 = ' -DETECT_MINAREA '+detecminarea+' -DETECT_THRESH '+detectthresh \
			+' -FILTER Y '+'-FILTER_NAME '+convfile
	option3 = ' -PHOT_APERTURES '+aperture \
			+' -GAIN '+'%.1f'%(gain)+' -PIXEL_SCALE '+'%.1f'%(pixscale)
	option4 = ' -SEEING_FWHM '+'%.3f'%(seeing)+' -STARNNW_NAME '+nnwfile
	option5 = ' -BACK_SIZE '+ backsize \
			+ ' -BACK_FILTERSIZE '+ backfiltersize+' -BACKPHOTO_TYPE LOCAL'
	option6 = ' -CHECKIMAGE_TYPE SEGMENTATION,APERTURES,BACKGROUND,-BACKGROUND'
	option7 = ' -CHECKIMAGE_NAME '+seg+','+aper+','+','+bkg+','+sub
	option8 = ' -PSF_NAME '+psf
	#   COMMAND
	#   detect = detection.fits is fine image show target significantly and have good seeing and many stars i
	dualcom ='sex -c '+configfile+' '+detect+' , '+inim+' -CATALOG_NAME '+cat+' -PARAMETERS_NAME '+paramfile+ ' '+option2+' '+option3+' '+option4+' '+option5+' '+option6+' '+option7+' '+option8

	os.system(dualcom)
	secat   = ascii.read(cat)
	# return secat, cat, seeing, fwhm_arcsec
	return secat, cat
#-------------------------------------------------------------------------#
def targetfind(tra, tdec, refra, refdec, sep):
	import astropy.units as u
	from astropy.coordinates import SkyCoord
	targ_coord	= SkyCoord(tra, tdec, unit=(u.deg, u.deg))
	phot_coord	= SkyCoord(refra, refdec, unit=(u.deg, u.deg))
	indx, d2d, d3d	= targ_coord.match_to_catalog_sky(phot_coord)
	return indx, d2d, d3d
#-------------------------------------------------------------------------#
def apass2med(incat, outcat, sedcat='/home/paek/table/stellar_sed_template_phot4med.dat'):
	from astropy.modeling import models, fitting
	from astropy.io import ascii
	#============================================================
	#	Function
	#------------------------------------------------------------
	def chi_sq(obs, exp):
		'''
		obs : observation value
		exp : expectation value
		'''
		return np.sum( ((obs-exp)**2)/exp )
	#------------------------------------------------------------
	#	Path
	#------------------------------------------------------------
	# path_plot = f'{path_base}/plot'
	#------------------------------------------------------------
	#	Table
	#------------------------------------------------------------
	reftbl = ascii.read(incat)
	sedtbl = ascii.read(sedcat)
	#------------------------------------------------------------
	#	Advance Preparation
	#------------------------------------------------------------
	#	med
	# filterlist_med = ['m575', 'm625', 'm675', 'm725', 'm775']
	filterlist_med = ['m425', 'm475', 'm525', 'm575', 'm625', 'm675', 'm725', 'm775', 'm825', 'm875', 'm925', 'm975', 'm1025', 'n6780', 'n6830',]
	#	APASS
	filterlist_apass = ['B', 'V', 'g', 'r', 'i']
	#	Vega --> AB
	reftbl['Bmag'] = reftbl['Bmag']+(-0.09)
	reftbl['Vmag'] = reftbl['Vmag']+(-0.02)
	#	Generate table space on the reference catalog
	for filte in filterlist_med:
		reftbl[f'{filte}mag'] = 0.0
		reftbl[f'e_{filte}mag'] = 0.0
	reftbl['n_sed'] = 0
	reftbl['mktype'] = ' '*10
	#	Dummy column
	sedtbl['chisq'] = 0.0
	sedtbl['const'] = 0.0
	#------------------------------------------------------------
	#	Main body
	#------------------------------------------------------------
	# i = 2
	print('-'*60)
	print(f'#\tConvert APASS catalog (stars:{len(reftbl)})')
	print('-'*60)
	for i in range(len(reftbl)):
		n_apass = reftbl['NUMBER'][i]
		ms_apass = np.array([reftbl[f'{filte}mag'][i] for filte in filterlist_apass])
		e_ms_apass = np.array([reftbl[f'e_{filte}mag'][i] for filte in filterlist_apass])
		#	Variable for fitting
		y = ms_apass
		yerr = e_ms_apass
		# j = 60
		for j in range(len(sedtbl)):
			ms_sed = np.array([sedtbl[f'{filte}'][j].item() for filte in filterlist_apass])
			#	Variable for fitting
			x = ms_sed
			#	Initial parameter
			m_dif = np.median(y-x)
			#------------------------------------------------------------
			#	Fitting
			#------------------------------------------------------------
			# initialize a linear fitter
			fit = fitting.LinearLSQFitter()
			# initialize a linear model
			line_init = models.Linear1D()
			line_init.slope.fixed = True
			# fit the data with the fitter
			# fitted_line = fit(line_init, x, y, weights=e_ms_apass)
			fitted_line = fit(line_init, x, y, )
			#------------------------------------------------------------
			#	Fitting results
			#------------------------------------------------------------
			chisq = chi_sq(ms_apass, fitted_line(x))
			sedtbl['chisq'][j] = chisq
			sedtbl['const'][j] = fitted_line.intercept.value
		#------------------------------------------------------------
		#	Pick optimized index
		#------------------------------------------------------------
		indx_opt = np.where(sedtbl['chisq']==np.min(sedtbl['chisq']))
		#------------------------------------------------------------
		#	Optimized values
		#------------------------------------------------------------
		chisq_opt = sedtbl['chisq'][indx_opt].item()
		const_opt = sedtbl['const'][indx_opt].item()
		ms_sed = np.array([sedtbl[f'{filte}'][indx_opt].item() for filte in filterlist_apass])
		x = ms_sed
		'''
		#------------------------------------------------------------
		# 	Plot
		#------------------------------------------------------------
		plt.close()
		fig, (ax1, ax2) = plt.subplots(2, 1)
		#	axis 1
		ax1.plot(x, y, 'ko', label='APASS')
		ax1.plot(x, x+const_opt, 'r-', label=f'Fitted Model, chisq={round(chisq_opt, 3)}')
		ax1.legend()
		ax1.set_title(f'APASS : {i}')

		ax1.set_ylabel(r'$\rm m_{APASS}$ [mag]')
		#	axis 2
		ax2.plot(x, y - (x+const_opt), 'ko')
		ax2.axhline(y=0, ls='--', color='grey')
		ax2.set_xlabel(r'$\rm m_{spec}$ [mag]')
		ax2.set_ylabel('Residual')
		ax2.set_ylim([-0.5, +0.5])

		plt.savefig(f'{path_plot}/plot_{i}.png')#, overwrite=True)

		#	Put med-band photometries on the APASS catalog
		for filte in filterlist_med:
			reftbl[f'{filte}mag'][i] = sedtbl[f'{filte}'][indx_opt].item()+const_opt
		reftbl['n_sed'][i] = sedtbl['number'][indx_opt].item()
		reftbl['mktype'][i] = sedtbl['mktype'][indx_opt].item()
		'''
		#	Put med-band photometries on the APASS catalog
		for filte in filterlist_med:
			reftbl[f'{filte}mag'][i] = sedtbl[f'{filte}'][indx_opt].item()+const_opt
		reftbl['n_sed'][i] = sedtbl['number'][indx_opt].item()
		reftbl['mktype'][i] = sedtbl['mktype'][indx_opt].item()

	reftbl.write(f'{outcat}', format='ascii.tab', overwrite=True)
	return reftbl

import numpy,from numpy,from numba,import astropy,from astropy


import os, sys, glob
import argparse
from struct import *
from pathlib import Path
from typing import List, Dict, Optional, Union, Type

import numpy as np
import pandas as pd

import astropy
from astropy import units as un
from astropy import constants as const
from astropy.cosmology import z_at_value, LambdaCDM

from astrild.simulation import Simulation
from astrild.particles.ecosmog import Ecosmog
from astrild.particles.halo import Halos
from astrild.utils.geometrical_transforms import (
    transform_box_to_lc_cart_coords,
    radial_coordinate_in_lc,
    angular_coordinate_in_lc,
)
from astrild.io import IO

dir_src = Path(__file__).parent.absolute()
default_config_file_df = dir_src / "configs/ray_snapshot_info.h5"

class RayRamsesWarning(BaseException):
    pass

class RayRamses(Simulation):
    """
    Attributes:
        dir_sim: directory of simulation
        dir_out: directory for results/returns
        file_dsc: file describtion for identification
        dir_root:
        opening_angle: light-cone opening angle; [deg^2]
        npix: nr. of cells along the four light-cone edges
        config: contains snapshot infos (e.g. redshifts, comov. dist, ...)

    Methods:
        compress_snapshot:
        sum_snapshots:
        find_halos_in_raytracing_box:
        find_halos_in_raytracing_snapshot:
    """

    def __init__(
        self,
        config: Union[None, pd.DataFrame],
        dir_sim: str,
        dir_out: str,
        file_dsc: dict = {"root": "Ray_maps", "extension": "dat"},
        dir_root: str = None,
        opening_angle: float = 20.0,
        npix: int = 8192,
        cosmo: astropy.cosmology = LambdaCDM,
    ):
        super().__init__(dir_sim, dir_out, file_dsc, dir_root)
        self.opening_angle = opening_angle
        self.npix = npix
        self.config = config
        self.cosmo = cosmo

    def compress_snapshot(
        self,
        fields: list,
        dir_out: str = None,
        convert: bool = False,
        cosmo: astropy.cosmology = LambdaCDM,
        save: bool = True,
    ) -> None:
        """
        Combines the ray-tracing outputs of individual CPUs at a given snapshot
        into one pandas .h5 file.

        Args:
        Returns:
        """
        self.cosmology = cosmo

        self.file_nrs = self.get_file_nrs(
            self.file_dsc, self.dirs["sim"], "min", True
        )
        self.files = {
            self.file_dsc["root"]: self.get_file_paths(
                self.file_dsc, self.dirs["sim"], uniques="min"
            )
        }
        # run through ray-ramses snapshots
        for ray_nr in np.unique(self.file_nrs):
            print("Ray-Nr %d" % ray_nr)
            idx = (np.where(self.file_nrs == ray_nr)[0]).astype(int)
            cpu_files = [
                self.files[self.file_dsc["root"]][idx]
                for idx, rnr in enumerate(self.file_nrs)
                if rnr == ray_nr
            ]
            first = True

            # run through cpu-files for snapshot
            for cpu_file in cpu_files:
                print("    %s" % cpu_file.split("/")[-1])

                ray_df = pd.read_csv(
                    cpu_file,
                    delim_whitespace=True,
                    skipinitialspace=True,
                    names=fields,
                    header=None,
                    lineterminator="\n",
                )

                if convert is True:
                    # convert comoving distance to [Gpc/h]
                    ray_df["chi_co"] /= self.cosmology.H0.value / 100

                    # Correct gamma_1 and gamma_2 to put in the form
                    # that normally appears in the literature
                    ray_df["shear_y"] *= 2.0 * sin(ray_df["the_co"])
                    gamm1_corr = -ray_df["shear_x"] * cos(
                        2.0 * ray_df["phi_co"]
                    ) - ray_df["shear_y"] * sin(2.0 * ray_df["phi_co"])
                    gamm2_corr = -ray_df["shear_x"] * sin(
                        2.0 * ray_df["phi_co"]
                    ) + ray_df["shear_y"] * sin(2.0 * ray_df["phi_co"])
                    ray_df["shear_x"] = gamm1_corr
                    ray_df["shear_y"] = gamm2_corr

                if first:
                    ray_collect_df = ray_df
                    first = False
                else:
                    ray_collect_df = ray_collect_df.append(ray_df)

            ray_collect_df = ray_collect_df.sort_values(
                by=["rayid"], axis=0, ascending=True
            )
            ray_collect_df = ray_collect_df.set_index("rayid")
            file_out = self.dirs["sim"] + "%s_output%05d.h5" % (
                self.file_dsc["root"],
                ray_nr,
            )
            ray_collect_df.to_hdf(file_out, key="df", mode="w")

    def sum_snapshots(
        self,
        dir_out: str,
        columns: list,
        columns_z_shift: list,
        integration_range: dict,
        ray_file_root: str = "Ray_maps_output%05d.h5",
        sim_folder_root: str = "box%d",
        z_src: float = None,
        z_src_shift: float = None,
    ) -> None:
        """
        Adds different ray-tracing outputs together. This can give you the
        integrated ray-tracing quantities between arbitrary redshifts along
        the ligh-cone. The ray-tracing outputs must either have the format of
        RayRamses outputs in pd.DataFrames or np.ndarray images.

        Args:
            dir_out:
            columns:
            columns_z_shift:
            integration_range:
            ray_file_root:
            sim_folder_root:
            z_src:
            z_src_shift:
        """
        file_name = self.dirs["lc"] + "ray_snapshot_info.h5"
        if not os.path.isfile(file_name):
            raise RayRamsesWarning(
                "The file 'ray_snapshot_info.h5' does note exist"
            )
        self.ray_info_df = pd.read_hdf(file_name, key="s")
        sim_folder_root = self.dirs["lc"] + sim_folder_root

        box_ray_nrs = _get_box_and_ray_nrs(integration_range)

        first = True
        for box_nr, ray_nr in box_ray_nrs:
            sim_info_df = self.ray_info_df.loc[(box_nr, ray_nr)]
            self.dirs["sim"] = sim_folder_root % box_nr + "/"
            ray_file = self.dirs["sim"] + ray_file_root % ray_nr
            ray_map_df = pd.read_hdf(ray_file)

            print(
                "Box Nr. %d; %s; Redshift %.3f"
                % (box_nr, os.path.basename(ray_file), sim_info_df["redshift"]),
                len(ray_map_df.index.values),
            )

            if (
                z_src_shift is not None
                and sim_info_df["redshift"] <= z_src_shift
            ):
                raise RayRamsesWarning(
                    "Redshift shift has not correct data structure"
                )
                # what snapshot to use if end of lightcone-box is reached
                if (ray_box_info_df.name[1] == ray_nrs.min()) and (box_nr < 4):
                    z_next = self.ray_info_df.loc[(box_nr + 1, 1)]["redshift"]
                else:
                    z_next = ray_box_info_df.iloc[ii + 1]["redshift"]

                # Shift redshift of light source
                # only of kappa but not of iswrs !!!
                ray_map_df["kappa_2"] = self._translate_redshift(
                    ray_map_df["kappa_2"],
                    sim_info_df["redshift"],
                    z_next,
                    z_src,
                    z_src_shift,
                )

            if first is True:
                ray_df_sum = ray_map_df
                first = False

            else:
                for column in columns:
                    ray_df_sum[column] = (
                        ray_df_sum[column].values + ray_map_df[column].values
                    )

        self._merged_snapshots_to_file(ray_df_sum, dir_out, integration_range)

    def _get_box_and_ray_nrs(self, integration_range: dict) -> np.ndarray:
        """
        Get all box and ray-snapshot numbers for selected range.

        Args:
        Returns:
        """
        if not integration_range["z"]:
            if integration_range["box"][0] == 0:
                print("Integrate over whole light-cone")
                self.complete_lc = True
            elif integration_range["ray"][0] == 0:
                print("Integrate over box", integration_range["box"])
                self.ray_info_df = ray_info_df[
                    ray_info_df.index.get_level_values(0).isin(
                        integration_range["box"]
                    )
                ]
                self.complete_lc = False
        else:
            print("Integrate over redshift-range", integration_range["z"])
            # if merging based on redshift
            z_range = np.asarray(integration_range["z"])
            self.ray_info_df = self.ray_info_df[
                (z_range.min() < self.ray_info_df["redshift"])
                & (self.ray_info_df["redshift"] < z_range.max())
            ]
            self.complete_lc = False

        return self.ray_info_df.index.values

    def _translate_redshift(
        self,
        quantity: str,
        z_near: float,
        z_far: float,
        z_src: float,
        z_src_shift: float,
    ) -> float:
        """
        Args:
            quantity pandas.DataSeries:
                ray-ramses output quantity
            x_near np.float:
                comoving distance closer to observer
            x_far np.float:
                comoving distance further from observer
            x_src np.float:
                source redshift used in ray-ramses simulation

        Returns:
        """
        x_far = self.cosmology.comoving_distance(z_far).to_value("Mpc")
        x_near = self.cosmology.comoving_distance(z_near).to_value("Mpc")
        x_src = self.cosmology.comoving_distance(z_src).to_value("Mpc")

        if z_far > z_src_shift:
            # if z of next snapshot larger than new source z, set the new source
            # equal to it, so that a distance of 150[Mpc/h] is maintained
            x_src_shift = self.cosmology.comoving_distance(z_far).to_value(
                "Mpc"
            )
        else:
            x_src_shift = self.cosmology.comoving_distance(
                z_src_shift
            ).to_value("Mpc")

        x_mid = 0.5 * (x_far + x_near)

        quantity_shift = (
            quantity
            * self._kernel_function(x_mid, x_src_shift)
            / self._kernel_function(x_mid, x_src)
        )
        return quantity_shift

    def _kernel_function(self, x: float, x_s: float) -> float:
        """
        Args:
            x np.float:
                comoving distance
            x_s np.float:
                comoving distance to source

        Returns:
        """
        g = (x_s - x) * x / x_s
        return g

    def _merged_snapshots_to_file(
        self, ray_df_sum: pd.DataFrame, dir_out: str, integration_range: dict
    ) -> None:
        """
        Write merged ray-tracing pd.DataFrame to .h5 file

        Args:
        """
        if not integration_range["z"]:
            if integration_range["box"][0] == 0:
                fout = dir_out + "Ray_maps_lc.h5"
                print("Save in %s" % fout)
                ray_df_sum.to_hdf(fout, key="df", mode="w")
            elif integration_range["ray"][0] == 0:
                fout = dir_out + "Ray_maps_box%d.h5" % box_nr
                print("Save in %s" % fout)
                ray_df_sum.to_hdf(fout, key="df", mode="w")
        else:
            fout = dir_out + "Ray_maps_zrange_%.2f_%.2f.h5" % (
                self.ray_info_df["redshift"].values.min(),
                self.ray_info_df["redshift"].values.max(),
            )
            print("Save in %s" % fout)
            ray_df_sum.to_hdf(fout, key="df", mode="w")

    @staticmethod
    def find_subfind_halos_in_raytracing_snapshot(
        ecosmog: Type[Ecosmog],
        box_nr: int,
        snap_nr: int,
        ray_nr: int,
        boxdist: float,
        boxsize: float,
        opening_angle: float,
        snaplimit: tuple,
        hubble: float,
    ) -> Union[None, pd.DataFrame]:
        coeff = hubble / 1e3
        halos = Halos.from_subfind(snap_nr, ecosmog)

        if halos.data is None:
            return None
        else:
            halocat = halos.filter_nonzero_subfind_halos_size(halos.data).cat

        halocatindex = np.arange(len(halocat["Group_M_Crit200"][:]))
        print(f"There are {len(halocatindex)} halos in box {box_nr} snapshot {snap_nr}")

        # cartesian coord. in light-cone [Mpc/h]
        pos = halocat["GroupPos"][:, :] * coeff
        pos[:, 0] = pos[:, 0] - boxsize / 2
        pos[:, 1] = pos[:, 1] - boxsize / 2
        pos[:, 2] = pos[:, 2] + boxdist
        # radial distance from observer [Mpc/h]
        rad_dist = np.sqrt(pos[:, 0]**2 + pos[:, 1]**2 + pos[:, 2]**2)
        # angular coord [deg]
        theta1_deg = np.arctan(pos[:, 0] / pos[:, 2]) * 180/np.pi
        theta2_deg = np.arctan(pos[:, 1] / pos[:, 2]) * 180/np.pi

        # index of halos in light-cone
        indx = np.where(
            (rad_dist >= np.min(snaplimit))
            & (rad_dist <= np.max(snaplimit))
            & (np.abs(theta1_deg) <= opening_angle / 2)
            & (np.abs(theta2_deg) <= opening_angle / 2)
        )[0]
        print(f"There are {len(indx)} halos in light-cone in box {box_nr} snapshot {snap_nr}")

        # project 3D velocity along line-of-sight in cart. coord.
        pos_norm = np.linalg.norm(pos[indx, :], axis=1)
        vr = (
            (
                halocat["GroupVel"][indx, :] * pos[indx, :] # element-wise dot-product
            ).sum(axis=1) / (pos_norm**2)
        )[:, np.newaxis] * pos[indx, :]
        # project 3D velocity along transverse to line-of-sight direction
        # in cart. coord.
        vt = halocat["GroupVel"][indx, :] - vr
        # small angle approximation
        # -> TODO improve by projecting on spher. coord. unit vectors e_theta and e_phi
        vel_x = vt[:, 0]
        vel_y = vt[:, 1]

        r200_deg = (
            np.arctan(halocat["Group_R_Crit200"][indx] * coeff / rad_dist[indx])
            * 180 / np.pi
        )

        halo_id = [
            int(f"{box_nr}{snap_nr}{ii}")
            for ii in halocatindex[indx].astype(int)
        ]
        halos_dict = {
            "id": halo_id,
            "x": pos[indx, 0],
            "y": pos[indx, 1],
            "z": pos[indx, 2],
            "rad_dist": rad_dist[indx],
            "theta1_deg": theta1_deg[indx] + opening_angle / 2,
            "theta1_pix": _degree_to_pixel(
                theta1_deg[indx] + opening_angle / 2
            ),
            "theta2_deg": theta2_deg[indx] + opening_angle / 2,
            "theta2_pix": _degree_to_pixel(
                theta2_deg[indx] + opening_angle / 2
            ),
            "x_vel": halocat["GroupVel"][indx, 0],
            "y_vel": halocat["GroupVel"][indx, 1],
            "z_vel": halocat["GroupVel"][indx, 2],
            "theta1_tv": vel_x,
            "theta2_tv": vel_y,
            "m200": halocat["Group_M_Crit200"][indx],
            "r200_deg": r200_deg,
            "r200_pix": _degree_to_pixel(r200_deg),
            "ray_nr": [ray_nr + 1] * len(indx),
            "snap_nr": [snap_nr] * len(indx),
        }
        halos_df = pd.DataFrame(data=halos_dict)
        return halos_df

    def find_halos_in_raytracing_box(
        self,
        ecosmog: Type[Ecosmog],
        snapdist: float,
        box_nr: int,
        boxsize: float,
        hubble: float,
    ) -> Union[None, pd.DataFrame]:
        """
        Args:

        Returns:
        """
        boxdist = snapdist[-1]
        first = True
        # run through ray-ramses snapshots
        for ray_nr in np.unique(self.file_nrs)[:-1]:
            snap_nr = (
                ray_nr + len(ecosmog.config.index) - len(self.config.index)
            )
            snaplimit = (snapdist[ray_nr - 1], snapdist[ray_nr])

            if "fof" in list(ecosmog.files.keys()):
                halos_df = self.find_subfind_halos_in_raytracing_snapshot(
                    ecosmog,
                    box_nr,
                    snap_nr,
                    ray_nr,
                    boxdist,
                    boxsize,
                    self.opening_angle,
                    self.npix,
                    snaplimit,
                    hubble,
                )
            elif "halos" in list(ecosmog.files.keys()):
                halos_df = self.find_rockstar_halos_in_raytracing_snapshot(
                    ecosmog,
                    box_nr,
                    snap_nr,
                    ray_nr,
                    boxdist,
                    boxsize,
                    self.opening_angle,
                    self.npix,
                    snaplimit,
                )

            if (first is True) and (halos_df is not None):
                halos_df_sum = halos_df
                first = False
            elif first is False:
                halos_df_sum = halos_df_sum.append(halos_df, ignore_index=True)
        return halos_df_sum

    @staticmethod
    def find_subfind_halos_in_raytracing_snapshot(
        ecosmog: Type[Ecosmog],
        box_nr: int,
        snap_nr: int,
        ray_nr: int,
        boxdist: float,
        boxsize: float,
        opening_angle: float,
        npix: int,
        snaplimit: tuple,
        hubble: float,
    ) -> Union[None, pd.DataFrame]:
        coeff = hubble / 1e3
        halos = Halos.from_subfind(snap_nr, ecosmog)

        if halos.data is None:
            return None
        else:
            halocat = halos.filter_nonzero_subfind_halos_size(halos.data).cat

        halocatindex = np.arange(len(halocat["Group_M_Crit200"][:]))
        print(f"There are {len(halocatindex)} halos in box {box_nr} snapshot {snap_nr}")

        # cartesian coord. in light-cone [Mpc/h]
        pos = halocat["GroupPos"][:, :] * coeff
        pos = transform_box_to_lc_cart_coords(pos, boxsize, boxdist)
        rad_dist = radial_coordinate_in_lc(pos)
        theta1_deg, theta2_deg = angular_coordinate_in_lc(pos, unit="deg")

        # index of halos in light-cone
        indx = np.where(
            (rad_dist >= np.min(snaplimit))
            & (rad_dist <= np.max(snaplimit))
            & (np.abs(theta1_deg) <= opening_angle / 2)
            & (np.abs(theta2_deg) <= opening_angle / 2)
        )[0]
        print(f"There are {len(indx)} halos in light-cone in box {box_nr} snapshot {snap_nr}")

        # project 3D velocity along line-of-sight in cart. coord.
        pos_norm = np.linalg.norm(pos[indx, :], axis=1)
        vr = (
            (
                halocat["GroupVel"][indx, :] * pos[indx, :] # element-wise dot-product
            ).sum(axis=1) / (pos_norm**2)
        )[:, np.newaxis] * pos[indx, :]
        # project 3D velocity along transverse to line-of-sight direction
        # in cart. coord.
        vt = halocat["GroupVel"][indx, :] - vr
        # small angle approximation
        # -> TODO improve by projecting on spher. coord. unit vectors e_theta and e_phi
        vel_x = vt[:, 0]
        vel_y = vt[:, 1]

        r200_deg = (
            np.arctan(halocat["Group_R_Crit200"][indx] * coeff / rad_dist[indx])
            * 180 / np.pi
        )

        halo_id = [
            int(f"{box_nr}{snap_nr}{ii}")
            for ii in halocatindex[indx].astype(int)
        ]
        halos_dict = {
            "id": halo_id,
            "x": pos[indx, 0],
            "y": pos[indx, 1],
            "z": pos[indx, 2],
            "rad_dist": rad_dist[indx],
            "theta1_deg": theta1_deg[indx] + opening_angle / 2,
            "theta1_pix": _degree_to_pixel(
                theta1_deg[indx] + opening_angle / 2, opening_angle, npix
            ),
            "theta2_deg": theta2_deg[indx] + opening_angle / 2,
            "theta2_pix": _degree_to_pixel(
                theta2_deg[indx] + opening_angle / 2, opening_angle, npix
            ),
            "x_vel": halocat["GroupVel"][indx, 0],
            "y_vel": halocat["GroupVel"][indx, 1],
            "z_vel": halocat["GroupVel"][indx, 2],
            "theta1_tv": vel_x,
            "theta2_tv": vel_y,
            "m200": halocat["Group_M_Crit200"][indx],
            "r200_deg": r200_deg,
            "r200_pix": _degree_to_pixel(r200_deg, opening_angle, npix),
            "ray_nr": [ray_nr + 1] * len(indx),
            "snap_nr": [snap_nr] * len(indx),
        }
        halos_df = pd.DataFrame(data=halos_dict)
        return halos_df

    @staticmethod
    def find_rockstar_halos_in_raytracing_snapshot(
        ecosmog: Type[Ecosmog],
        box_nr: int,
        snap_nr: int,
        ray_nr: int,
        boxdist: float,
        boxsize: float,
        opening_angle: float,
        npix: int,
        snaplimit: tuple,
    ) -> Union[None, pd.DataFrame]:
        halos = Halos.from_rockstar(snap_nr, ecosmog)

        if halos.data is None:
            return None
        else:
            halocat = halos.data

        halocatindex = np.arange(len(halocat["m200c"].values))
        print(f"There are {len(halocatindex)} halos in box {box_nr} snapshot {snap_nr}")

        # cartesian coord. in light-cone [Mpc/h]
        pos = halocat[["x", "y", "z"]].values
        pos = transform_box_to_lc_cart_coords(pos, boxsize, boxdist)
        rad_dist = radial_coordinate_in_lc(pos)
        theta1_deg, theta2_deg = angular_coordinate_in_lc(pos, unit="deg")

        # index of halos in light-cone
        indx = np.where(
            (rad_dist >= np.min(snaplimit))
            & (rad_dist <= np.max(snaplimit))
            & (np.abs(theta1_deg) <= opening_angle / 2)
            & (np.abs(theta2_deg) <= opening_angle / 2)
        )[0]
        halocat = halocat.iloc[halocatindex[indx]]
        if len(indx) == 0: return None

        pos = pos[halocatindex[indx], :]
        rad_dist = rad_dist[halocatindex[indx]]
        theta1_deg = theta1_deg[halocatindex[indx]]
        theta2_deg = theta2_deg[halocatindex[indx]]
        print(f"There are {len(halocat.index.values)} halos in light-cone in box {box_nr} snapshot {snap_nr}")

        # get redshift
        redshift = Dc_to_redshift(ecosmog.cosmo, rad_dist * un.Mpc)

        # get angular distance

        # project 3D velocity along line-of-sight in cart. coord.
        pos_norm = np.linalg.norm(pos, axis=1)
        vr = (
            (
                halocat[["vx", "vy", "vz"]].values * pos # element-wise dot-product
            ).sum(axis=1) / (pos_norm**2)
        )[:, np.newaxis] * pos
        # project 3D velocity along transverse to line-of-sight direction
        # in cart. coord.
        vt = halocat[["vx", "vy", "vz"]].values - vr
        # small angle approximation
        # -> TODO improve by projecting on spher. coord. unit vectors e_theta and e_phi

        r200_deg = (
            np.arctan(halocat["r200c"].values / (rad_dist * 1e3))
            * 180 / np.pi
        )

        halo_id = [
            int(f"{box_nr}{snap_nr}{ii}")
            for ii in halocatindex[indx].astype(int)
        ]
        halos_dict = {
            "id": halo_id,
            "x": pos[:, 0],
            "y": pos[:, 1],
            "z": pos[:, 2],
            "Dc": rad_dist,
            "Da": rad_dist / (1 + redshift),
            "redshift": redshift,
            "theta1_deg": theta1_deg + opening_angle / 2,
            "theta1_pix": _degree_to_pixel(theta1_deg + opening_angle/2, opening_angle, npix),
            "theta2_deg": theta2_deg + opening_angle / 2,
            "theta2_pix": _degree_to_pixel(theta2_deg + opening_angle/2, opening_angle, npix),
            "x_vel": halocat["vx"].values,
            "y_vel": halocat["vy"].values,
            "z_vel": halocat["vz"].values,
            "theta1_tv": vt[:, 0],
            "theta2_tv": vt[:, 1],
            "m200": halocat["m200c"].values,
            "c_NFW": halocat["r200c"].values / halocat["Rs"].values,
            "r200_deg": r200_deg,
            "r200_pix": _degree_to_pixel(r200_deg, opening_angle, npix),
            "ray_nr": [ray_nr + 1] * len(pos),
            "snap_nr": [snap_nr] * len(pos),
        }
        halos_df = pd.DataFrame(data=halos_dict)
        return halos_df

def _degree_to_pixel(deg: np.ndarray, opening_angle, npix) -> np.ndarray:
    """ Convert degree to pixel position """
    return np.ceil(deg * npix / opening_angle).astype(int)

def Dc_to_redshift(cosmo: astropy.cosmology, Dc: un.quantity.Quantity,) -> np.array:
    """ Return redshift at comoving distance [Mpc] """
    return np.array([z_at_value(cosmo.comoving_distance, dist) for dist in Dc])

import numpy,import astropy,from astropy


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
NTSC RGB Colourspace
====================

Defines the *NTSC RGB* colourspace:

-   :attr:`NTSC_RGB_COLOURSPACE`.

See Also
--------
`RGB Colourspaces IPython Notebook
<http://nbviewer.ipython.org/github/colour-science/colour-ipython/blob/master/notebooks/models/rgb.ipynb>`_  # noqa

References
----------
.. [1]  `Recommendation ITU-R BT.470-6 - Conventional Television Systems
        <http://www.itu.int/dms_pubrec/itu-r/rec/bt/R-REC-BT.470-6-199811-S!!PDF-E.pdf>`_  # noqa
        (Last accessed 13 April 2014)
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models import RGB_Colourspace, normalised_primary_matrix

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013 - 2014 - Colour Developers'
__license__ = 'New BSD License - http://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-science@googlegroups.com'
__status__ = 'Production'

__all__ = ['NTSC_RGB_PRIMARIES',
           'NTSC_RGB_WHITEPOINT',
           'NTSC_RGB_TO_XYZ_MATRIX',
           'XYZ_TO_NTSC_RGB_MATRIX',
           'NTSC_RGB_TRANSFER_FUNCTION',
           'NTSC_RGB_INVERSE_TRANSFER_FUNCTION',
           'NTSC_RGB_COLOURSPACE']

NTSC_RGB_PRIMARIES = np.array(
    [[0.67, 0.33],
     [0.21, 0.71],
     [0.14, 0.08]])
"""
*NTSC RGB* colourspace primaries.

NTSC_RGB_PRIMARIES : ndarray, (3, 2)
"""

NTSC_RGB_WHITEPOINT = ILLUMINANTS.get(
    'CIE 1931 2 Degree Standard Observer').get('C')
"""
*NTSC RGB* colourspace whitepoint.

NTSC_RGB_WHITEPOINT : tuple
"""

NTSC_RGB_TO_XYZ_MATRIX = normalised_primary_matrix(NTSC_RGB_PRIMARIES,
                                                   NTSC_RGB_WHITEPOINT)
"""
*NTSC RGB* colourspace to *CIE XYZ* colourspace matrix.

NTSC_RGB_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_NTSC_RGB_MATRIX = np.linalg.inv(NTSC_RGB_TO_XYZ_MATRIX)
"""
*CIE XYZ* colourspace to *NTSC RGB* colourspace matrix.

XYZ_TO_NTSC_RGB_MATRIX : array_like, (3, 3)
"""

NTSC_RGB_TRANSFER_FUNCTION = lambda x: x ** (1 / 2.2)
"""
Transfer function from linear to *NTSC RGB* colourspace.

NTSC_RGB_TRANSFER_FUNCTION : object
"""

NTSC_RGB_INVERSE_TRANSFER_FUNCTION = lambda x: x ** 2.2
"""
Inverse transfer function from *NTSC RGB* colourspace to linear.

NTSC_RGB_INVERSE_TRANSFER_FUNCTION : object
"""

NTSC_RGB_COLOURSPACE = RGB_Colourspace(
    'NTSC RGB',
    NTSC_RGB_PRIMARIES,
    NTSC_RGB_WHITEPOINT,
    NTSC_RGB_TO_XYZ_MATRIX,
    XYZ_TO_NTSC_RGB_MATRIX,
    NTSC_RGB_TRANSFER_FUNCTION,
    NTSC_RGB_INVERSE_TRANSFER_FUNCTION)
"""
*NTSC RGB* colourspace.

NTSC_RGB_COLOURSPACE : RGB_Colourspace
"""


# -*- coding: utf-8 -*-
"""
This module contains legacy code for the cellulosic ethanol biorefinery
based on the 2011 NREL report [1]_.

References
----------
.. [1] Humbird, D., Davis, R., Tao, L., Kinchin, C., Hsu, D., Aden, A.,
    Dudgeon, D. (2011). Process Design and Economics for Biochemical
    Conversion of Lignocellulosic Biomass to Ethanol: Dilute-Acid
    Pretreatment and Enzymatic Hydrolysis of Corn Stover
    (No. NREL/TP-5100-47764, 1013269). https://doi.org/10.2172/1013269

"""

# %% Legacy code

### DO NOT DELELTE:
### This code is used to generate splits for unit operations in wastewater
### treatment of a cellulosic ethanol biorefinery.
import thermosteam as tmo
import numpy as np
from biorefineries.cornstover import chemicals
tmo.settings.set_thermo(chemicals)
chemical_groups = dict(
        OtherSugars = ('Arabinose',
                        'Mannose',
                        'Galactose',
                        'Cellobiose',
                        'Sucrose'),
        SugarOligomers = ('GlucoseOligomer',
                          'XyloseOligomer',
                          'GalactoseOligomer',
                          'ArabinoseOligomer',
                          'MannoseOligomer'),
        OrganicSolubleSolids = ('AmmoniumAcetate',
                                'SolubleLignin',
                                'Extract',
                                'LacticAcid',
                                'Cellulase'),
        InorganicSolubleSolids = ('AmmoniumSulfate',
                                  'DAP',
                                  'NaOH',
                                  'HNO3',
                                  'NaNO3'),
        Furfurals = ('Furfural',
                      'HMF'),
        OtherOrganics = ('Glycerol',
                          'Denaturant',
                          'Oil',
                          'SuccinicAcid',
                          'Xylitol'),
        COxSOxNOxH2S = ('NO',
                        'NO2',
                        'SO2',
                        'CO',
                        'H2S'),
        Protein = ('Protein',
                    'Enzyme',
                    'DenaturedEnzyme'),
        CellMass = ('WWTsludge',
                    'Z_mobilis',
                    'T_reesei'),
        OtherInsolubleSolids = ('Tar',
                                'Ash',
                                'Lime'),
        OtherStructuralCarbohydrates = ('Arabinan',
                                        'Mannan',
                                        'Galactan')
)

def find_split(IDs, flow0, flow1):
    flow0 = np.asarray(flow0)
    splits = flow0/(flow0 + np.asarray(flow1))
    chemicals = tmo.settings.get_chemicals()
    array = np.zeros(chemicals.size)
    for ID, split in zip(IDs, splits):
        if ID in chemical_groups:
            array[chemicals.get_index(chemical_groups[ID])] = split
        else:
            array[chemicals.index(ID)] = split
    return array

splits = [
    ('Ethanol', 1, 15),
    ('Water', 27158, 356069),
    ('Glucose', 3, 42),
    ('Xylose', 7, 85),
    ('OtherSugars', 13, 175),
    ('SugarOligomers', 10, 130),
    ('OrganicSolubleSolids', 182, 2387),
    ('InorganicSolubleSolids', 8, 110),
    ('Ammonia', 48, 633),
    ('AceticAcid', 0, 5),
    ('Furfurals', 5, 70),
    ('OtherOrganics', 9, 113),
    ('Cellulose', 19, 6),
    ('Xylan', 6, 2),
    ('OtherStructuralCarbohydrates', 1, 0),
    ('Lignin', 186, 64),
    ('Protein', 51, 18),
    ('CellMass', 813, 280),
    ('OtherInsolubleSolids', 68, 23)
]

anaerobic_bioreactor_sludge_splits = find_split(*zip(*splits))

splits = [
    ('Glucose', 19, 502),
    ('Xylose', 40, 1022),
    ('OtherSugars', 81, 2175),
    ('SugarOligomers', 60, 1552),
    ('OrganicSolubleSolids', 612, 15808),
    ('InorganicSolubleSolids', 97, 2513),
    ('Furfurals', 19, 513),
    ('OtherOrganics', 52, 1348),
    ('Glucan', 1230, 25),
    ('Xylan', 415, 8),
    ('OtherStructuralCarbohydrates', 94, 2),
    ('Lignin', 12226, 250),
    ('Protein', 3376, 69),
    ('CellMass', 925, 19),
    ('OtherInsolubleSolids', 4489, 92)
]

pressure_filter_splits = find_split(*zip(*splits))

splits = [
    ('Ethanol', 0, 1),
    ('Water', 381300, 2241169),
    ('Glucose', 0, 2),
    ('Xylose', 1, 3),
    ('OtherSugars', 1, 7),
    ('SugarOligomers', 1, 6),
    ('OrganicSolubleSolids', 79, 466),
    ('InorganicSolubleSolids', 4828, 28378),
    ('Ammonia', 3, 16),
    ('Furfurals', 0, 3),
    ('OtherOrganics', 1, 7),
    ('CarbonDioxide', 6, 38),
    ('O2', 3, 17),
    ('N2', 5, 32),
    ('Cellulose', 0, 194),
    ('Xylan', 0, 65),
    ('OtherStructuralCarbohydrates', 0, 15),
    ('Lignin', 0, 1925),
    ('Protein', 0, 90),
    ('CellMass', 0, 19778),
    ('OtherInsolubleSolids', 0, 707)
]

membrane_bioreactor_splits = find_split(*zip(*splits))

centrifuge_species = ('Water', 'Glucose', 'Xylose', 'OtherSugars',
                      'SugarOligomers', 'OrganicSolubleSolids',
                      'InorganicSolubleSolids', 'Ammonia', 'Furfurals',
                      'OtherOrganics', 'CO2', 'COxSOxNOxH2S', 'Cellulose',
                      'Xylan', 'OtherStructuralCarbohydrates', 'Lignin',
                      'Protein', 'CellMass', 'OtherInsolubleSolids')
S623_flow = np.array([7708, 0, 0, 1, 1, 13, 75, 3, 0, 1, 1, 2, 25, 8, 2, 250, 52, 1523, 92])
S616_flow = np.array([109098, 3, 6, 13, 9, 187, 1068, 46, 5, 8, 14, 31, 1, 0, 0, 13, 3, 80, 5])

sludge_centrifuge_splits = find_split(centrifuge_species, S616_flow, S623_flow)


"""
toy_engine.py

"""

bd_results_string = \
"""<rates>
  <solvent>
    <kT> 1 </kT>
    <debye_length> 1.0 </debye_length>
    <dielectric> 78 </dielectric>
    <vacuum_permittivity> 0.000142 </vacuum_permittivity>
    <water_viscosity> 0.243 </water_viscosity>
    <relative_viscosity> 1 </relative_viscosity>
  </solvent>
    <hydrodynamic_interactions> 0 </hydrodynamic_interactions>
  <time_step_tolerances>
    <minimum_core_dt> 0.2 </minimum_core_dt>
    <minimum_core_rxn_dt> 0.05 </minimum_core_rxn_dt>
    <minimum_chain_dt> 0 </minimum_chain_dt>
    <minimum_chain_rxn_dt> 0 </minimum_chain_rxn_dt>
  </time_step_tolerances>
  <molecule_info>
    <b_radius> {b_radius} </b_radius>
    <b_reaction_rate> {b_reaction_rate} </b_reaction_rate>
    <n_reactions> 1 </n_reactions>
  </molecule_info>
  <reactions>
    <n_trajectories> {n_trajectories} </n_trajectories>
    <stuck> 0 </stuck>
    <escaped> {escaped} </escaped>
    <completed>
      <name> {name1} </name>
      <n> {n1} </n>
    </completed>
    <completed>
      <name> {name2} </name>
      <n> {n2} </n>
    </completed>
  </reactions>
  <n_bd_steps> 0 </n_bd_steps>
</rates>
"""

import numpy as np

def write_b_surface_output_file(output_file_name, k_on_src, b_transition_probs,
                                milestone_transition_probs):
    CONV_FACTOR = 602000000.0
    b_reaction_rate = k_on_src / CONV_FACTOR
    N_TRAJS = 1000000000
    for key in b_transition_probs:
        if key == "escaped":
            escaped = int(b_transition_probs["escaped"] * N_TRAJS)
        else:
            key1 = key
            name1 = "b_%s" % key
            n1 = int(b_transition_probs[key] * N_TRAJS)

    for key in milestone_transition_probs:
        if key == "escaped":
            escaped2 = int(milestone_transition_probs["escaped"] * n1)
            pass
        else:
            name2 = "%s_%s" % (key1, key)
            n2 = int(milestone_transition_probs[key] * n1)

    new_result = bd_results_string.format(
        b_radius=0.0, b_reaction_rate=b_reaction_rate, n_trajectories=N_TRAJS,
        escaped=escaped+escaped2, name1=name1, n1=n1, name2=name2, n2=n2)
    #print("new_result:", new_result)
    with open(output_file_name, "w") as f:
        f.write(new_result)

    return

class BrownianParticle():
    def __init__(self, mass, dimensions, position, velocity, diffusion,
                 potential_energy_function):
        self.mass = mass
        self.dimensions = dimensions
        assert len(position.shape) == dimensions
        self.position = position
        self.velocity = velocity
        self.diffusion = diffusion
        self.potential_energy_function = potential_energy_function
        return

class ToyIntegrator():
    """
    Base class for toy integrators.
    """
    def __init__(self):
        pass

class SmoluchowskiSphericalMMVTIntegrator(ToyIntegrator):
    """
    An integrator that generates MMVT output files based on a Smoluchowski
    region.
    """
    def __init__(self, calc, index, output_file_name, style="openmm"):
        self.calc = calc
        self.index = index
        self.style = style.lower()
        self.output_file_name = output_file_name
        self._generate_output_file_header()
        self.in_state = 1
        self.time = 0.0
        self.bounce_counter = 0
        self.timestep = 0.002
        self.N_i_j_alpha_dict, self.R_i_alpha_dict, self.N_alpha_beta_dict, self.T_alpha \
            = self.calc.regions[self.index].produce_mmvt_statistics(self.index)
        return

    def _generate_output_file_header(self):
        if self.style=="openmm":
            header = "#\"Bounced boundary ID\",\"bounce index\",\"total time (ps)\n"
        elif self.style=="namd":
            header = "# NAMD TEST OUTPUT\n"
        else:
            header = "UNKNOWN OUTPUT STYLE"
        with open(self.output_file_name, "w") as f:
            f.write(header)
        return

    def _write_bounce_to_output_file(self, line):
        with open(self.output_file_name, "a") as f:
            f.write(line)
        return
    """ # TODO: remove if the other version is good enough
    def _write_random_transition(self, starting_step, num_steps):
        # TODO: perhaps this algorithm can be improved by not sampling times
        #  for both transitions and bounces, but rather using the number
        #  of transitions per bounce to decide when a bounce becomes a
        #  transition?
        MAX_ITER = 1e9
        step_time = 0.0
        N_i_j_alpha_dict, R_i_alpha_dict, N_alpha_beta_dict, T_alpha \
            = self.calc.regions[self.index].produce_mmvt_statistics(self.index)
        available_transition_dict = {}
        available_transitions = []
        transition_probabilities = []
        total_transitions_out_of_state = 0
        for key in N_i_j_alpha_dict:
            if key[0] == self.in_state:
                available_transition_dict[key[1]] = N_i_j_alpha_dict[key]
                total_transitions_out_of_state += N_i_j_alpha_dict[key]

        probability_sum = sum(available_transition_dict.values())
        for key in available_transition_dict:
            available_transitions.append(key)
            transition_probabilities.append(available_transition_dict[key]/probability_sum)

        if len(available_transitions) == 0:
            time_in_this_transition = 1e9
            avg_time_between_self_bounces = T_alpha / N_alpha_beta_dict[self.in_state]
        else:
            next_state = np.random.choice(a=np.array(available_transitions),
                                          p=np.array(transition_probabilities))
            avg_transition_time = R_i_alpha_dict[self.in_state] / total_transitions_out_of_state
            time_in_this_transition = np.random.exponential(avg_transition_time)
            avg_time_between_self_bounces = R_i_alpha_dict[self.in_state] / N_alpha_beta_dict[self.in_state]

        # sample times in while loop until time_between_transitions is exceeded
        counter = 0
        while True:
            time_in_this_self_bounce = np.random.exponential(avg_time_between_self_bounces)
            step_time += time_in_this_self_bounce
            if step_time > time_in_this_transition:
                break
            if self.style=="openmm":
                bounce_str = "{},{},{}\n".format(self.in_state, self.bounce_counter, self.time+step_time)

            #print(bounce_str)
            self._write_bounce_to_output_file(bounce_str)
            self.bounce_counter += 1
            if self.bounce_counter >= starting_step + num_steps:
                return
            if counter > MAX_ITER:
                raise Exception("Max iterations reached.")
            counter += 1

        self.time += time_in_this_transition
        self.in_state = next_state
        if self.style=="openmm":
            bounce_str = "{},{},{}\n".format(self.in_state, self.bounce_counter, self.time)
        #print(bounce_str)
        self._write_bounce_to_output_file(bounce_str)
        self.bounce_counter += 1
        return
    """

    def _write_random_transition(self, starting_step, num_steps):
        # TODO: perhaps this algorithm can be improved by not sampling times
        #  for both transitions and bounces, but rather using the number
        #  of transitions per bounce to decide when a bounce becomes a
        #  transition?
        MAX_ITER = 1e9
        step_time = 0.0

        available_transition_dict = {}
        available_transitions = []
        transition_probabilities = []
        total_transitions_out_of_state = 0
        for key in self.N_i_j_alpha_dict:
            if key[0] == self.in_state:
                available_transition_dict[key[1]] = self.N_i_j_alpha_dict[key]
                total_transitions_out_of_state += self.N_i_j_alpha_dict[key]

        probability_sum = sum(available_transition_dict.values())
        for key in available_transition_dict:
            available_transitions.append(key)
            transition_probabilities.append(available_transition_dict[key]/probability_sum)

        if len(available_transitions) == 0:
            time_in_this_transition = 1e9
            avg_time_between_self_bounces = self.T_alpha / self.N_alpha_beta_dict[self.in_state]
            bounces_per_transition = 1e9
        else:
            next_state = np.random.choice(a=np.array(available_transitions),
                                          p=np.array(transition_probabilities))
            avg_transition_time = self.R_i_alpha_dict[self.in_state] / total_transitions_out_of_state
            time_in_this_transition = np.random.exponential(avg_transition_time)
            avg_time_between_self_bounces = self.R_i_alpha_dict[self.in_state] / self.N_alpha_beta_dict[self.in_state]
            avg_bounces_per_transition = self.N_alpha_beta_dict[self.in_state] / total_transitions_out_of_state
            bounces_per_transition = np.random.exponential(avg_bounces_per_transition)

        # sample times in while loop until time_between_transitions is exceeded
        counter = 1
        while True:
            if counter > bounces_per_transition:
                break
            time_in_this_self_bounce = np.random.exponential(avg_time_between_self_bounces)
            step_time += time_in_this_self_bounce

            if self.style=="openmm":
                bounce_str = "{},{},{}\n".format(self.in_state, self.bounce_counter, self.time+step_time)
            elif self.style=="namd":
                if self.in_state == 1:
                    new_anchor = self.index - 1
                    new_milestone = self.index - 1
                elif self.in_state == 2:
                    new_anchor = self.index + 1
                    new_milestone = self.index

                if new_anchor == -1:
                    new_anchor = 1
                    new_milestone = self.index

                template="SEEKR: Cell Collision: current: {}, new: {}, stepnum: {}\n"
                bounce_str = template.format(self.index, new_anchor, int((self.time+step_time)/self.timestep))
                if self.bounce_counter == 0:
                    template="SEEKR: Milestone Transition: anchor: {}, source: {}, destination: {}, stepnum: {}, incubation steps: {}\n"
                    bounce_str += template.format(
                        self.index, "none", new_milestone,
                        int((self.time+step_time)/self.timestep),
                        int(step_time/self.timestep))

            self._write_bounce_to_output_file(bounce_str)
            self.bounce_counter += 1

            if self.bounce_counter >= starting_step + num_steps:
                return

            if counter > MAX_ITER:
                raise Exception("Max iterations exceeded.")
            counter += 1

        self.time += step_time
        prev_state = self.in_state
        self.in_state = next_state
        if self.style=="openmm":
            bounce_str = "{},{},{}\n".format(self.in_state, self.bounce_counter, self.time)
        elif self.style=="namd":
            if self.in_state == 1:
                new_anchor = self.index - 1
                old_milestone = self.index
                new_milestone = self.index - 1
            elif self.in_state == 2:
                new_anchor = self.index + 1
                old_milestone = self.index - 1
                new_milestone = self.index
            if self.bounce_counter == 0:
                old_milestone = "none"
            template="SEEKR: Cell Collision: current: {}, new: {}, stepnum: {}\n"\
                     +"SEEKR: Milestone Transition: anchor: {}, source: {}, destination: {}, stepnum: {}, incubation steps: {}\n"
            bounce_str = template.format(self.index, new_anchor, int((self.time)/self.timestep),
                                         self.index, old_milestone, new_milestone, int((self.time)/self.timestep), int(step_time/self.timestep))
        self._write_bounce_to_output_file(bounce_str)
        #print("bounce_str:", bounce_str)
        self.bounce_counter += 1
        return

    def step(self, number):
        starting_step = self.bounce_counter
        while self.bounce_counter < number:
            self._write_random_transition(starting_step, number)

        return

class SmoluchowskiSphericalElberIntegrator(ToyIntegrator):
    """
    An integrator that generates MMVT output files based on a Smoluchowski
    region.
    """
    def __init__(self, calc, index, output_file_name, style="openmm"):
        self.calc = calc
        self.index = index
        self.style = style.lower()
        self.output_file_name = output_file_name
        self._generate_output_file_header()
        self.in_state = index
        self.time = 0.0
        self.bounce_counter = 0
        self.timestep = 1.0
        self.elberN_ij, self.elberR_i = self.calc.produce_elber_statistics()
        return

    def _generate_output_file_header(self):
        if self.style=="openmm":
            header = "#\"Bounced boundary ID\",\"bounce index\",\"total time (ps)\n"
        else:
            header = "UNKNOWN OUTPUT STYLE"
        with open(self.output_file_name, "w") as f:
            f.write(header)
        return

    def _write_bounce_to_output_file(self, line):
        with open(self.output_file_name, "a") as f:
            f.write(line)
        return

    def _write_random_transition(self, starting_step, num_steps):
        # TODO: perhaps this algorithm can be improved by not sampling times
        #  for both transitions and bounces, but rather using the number
        #  of transitions per bounce to decide when a bounce becomes a
        #  transition?

        available_transition_dict = {}
        available_transitions = []
        transition_probabilities = []
        total_transitions_out_of_state = 0
        for key in self.elberN_ij:
            if key[0] == self.in_state:
                available_transition_dict[key[1]] = self.elberN_ij[key]
                total_transitions_out_of_state += self.elberN_ij[key]

        probability_sum = sum(available_transition_dict.values())
        for key in available_transition_dict:
            available_transitions.append(key)
            transition_probabilities.append(available_transition_dict[key]/probability_sum)

        next_state = np.random.choice(a=np.array(available_transitions),
                                        p=np.array(transition_probabilities))
        avg_transition_time = self.elberR_i[self.in_state]
        step_time = np.random.exponential(avg_transition_time)

        transition_alias = next_state - self.in_state + 2
        self.time += step_time

        if self.style=="openmm":
            bounce_str = "{},{},{}\n".format(transition_alias, self.bounce_counter, step_time)
        #print(bounce_str)
        self._write_bounce_to_output_file(bounce_str)
        self.bounce_counter += 1
        return

    def step(self, number):
        starting_step = self.bounce_counter
        while self.bounce_counter < number:
            self._write_random_transition(starting_step, number)

        return


from corvus.structures import Handler, Exchange, Loop, Update
import corvutils.pyparsing as pp
import os, sys, subprocess, shutil #, resource
import re
from scipy.interpolate import CubicSpline
from scipy.integrate import quad
from scipy.signal import convolve
import numpy as np
# Debug: FDV
import pprint

pp_debug = pprint.PrettyPrinter(indent=4)

# Define dictionary of implemented calculations
implemented = {}
strlistkey = lambda L:','.join(sorted(L))
subs = lambda L:[{L[j] for j in range(len(L)) if 1<<j&k} for k in range(1,1<<len(L))]
#for s in subs(['cell_vectors', 'cell_struct_xyz_red', 'cell_scaling_iso', 'cell_scaling_abc', 'number_density']):
#    key = strlistkey(s)
#    autodesc = 'Get ' + ', '.join(s) + ' using cif2cell'
#    cost = 10
#    implemented[key] = {'type':'Exchange','out':list(s),'req':['cif_input'],
#                        'desc':autodesc,'cost':cost}

implemented['mbxanes'] = {'type':'Exchange','out':['mbxanes'],'cost':0,
                        'req':['xanes_cfavg','spectralFunction'],'desc':'Calculate many-body xanes from xanes and spectral function.'}
#'req':['xanes','spectal_function'],'desc':'Calculate supercell from cif input.'}

class mbconv(Handler):
    def __str__(self):
        return 'mbconv Handler'

    @staticmethod
    def canProduce(output):
        if isinstance(output, list) and output and isinstance(output[0], str):
            return strlistkey(output) in implemented
        elif isinstance(output, str):
            return output in implemented
        else:
            raise TypeError('Output should be token or list of tokens')

    @staticmethod
    def requiredInputFor(output):
        if isinstance(output, list) and output and isinstance(output[0], str):
            unresolved = {o for o in output if not mbconv.canProduce(o)}
            canProduce = (o for o in output if mbconv.canProduce(o))
            additionalInput = (set(implemented[o]['req']) for o in canProduce)
            return list(set.union(unresolved,*additionalInput))
        elif isinstance(output, str):
            if output in implemented:
                return implemented[output]['req']
            else:
                return [output]
        else:
            raise TypeError('Output should be token or list of tokens')

    @staticmethod
    def cost(output):
        if isinstance(output, list) and output and isinstance(output[0], str):
            key = strlistkey(output)
        elif isinstance(output, str):
            key = output
        else:
            raise TypeError('Output should be token or list of tokens')
        if key not in implemented:
            raise LookupError('Corvus cannot currently produce ' + key + ' using FEFF')
        return implemented[key]['cost']

    @staticmethod
    def sequenceFor(output,inp=None):
        if isinstance(output, list) and output and isinstance(output[0], str):
            key = strlistkey(output)
        elif isinstance(output, str):
            key = output
        else:
            raise TypeError('Output should be token of list of tokens')
        if key not in implemented:
            raise LookupError('Corvus cannot currently produce ' + key + ' using FEFF')
        f = lambda subkey : implemented[key][subkey]
        required = f('req')
        # JJK - Need to add requirements of internal workflow here.
        if 'mbconv' in list(inp.keys()):
            required.extend()

        if f('type') is 'Exchange':
            return Exchange(mbconv, f('req'), f('out'), cost=f('cost'), desc=f('desc'))

    @staticmethod
    def prep(config):
        subdir = config['pathprefix'] + str(config['xcIndex']) + '_MBXANES'
        xcDir = os.path.join(config['cwd'], subdir)
        # Make new output directory if if doesn't exist
        if not os.path.exists(xcDir):
            os.mkdir(xcDir)
        # Store current Exchange directory in configuration
        config['xcDir'] = xcDir

    #@staticmethod
    #def setDefaults(input,target):

    @staticmethod
    def run(config, input, output):

        # Loop over targets in output.
        if 'mbxanes' in output:
            # In future use file_reader handler to read in XANES and spectral function if already calculated.
            w  = np.array(input.get('xanes_cfavg')[0])
            mu0= np.array(input.get('xanes_cfavg')[1])
            wsf= np.flip(-1.0*np.array(input.get('spectralFunction')[0]))
            sf = np.flip(np.array(input.get('spectralFunction')[1]))
            # Interpolate both XANES and spectral function onto an even grid
            #w, mu0 = np.loadtxt('xanes.dat',usecols = (0,1)).T
            #wsf,sf = np.loadtxt('spfcn.dat',usecols = (0,1)).T
            min_diff = np.amin(np.ediff1d(w))
            min_diff = min(min_diff,np.amin(np.ediff1d(wsf)))

            mu0_cs = CubicSpline(w,mu0)
            spfcn_cs = CubicSpline(wsf,sf)
            # Use larger of two ranges to specify range
            w_terp = np.arange(w[0],w[-1],min_diff)
            wsf_terp = np.arange(wsf[0],wsf[-1],min_diff)
            mu0_terp = mu0_cs(w_terp)
            spfcn_terp = spfcn_cs(wsf_terp)

            mu_mb = convolve(mu0_terp,spfcn_terp,mode='full')*min_diff

            # If extra broadening is requested, perform a convolution of that as well.
            if 'mbconv.extra_broadening' in input:
                gam = input['mbconv.extra_broadening'][0][0]
                A_br = gam/np.pi*1.0/(wsf_terp**2 + gam**2)
                mu_mb = np.convolve(mu_mb,A_br,mode='same')*min_diff

            scale=w_terp[-1] - w_terp[0] + wsf_terp[-1] - wsf_terp[0]
            first = w_terp[0] + wsf_terp[0]
            w_terp = np.linspace(0.0,scale,mu_mb.size)
            w_terp = w_terp + first
            mu0_terp = mu0_cs(w_terp)
            output['mbxanes'] = [w_terp,mu_mb]
            np.savetxt('mbxanes.dat',np.array([w_terp, mu_mb, mu0_terp]).transpose())

    @staticmethod
    def cleanup(config):
        pass