#! /usr/bin/env python3
#
#  Copyright 2018 California Institute of Technology
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# ISOFIT: Imaging Spectrometer Optimal FITting
# Author: David R Thompson, david.r.thompson@jpl.nasa.gov
#

import logging
import numpy as np
from scipy.interpolate import interp1d
from scipy.signal import convolve
from scipy.io import loadmat

from .common import eps, spectral_response_function, load_wavelen, resample_spectrum, emissive_radiance
from isofit.configs import Config

### Variables ###

# Max. wavelength difference (nm) that does not trigger expensive resampling
wl_tol = 0.01

### Classes ###

class Instrument:

    def __init__(self, full_config: Config):
        """A model of the spectrometer instrument, including spectral
        response and noise covariance matrices. Noise is typically calculated
        from a parametric model, fit for the specific instrument.  It is a
        function of the radiance level."""

        config = full_config.forward_model.instrument

        # If needed, skip first index column and/or convert to nanometers
        self.wl_init, self.fwhm_init = load_wavelen(config.wavelength_file)
        self.n_chan = len(self.wl_init)

        self.fast_resample = config.fast_resample

        self.bounds = config.statevector.get_all_bounds()
        self.scale = config.statevector.get_all_scales()
        self.init = config.statevector.get_all_inits()
        self.prior_mean = np.array(config.statevector.get_all_prior_means())
        self.prior_sigma = np.array(config.statevector.get_all_prior_sigmas())
        self.statevec_names = config.statevector.get_element_names()
        self.n_state = len(self.statevec_names)

        if config.SNR is not None:
            self.model_type = 'SNR'
            self.snr = config.SNR
        elif config.parametric_noise_file is not None:
            self.model_type = 'parametric'
            self.noise_file = config.parametric_noise_file
            coeffs = np.loadtxt(
                self.noise_file, delimiter=' ', comments='#')
            p_a, p_b, p_c = [interp1d(coeffs[:, 0], coeffs[:, col],
                                      fill_value='extrapolate') for col in (1, 2, 3)]
            self.noise = np.array([[p_a(w), p_b(w), p_c(w)]
                                   for w in self.wl_init])
            self.integrations = config.integrations

        elif config.pushbroom_noise_file is not None:
            self.model_type = 'pushbroom'
            self.noise_file = config.pushbroom_noise_file
            D = loadmat(self.noise_file)
            self.ncols = D['columns'][0, 0]
            if self.n_chan != np.sqrt(D['bands'][0, 0]):
                logging.error('Noise model mismatches wavelength # bands')
                raise ValueError('Noise model mismatches wavelength # bands')
            cshape = ((self.ncols, self.n_chan, self.n_chan))
            self.covs = D['covariances'].reshape(cshape)
            self.integrations = config.integrations

        elif config.nedt_noise_file is not None:
            self.model_type = 'NEDT'
            self.noise_file = config.nedt_noise_file
            self.noise_data = np.loadtxt(
                self.noise_file, delimiter=',', skiprows=8)
            noise_data_w_nm = self.noise_data[:, 0] * 1000
            noise_data_NEDT = self.noise_data[:, 1]
            nedt = interp1d(noise_data_w_nm, noise_data_NEDT)(self.wl_init)

            T, emis = 300., 0.95  # From Glynn Hulley, 2/18/2020
            _, drdn_dT = emissive_radiance(emis, T, self.wl_init)
            self.noise_NESR = nedt * drdn_dT

        else:
            raise IndexError('Please define the instrument noise.')
            # This should never be reached, as an error is designated in the config read

        # We track several unretrieved free variables, that are specified
        # in a fixed order (always start with relative radiometric
        # calibration)
        self.bvec = ['Cal_Relative_%04i' % int(w) for w in self.wl_init] + \
            ['Cal_Spectral', 'Cal_Stray_SRF']
        self.bval = np.zeros(self.n_chan+2)

        if config.unknowns is not None:

            # First we take care of radiometric uncertainties, which add
            # in quadrature.  We sum their squared values.  Systematic
            # radiometric uncertainties account for differences in sampling
            # and radiative transfer that manifest predictably as a function
            # of wavelength.
            if config.unknowns.channelized_radiometric_uncertainty_file is not None:
                f = config.unknowns.channelized_radiometric_uncertainty_file
                u = np.loadtxt(f, comments='#')
                if (len(u.shape) > 0 and u.shape[1] > 1):
                    u = u[:, 1]
                self.bval[:self.n_chan] = self.bval[:self.n_chan] + \
                    pow(u, 2)

            # Uncorrelated radiometric uncertainties are consistent and
            # independent in all channels.
            if config.unknowns.uncorrelated_radiometric_uncertainty is not None:
                u = config.unknowns.uncorrelated_radiometric_uncertainty
                self.bval[:self.n_chan] = self.bval[:self.n_chan] + \
                    pow(np.ones(self.n_chan) * u, 2)

            # Radiometric uncertainties combine via Root Sum Square...
            # Be careful to avoid square roots of zero!
            small = np.ones(self.n_chan)*eps
            self.bval[:self.n_chan] = np.maximum(self.bval[:self.n_chan], small)
            self.bval[:self.n_chan] = np.sqrt(self.bval[:self.n_chan])

            # Now handle spectral calibration uncertainties
            if config.unknowns.wavelength_calibration_uncertainty is not None:
                self.bval[-2] = config.unknowns.wavelength_calibration_uncertainty
            if config.unknowns.stray_srf_uncertainty is not None:
                self.bval[-1] = config.unknowns.stray_srf_uncertainty

        # Determine whether the calibration is fixed.  If it is fixed,
        # and the wavelengths of radiative transfer modeling and instrument
        # are the same, then we can bypass computationally expensive sampling
        # operations later.
        self.calibration_fixed = True
        if config.statevector.GROW_FWHM is not None or config.statevector.WL_SHIFT is not None or \
                config.statevector.WL_SPACE is not None:
            self.calibration_fixed = False

    def xa(self):
        """Mean of prior distribution, calculated at state x."""

        return self.init.copy()

    def Sa(self):
        """Covariance of prior distribution (diagonal)."""

        if self.n_state == 0:
            return np.zeros((0, 0), dtype=float)
        return np.diagflat(np.power(self.prior_sigma, 2))

    def Sy(self, meas, geom):
        """Calculate measuremment error covariance.  Kelvin Man Yiu Leung and
            Jayanth Jagalur Mohan (MIT) developed the noise clipping strategy.

        Input: meas, the instrument measurement
        Returns: Sy, the measurement error covariance due to instrument noise
        """
        if self.model_type == 'SNR':
            nedl = (1.0 / self.snr) * meas
            minimum_noise = np.sqrt(1e-7)
            bad = nedl < minimum_noise
            if np.any(bad):
                logging.debug('SNR noise model found noise <= 0 - adjusting to slightly positive to avoid /0.')
            nedl[bad] = minimum_noise
            return np.diagflat(np.power(nedl,2))

        elif self.model_type == 'parametric':
            noise_plus_meas = self.noise[:, 1]+meas
            if np.any(noise_plus_meas <=0):
                noise_plus_meas[noise_plus_meas <= 0] = 1e-5
                logging.debug('Parametric noise model found noise <= 0 - adjusting to slightly positive to avoid /0.')
            nedl = np.abs(self.noise[:, 0]*np.sqrt(noise_plus_meas)+self.noise[:, 2])
            nedl = nedl/np.sqrt(self.integrations)
            return np.diagflat(np.power(nedl,2))

        elif self.model_type == 'pushbroom':
            if geom.pushbroom_column is None:
                C = np.squeeze(self.covs.mean(axis=0))
            else:
                C = self.covs[geom.pushbroom_column, :, :]
            return C / np.sqrt(self.integrations)

        elif self.model_type == 'NEDT':
            return np.diagflat(np.power(self.noise_NESR,2))

    def dmeas_dinstrument(self, x_instrument, wl_hi, rdn_hi):
        """Jacobian of measurement with respect to the instrument
           free parameter state vector. We use finite differences for now."""

        dmeas_dinstrument = np.zeros((self.n_chan, self.n_state), dtype=float)
        if self.n_state == 0:
            return dmeas_dinstrument

        meas = self.sample(x_instrument, wl_hi, rdn_hi)
        for ind in range(self.n_state):
            x_instrument_perturb = x_instrument.copy()
            x_instrument_perturb[ind] = x_instrument_perturb[ind]+eps
            meas_perturb = self.sample(x_instrument_perturb, wl_hi, rdn_hi)
            dmeas_dinstrument[:, ind] = (meas_perturb - meas) / eps
        return dmeas_dinstrument

    def dmeas_dinstrumentb(self, x_instrument, wl_hi, rdn_hi):
        """Jacobian of radiance with respect to the instrument parameters
        that are unknown and not retrieved, i.e., the inevitable persisting
        uncertainties in instrument spectral and radiometric calibration.

        Input: meas, a vector of size n_chan
        Returns: Kb_instrument, a matrix of size [n_measurements x nb_instrument]
        """

        # Uncertainty due to radiometric calibration
        meas = self.sample(x_instrument, wl_hi, rdn_hi)
        dmeas_dinstrument = np.hstack(
            (np.diagflat(meas), np.zeros((self.n_chan, 2))))

        # Uncertainty due to spectral calibration
        if self.bval[-2] > 1e-6:
            dmeas_dinstrument[:, -2] = self.sample(x_instrument, wl_hi,
                                                   np.hstack((np.diff(rdn_hi), np.array([0]))))

        # Uncertainty due to spectral stray light
        if self.bval[-1] > 1e-6:
            ssrf = spectral_response_function(np.arange(-10, 11), 0, 4)
            blur = convolve(meas, ssrf, mode='same')
            dmeas_dinstrument[:, -1] = blur - meas

        return dmeas_dinstrument

    def sample(self, x_instrument, wl_hi, rdn_hi):
        """Apply instrument sampling to a radiance spectrum, returning predicted measurement."""

        if self.calibration_fixed and (len(self.wl_init) == len(wl_hi)) and \
                    all((self.wl_init - wl_hi) < wl_tol):
            return rdn_hi
        wl, fwhm = self.calibration(x_instrument)
        if rdn_hi.ndim == 1:
            return resample_spectrum(rdn_hi, wl_hi, wl, fwhm)
        else:
            resamp = []
            # The "fast resample" option approximates a complete resampling
            # by a convolution with a uniform FWHM.
            if self.fast_resample:
                for i, r in enumerate(rdn_hi):
                    ssrf = spectral_response_function(np.arange(-10, 11), 0, fwhm[0])
                    blur = convolve(r, ssrf, mode='same')
                    resamp.append(interp1d(wl_hi, blur)(wl))
            else:
                for i, r in enumerate(rdn_hi):
                    r2 = resample_spectrum(r, wl_hi, wl, fwhm)
                    resamp.append(r2)
            return np.array(resamp)

    def simulate_measurement(self, meas, geom):
        """Simulate a measurement by the given sensor, for a true radiance
        sampled to instrument wavelengths. This basically just means
        drawing a sample from the noise distribution."""

        Sy = self.Sy(meas, geom)
        mu = np.zeros(meas.shape)
        rdn_sim = meas + np.random.multivariate_normal(mu, Sy)
        return rdn_sim

    def calibration(self, x_instrument):
        """Calculate the measured wavelengths."""

        wl, fwhm = self.wl_init, self.fwhm_init
        space_orig = wl - wl[0]
        offset = wl[0]
        if 'GROW_FWHM' in self.statevec_names:
            ind = self.statevec_names.index('GROW_FWHM')
            fwhm = fwhm + x_instrument[ind]

        if 'WL_SPACE' in self.statevec_names:
            ind = self.statevec_names.index('WL_SPACE')
            space = x_instrument[ind]
        else:
            space = 1.0

        if 'WL_SHIFT' in self.statevec_names:
            ind = self.statevec_names.index('WL_SHIFT')
            shift = x_instrument[ind]
        else:
            shift = 0.0

        wl = offset + shift + space_orig * space
        return wl, fwhm

    def summarize(self, x_instrument, geom):
        """Summary of state vector."""

        if len(x_instrument) < 1:
            return ''
        return 'Instrument: '+' '.join(['%5.3f' % xi for xi in x_instrument])


import torch
import numpy as np
from torch.nn import functional as F

# =================================-----FOR DATASET------============================

def gaussian2D(shape, sigma=1):
    '''
    对输入shape(radius,radius)生产一个高斯核
    :param shape: (diameter,diameter)
    :param sigma:
    :return: (radius*2+1,radius*2+1)
    '''
    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m + 1, -n:n + 1]
    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
    # 过于小的数设置为0
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h

def draw_umich_gaussian(heatmap, center, radius, k=1):
    '''
    在heatmap上对中心点center半径为radius画高斯分布
    :param heatmap: 特征图(128*128)
    :param center: (x,y)
    :param radius: 半径
    :param k:
    :return:
    '''
    # 从中心开始扩展,长度为2 * radius + 1,文章说是sigma=radius/3,
    # 这里也解决了radius=0，导致sigma除数为0的问题
    diameter = 2 * radius + 1
    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)

    x, y = center[0], center[1]

    height, width = heatmap.shape[0:2]

    # 越界处理
    left, right = min(x, radius), min(width - x, radius + 1)
    top, bottom = min(y, radius), min(height - y, radius + 1)

    # 对齐处理
    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:  # TODO debug
        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
    return heatmap

def gaussian_radius(det_size, min_overlap=0.7):
    '''
    求高斯半径
    方法来自于CornerNet:https://arxiv.org/pdf/1808.01244.pdf
    原理就是对三种情况(1内扩1外扩,2内扩,2外扩)解一元二次方程:https://github.com/princeton-vl/CornerNet/issues/110
    :param det_size: bbox在特征图的大小(h,w)
    :param min_overlap: 最小的IOU
    :return: 最小的半径，其保证iou>=min_overlap
    '''
    height, width = det_size

    a1 = 1
    b1 = (height + width)
    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
    r1 = (b1 + sq1) / 2

    a2 = 4
    b2 = 2 * (height + width)
    c2 = (1 - min_overlap) * width * height
    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
    r2 = (b2 + sq2) / 2

    a3 = 4 * min_overlap
    b3 = -2 * min_overlap * (height + width)
    c3 = (min_overlap - 1) * width * height
    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
    r3 = (b3 + sq3) / 2
    return min(r1, r2, r3)

# =================================-----FOR TEST-----============================

def hm_topk(hm, k):
    # 使用max_pool获取峰点
    batch, cls, h, w = hm.size()
    out = F.max_pool2d(hm, 3, 1, 1)
    keep_max = (out == hm).float()
    hm = keep_max * hm
    # 在heatmap中取每个类别的topk(hm经sigmoid后) topk_indexs的值在0~h*w
    topk_scores, topk_indexs = hm.view(batch, cls, -1).topk(k)  # (batch,cls,k)
    # 所有类别取得分最高的topk
    topk_scores, topk_ind = topk_scores.view(batch, -1).topk(k)  # (batch,k)
    # 获取得分最高的类别topk,topk_scores每个类别有得分最高的k个，topk_ind除k取下整即为class
    topk_cls = topk_ind // k
    # 若topk_indexs.size=(batch,cls*k),topk_indexs[topk_ind]即为最终的index
    topk_indexs = topk_indexs.view(batch, -1).gather(1, topk_ind)
    # 获取所有类别中最高得分topk_indexs对应的横纵坐标,即一维转二维
    topk_ys, topk_xs = topk_indexs // w, topk_indexs % w
    return topk_scores, topk_indexs, topk_cls, topk_xs, topk_ys

def heatmap_bbox(hm, wh, reg, k=100):
    scores, indexs, cls, xs, ys = hm_topk(hm.sigmoid_(), k)
    batch = reg.size(0)
    # 先转置便于取关键点对应的2个偏移量
    reg = reg.view(batch, 2, -1).transpose(2, 1).contiguous()  # (batch,w*h,2)
    reg_indexs = indexs.unsqueeze(2).expand(batch, -1, 2)  # (batch,k,2)
    reg = reg.gather(1, reg_indexs)  # (batch,k,2)
    xs = xs.float() + reg[:, :, 0]
    ys = ys.float() + reg[:, :, 1]
    # wh via reg_indexs
    wh = wh.view(batch, 2, -1).transpose(2, 1).contiguous().gather(1, reg_indexs)  # ((batch,k,2)
    # bbox via xs and wh
    bbox = xs - wh[:, :, 0] / 2, ys - wh[:, :, 1] / 2, xs + wh[:, :, 0] / 2, ys + wh[:, :, 1] / 2
    bbox = torch.stack(bbox, -1)  # (batch,k,4)
    return bbox, cls, scores

def area_of(left_top, right_bottom) -> torch.Tensor:
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N, 2): left top corner.
        right_bottom (N, 2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    hw = torch.clamp(right_bottom - left_top, min=0.0)
    return hw[..., 0] * hw[..., 1]

def iou_of(boxes0, boxes1, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        boxes0 (N, 4): ground truth boxes.
        boxes1 (N or 1, 4): predicted boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2])
    overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:])

    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)

def soft_nms(box_scores, score_threshold=0.5, sigma=0.5, top_k=-1):
    """Soft NMS implementation.

    References:
        https://arxiv.org/abs/1704.04503
        https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/cython_nms.pyx
        https://oldpan.me/archives/write-hard-nms-c

    Args:
        box_scores (N, 6): boxes in corner-form and probabilities. [x1,y1,x2,y2,cls,score]
        score_threshold: boxes with scores less than value are not considered.
        sigma: the parameter in score re-computation.
            scores[i] = scores[i] * exp(-(iou_i)^2 / simga)
        top_k: keep top_k results. If k <= 0, keep all the results.
    Returns:
         picked_box_scores (K, 5): results of NMS.
    """
    picked_box_scores = []
    while box_scores.size(0) > 0:
        max_score_index = torch.argmax(box_scores[:, -1])
        cur_box_prob = box_scores[max_score_index, :].clone().detach()
        picked_box_scores.append(cur_box_prob)
        if len(picked_box_scores) == top_k > 0 or box_scores.size(0) == 1:
            break
        cur_box = cur_box_prob[:-2]
        box_scores[max_score_index, :] = box_scores[-1, :]
        box_scores = box_scores[:-1, :]
        ious = iou_of(cur_box.unsqueeze(0), box_scores[:, :-2])

        box_scores[:, -1] = box_scores[:, -1] * torch.exp(-(ious * ious) / sigma)

        box_scores = box_scores[box_scores[:, -1] > score_threshold, :]
    if len(picked_box_scores) > 0:
        return torch.stack(picked_box_scores)
    else:
        return torch.tensor([])


# -*- coding: utf-8 -*-
"""
SMPTE 240M Colourspace
======================

Defines the *SMPTE 240M* colourspace:

-   :attr:`colour.models.SMPTE_240M_COLOURSPACE`.

References
----------
-   :cite:`SocietyofMotionPictureandTelevisionEngineers1999b` : Society of
    Motion Picture and Television Engineers. (1999). ANSI/SMPTE 240M-1995 -
    Signal Parameters - 1125-Line High-Definition Production Systems. Retrieved
    from http://car.france3.mars.free.fr/HD/INA- 26 jan 06/\
SMPTE normes et confs/s240m.pdf
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models.rgb import (RGB_Colourspace, normalised_primary_matrix,
                               oetf_SMPTE240M, eotf_SMPTE240M)

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

__all__ = [
    'SMPTE_240M_PRIMARIES', 'SMPTE_240M_WHITEPOINT_NAME',
    'SMPTE_240M_WHITEPOINT', 'SMPTE_240M_TO_XYZ_MATRIX',
    'XYZ_TO_SMPTE_240M_MATRIX', 'SMPTE_240M_COLOURSPACE'
]

SMPTE_240M_PRIMARIES = np.array([
    [0.6300, 0.3400],
    [0.3100, 0.5950],
    [0.1550, 0.0700],
])
"""
*SMPTE 240M* colourspace primaries.

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

SMPTE_240M_WHITEPOINT_NAME = 'D65'
"""
*SMPTE 240M* colourspace whitepoint name.

SMPTE_240M_WHITEPOINT_NAME : unicode
"""

SMPTE_240M_WHITEPOINT = (ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
    SMPTE_240M_WHITEPOINT_NAME])
"""
*SMPTE 240M* colourspace whitepoint.

SMPTE_240M_WHITEPOINT : ndarray
"""

SMPTE_240M_TO_XYZ_MATRIX = normalised_primary_matrix(SMPTE_240M_PRIMARIES,
                                                     SMPTE_240M_WHITEPOINT)
"""
*SMPTE 240M* colourspace to *CIE XYZ* tristimulus values matrix.

SMPTE_240M_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_SMPTE_240M_MATRIX = np.linalg.inv(SMPTE_240M_TO_XYZ_MATRIX)
"""
*CIE XYZ* tristimulus values to *SMPTE 240M* colourspace matrix.

XYZ_TO_SMPTE_240M_MATRIX : array_like, (3, 3)
"""

SMPTE_240M_COLOURSPACE = RGB_Colourspace(
    'SMPTE 240M',
    SMPTE_240M_PRIMARIES,
    SMPTE_240M_WHITEPOINT,
    SMPTE_240M_WHITEPOINT_NAME,
    SMPTE_240M_TO_XYZ_MATRIX,
    XYZ_TO_SMPTE_240M_MATRIX,
    oetf_SMPTE240M,
    eotf_SMPTE240M,
)
SMPTE_240M_COLOURSPACE.__doc__ = """
*SMPTE 240M* colourspace.

References
----------
:cite:`SocietyofMotionPictureandTelevisionEngineers1999b`,

SMPTE_240M_COLOURSPACE : RGB_Colourspace
"""


"""
Fluorescence Intensity calculations:
   Y. Choi

calculates fluorescence intensities for multiple elements at a fixed incident x-ray energy.
measured fluorescence depends on net fluorescence yield and transmission.
included are:
    (with incident energy dependence): atomic crosssection, fluorescence yield for each edge,
                                     transition probability for each emission line,
    (with emitted fluorescence energy dependence): transmission to vortex detector, detector efficiency.
global variables: re, Barn, Nav, pre_edge_margin, fluo_emit_min, det_res
AtNum2f1f2Xect: needs datafile readf1f2a.py which is from chantler table
cal_NetYield2: for XRF calculations.  major and minor lines are calculated
                output as 'Elemental_Sensitivity.txt'

sim_spectra: keeps individual emission lines with significant intensity without weight-averaging
             each element can have a different relative concentration.
             makes strongest emission to 100 and scales the rest accordingly.
             output as 'simSpectrum.txt'
             runs sim_GaussPeaks() to emulate measure intensity as gaussian peak
             output as 'simSpectrum.plot'
             runs cal_NetYield2()
5/24/2010:  self-absorption effect included, 4 options (top, bottom, all, surface)
            Each option makes different elemental distribution,
            but substrate elements are always distributed evenly.
            get_ChemName(), nominal_density() are in readf1f2a.py
5/26/2010:   incident x-ray attenuation effect added, and incident angle effect is also added.
            refraction effect added to the angle, below critical angle, th=1e-10.
6/4/2010:    SampleMatrix2 has two layers and each layer can have different composition
            Class ElemFY is for fluorescence elements with AtSym, Conc, tag attribues,
            inputs in SimSpect, Cal_NetFyield2 changed.
            Concentrations are normalized to the concentration of the substrate1 molecule.
            Ex) with CaCO3[2.71g/cc]/Fe2O3[5.26g/cc] substrate, Ca concentration is 1.0 meaning 1.63e22 Ca/cc
            and Fe concentration is 2.43 = 2*1.98e22/1.63e22 where 1.98e22 is number density (/cc) of Fe2O3
6/8/2010:    now using python2.6 instead of 2.5
4/14/2011:  fluo_det.py: detection dependent part in fluorescence yield.
            Original fluo.py was modified to fluo_new.py.
            readf1f2a.py is used instead of original readf1f2.py
            fluo_new.py is being split into two.
            fluo_elem.py is for elemtnal dependence: fluorescenc yield, crosssection.
            fluo_det.py is for detection dependnece: detector efficiency, attenuation, sample
            * for attenuation, total cross-section([4]) is used instead of photoelectric ([2]).
            ** for fluorescence yield, photoelectric is used.
"""

import math
import numpy
import sys
from larch import use_plugin_path
use_plugin_path('xray')
use_plugin_path('xsw')

from physical_constants import AVOGADRO, BARN

pre_edge_margin=150.    # FY calculated from 150 eV below the absorption edge.
fluo_emit_min=500.      # minimum energy for emitted fluorescence.  ignore fluorescence emissions below 500eV
det_res=100.            # detector resoltuion in eV, used in sim_GaussPeaks

'''
----------------------------------------------------------------------------------------------------------
class: Material, ElemFY, SampleMatrix2
----------------------------------------------------------------------------------------------------------
'''
class Material:
    def __init__(self, composition, density, thickness=1,_larch=None):
        self.composition=composition        # ex) SiO2
        self.density=density              # in g/cm^3
        self.thickness=thickness            # in cm
        self.la=1.0                         #absoprtion length in cm
        self.trans=0.5                      # transmission.
        self.absrp=0.5                 # absorption
        self.delta=0.1                  # index of refraction, real part
        self.beta=0.1                   # index of refraction, imaginary part
        # MN replace
        elements = chemparse(composition)
        AtomList = elements.keys()
        AtomIndex = elements.values()
        AtomWeight=0
        for (ii, atom) in enumerate(AtomList):
            # MN replace...
            AtWt= atomic_mass(atom)
            index=AtomIndex[ii]
            AtomWeight = AtomWeight + index*AtWt
        NumberDensity=density*AVOGARDRO/AtomWeight
        self.NumDen=NumberDensity       # number of molecules per cm^3
        self.AtWt=AtomWeight            # weight per mole
    def getLa(self, energy, NumLayer=1.0):    # get absorption legnth for incident x-ray, NumLayer for multiple layers
        # MN replace...
        temp= xray_delta_beta(self.composition, self.density, energy, _larch=_larch)
        # returns delta, beta, la_photoE, Nat, la_total
        self.delta=temp[0];     self.beta=temp[1]
        self.la=temp[2]  # temp[4] (total) instead of temp[2] (photoelectric)
        if NumLayer<0:  NumLayer=0     # Number of layers cannot be less than zero
        self.trans=math.exp(-self.thickness*NumLayer/self.la)  #la attenuation length cm, NumLayer for multiple layers
        self.absrp=1-self.trans

class ElemFY:  # fluorescing element
    def __init__(self, AtomicSymbol='Fe', Concentration=10e-6, tag='dilute'):
        temp=AtomicSymbol[0].upper()  # capitalize the first letter
        self.AtSym=temp+AtomicSymbol[1:]
        self.Conc=Concentration
        self.tag=tag  # dilute, substrate1, substrate2

class SampleMatrix2:  # sample matrix for self-absorption correction, 6/3: two layers with different compositions
    def __init__(self, composition1='Si', density1=2.33, thickness1=0.1, \
                       composition2='Si', density2=2.33, thickness2=0.1,
                       angle0=45.0, option='surface', _larch=None):
        self.composition1 = composition1     # ex) Fe2O3
        # MN replace:
        out= chemparse(composition1)         # output from get_ChemName
        self.ElemList1 = out.keys()          # list of elments in matrix: 'Fe', 'O' for Fe2O3
        self.ElemInd1 = out.values()         # list of index: 2, 3 for Fe2O3
        # self.ElemFrt1 = out[2]             # list of fraction: 0.4, 0.6 for Fe2O3
        self.density1 = density1             # in g/cm^3
        self.thickness1 = thickness1         # top layer thickness in cm
        self.composition2 = composition2
        # MN replace with chemparse()
        out=chemparse(composition2)
        self.ElemList2 = out.keys()              # for the bottom substrate
        self.ElemInd2 = out.values()
        # self.ElemFrt2 = out[2]
        self.density2 = density2
        self.thickness2 = thickness2         # bottom layer thickness in cm
        self.angle = angle0*math.pi/180.     # in radian, incident beam angle, surface normal =pi/2
        self.option = option                 # option for fluorescing element location, surface/top/bottom
        self.la1 = 1.0                       # absoprtion length in cm
        self.delta1 = 0.1                    # index of refraction, real part correction
        self.beta1 = 0.1                     # index of refraction, imaginary part
        self.Nat1 = 1.0                      # atomic number density of the Fe2O3, 1e22 Fe2O3 atoms /cm^3
        self.la2 = 1.0
        self.delta2 = 0.1
        self.beta2 = 0.1
        self.Nat2 = 1.0                     # atomic number density of the first element, 1e22 Fe2O3 atoms/cm^3
        self.scale = 1.0                    # weighted average over the depth range: sum of (trans*factors)
        self.scale1 = 1.0                   # sum of (trans*1.0) this is for substrate element in top layer
        self.scale2 = 1.0                   # sum of (trans*thickness2/thickness1) for substrate element in bottom layer
        self.ElemListFY =[]                 # fluorescing element
        self.ElemFrtFY =[]                  # fraction for fluorescing element
        self.Nat1=1                  # atomic number density in atoms/cc, for example # of Fe2O3/cm^3 for Fe2O3 substrate
        self.Nat2=1                  # atomic number density in atoms/cc
        substrate1_material = Material(composition1, density1, _larch=_larch)
        AtNumDen1 = substrate1_material.NumDen      #  substrate1
        substrate2_material = Material(composition2, density2, _larch=_larch)
        AtNumDen2 = substrate2_material.NumDen      #  substrate2, fixed 8/12/10
        self.Nat1=AtNumDen1
        self.Nat2=AtNumDen2
        self.txt=''
        text1='substrate1:%6.3e %s/cm^3' % (AtNumDen1, composition1)
        #print(text1)
        text2='substrate2:%6.3e %s/cm^3' % (AtNumDen2, composition2)
        self.txt=text1+'  '+text2
        print(self.txt)
        # atom.conc is normalized to the number density of substrate1 molecule
        for (ii, item) in enumerate(self.ElemList1):
            # MN replace:
            if atomic_number(item)>=12:       # ignore elements below Mg
                #atom=ElemFY(item, self.ElemFrt1[ii], 'substrate1')
                atom=ElemFY(item, self.ElemInd1[ii], 'substrate1')
                # eg. for Fe2O3, Fe concentration is 2 (=2 x Fe2O3 number density)
                self.ElemListFY.append(atom)
        for (ii, item) in enumerate(self.ElemList2):
            # MN replace:
            if atomic_number(item)>=12:       # ignore elements below Mg
                #atom=ElemFY(item, self.ElemFrt2[ii], 'substrate2')
                atom=ElemFY(item, self.ElemInd2[ii]*AtNumDen2/AtNumDen1, 'substrate2')
                self.ElemListFY.append(atom)
        numLayer=100                        # each layer is sliced into 100 sublayers.
        self.depths=[]                      # depth values for calculation
        for ii in range(numLayer):
            step=1.0/numLayer
            self.depths.append(step*ii*self.thickness1)
        for ii in range(numLayer):
            step=1.0/numLayer
            self.depths.append(step*ii*self.thickness2+self.thickness1)
        self.factors=[]                     # 1 or pre if element present at each depth
        pre=self.thickness2/self.thickness1 # prefactor, if two layers have different thickness
        if self.option=='surface':          # fluorescecing atoms present in only on the surface
            for ii in range(len(self.depths)):
                if ii==0:
                    self.factors.append(1.0)
                else:
                    self.factors.append(0.0)
        if self.option=='all':              # fluorescecing atoms present throughout
            for ii in range(len(self.depths)):
                if self.depths[ii]<self.thickness1:
                    self.factors.append(1.0)
                else:
                    self.factors.append(pre)
        if self.option=='top':              # fluorescecing atoms present only in the top layer
            for ii in range(len(self.depths)):
                if self.depths[ii]<self.thickness1:
                    self.factors.append(1.0)
                else:
                    self.factors.append(0.0)
        if self.option=='bottom':           # fluorescecing atoms present only in the bottom layer
            for ii in range(len(self.depths)):
                if self.depths[ii]<self.thickness1:
                    self.factors.append(0.0)
                else:
                    self.factors.append(pre)
        # note: fluorescing substrate atoms present throughout regardless of the option.
        self.trans=[]                       # transmission to surface for emitted fluorescence
        self.absrp=[]                       # absorption until surface for emitted fluorescence
        self.inten0=[]                      # incident x-ray intensity at each depth
        for ii in range(len(self.depths)):
            self.trans.append(0.5)
            self.absrp.append(0.5)
            self.inten0.append(0.5)
    def getPenetration(self, energy0):  # incident x-ray penetration(attenuation)
        # refraction at air/top layer interface
        # MN replace:
        temp=f1f2.get_delta(self.composition1, self.density1, energy0)  # energy0: incident x-ray energy
        delta1=temp[0];  beta1=temp[1]
        la1=temp[4]     # in cm, temp[4] instead of temp[2], using total instead of photoelectric
        self.la1=la1                                        # absorption length in microns at incident x-ray energy
        angle_critical1 = (2.0*delta1)**(0.5)               # in radian, critical angle for total external reflection
        if angle_critical1>=self.angle:                     # below critical angle, the corrected should be zero.
            angle_corrected1=1.0e-15                        # a smaller number instead of zero
        else:                                               # above critical angle
            angle_corrected1 = (self.angle**2.0 - angle_critical1**2.0)**(0.5)  # in radian
        # refraction at top/bottom layers interface
        # MN replace:
        temp=f1f2.get_delta(self.composition2, self.density2, energy0)  # energy0: incident x-ray energy
        delta2=temp[0];  beta2=temp[1];
        la2=temp[4]     # in cm, temp[4] instead of temp[2], using total instead of photoelectric
        self.la2=la2                                            # absorption length in cm at incident x-ray energy
        angle_corrected2 = ( 2.0-(1.0-delta1)/(1.0-delta2)*(2.0-angle_corrected1**2) )**0.5
        # using Snell's law, assume beta effect not ignificant, in radian
        for (ii, depth) in enumerate(self.depths):
            if self.depths[ii]<self.thickness1:                 # top layer
                beampath = depth/math.sin(angle_corrected1)     # in cm
                inten0 = math.exp(-beampath/la1)                # attenuated incident beam intensity at depths
            else:                                               # bottom layer
                beampath1 = self.thickness1/math.sin(angle_corrected1)
                beampath2 = (depth-self.thickness1)/math.sin(angle_corrected2)
                inten0 = math.exp(-beampath1/la1)*math.exp(-beampath2/la2)
            self.inten0[ii] = inten0                            # incident x-ray attenuation
    def getLa(self, energy, NumLayer=1.0):  # emitted fluorescence trasmission attenuation up to top surface
        transmitted=1.0
        # MN replace:
        temp=f1f2.get_delta(self.composition1, self.density1, energy)  # energy is for fluorescence
        self.delta1 = temp[0]
        self.beta1 = temp[1]
        self.la1 = temp[4]      # in cm, temp[4] instead of temp[2], using total instead of photoelectric
        # absorption length in cm at emitted fluorescence energy
        # temp[3]: atomic number density of the first element in atoms/cc
        # MN replace:
        temp=f1f2.get_delta(self.composition2, self.density2, energy)  # energy is for fluorescence
        self.delta2 = temp[0]
        self.beta2 = temp[1]
        self.la2 = temp[4]                                      # absorption length in cm at emitted fluorescence energy
        # in cm, temp[4] instead of temp[2], using total instead of photoelectric
        angle_exit = (math.pi/2.0 - self.angle)                 # becomes 90 at a small incident angle
        for (ii, depth) in enumerate(self.depths):
            if self.depths[ii]<self.thickness1:                 # top layer
                 transmitted=math.exp(-depth/math.sin(angle_exit)*NumLayer/self.la1)
            else:                                               # bottom layer
                transmitted2 = math.exp(-(depth-self.thickness1)/math.sin(angle_exit)*NumLayer/self.la2)
                transmitted1 = math.exp(-self.thickness1/math.sin(angle_exit)*NumLayer/self.la1)
                transmitted = transmitted2*transmitted1
        self.trans[ii] = transmitted
        self.absrp[ii] = 1.0 - transmitted
        scale = 0.0; scale1=0.0; scale2=0.0
        for (ii,trans) in enumerate(self.trans):
            scale = scale + trans*self.inten0[ii]*self.factors[ii]
            if self.depths[ii]<self.thickness1:
                scale1 = scale1 + trans*self.inten0[ii]*1.0
            else:   # if thickness2 is different from thickness1, weight differently
                scale2 = scale2 + trans*self.inten0[ii]*(self.thickness2/self.thickness1)
        # scale, scale1, scale2: emitted fluorescence transmission and depth profile
        self.scale = scale      # sum of (trans*factors) for nonsubstrate FY
        self.scale1 = scale1    # sum of (trans*factors) for substrate FY. factors=1, top layer
        self.scale2 = scale2    # sum of (trans*factors) for substrate FY. factors=pre, bttom layer

'''
----------------------------------------------------------------------------------------------------------
Detector efficiency, fluorescence attenuation
----------------------------------------------------------------------------------------------------------
'''
# WD30 used for XSW, sample-->He-->Kapton-->Collimator-->Detector
# WD60 used for XRM, sample-->Collimator-->Detector
# eV1 is fluorescence energy in eV
# xHe=1 means Helium gas from sample to WD60 collimator.
# xAl=1 means 1 layer of 1.5mil Al foil as attenuator
# xKapton=1 means 1 layer of 0.3 mil Kapton (
# WD=6 means working distance of 6cm for the detector.

def Assemble_QuadVortex(eV1):
    # quad vortex detector efficiency. eV1: fluo energy
    net=1.
    BeVortex=Material('Be', 1.85, 0.00125)
    SiO2Vortex=Material('SiO2', 2.2, 0.00001)
    SiVortex=Material('Si', 2.33, 0.035)
    BeVortex.getLa(eV1, 1)  # one Be layer in Vortex
    SiO2Vortex.getLa(eV1, 1)    # oxide layer on Si detection layer
    SiVortex.getLa(eV1, 1)  # Si detection layer, what's absorbed is counted.
    net=net*BeVortex.trans*SiO2Vortex.trans*SiVortex.absrp
    if (print2screen):
        print( '%.3f eV : BeVortex.trans=%.3e , SiO2Vortex.trans=%.3e, SiVortex.absrp=%.3e, Det_efficiency=%.3e' % (eV1, BeVortex.trans,  SiO2Vortex.trans, SiVortex.absrp, net))
    return net

def Assemble_Collimator(eV1, xHe=1, xAl=0,xKapton=0, WD=6.0, xsw=0):
    # from sample surface to detector
    # xsw=0/1/-1,  6cm, 3cm, no collimator
    # He_path depends on the collliator.
    if xHe==1:
        He_path=WD
    else:
        He_path=0.
    air_path = WD-He_path
    kapton_inside=0  # kapton inside collimator.  no collimator, no kapton_inside
    if xsw==0:  # WD60mm collimator
        kapton_inside=1                 # collimator has thin Kapton on the second aperture
        if xHe==1:
            air_path = 1.51             # 1.51cm between second aperture and Be of Vortex
        else:
            air_path=WD
        He_path = WD-air_path           # He is between sample surface to the second aperture of xrm collimator
    if xsw==1:  # WD30mm colllimator
        kapton_inside=1                 # collimator has thin Kapton on the second aperture
        if xHe==1:      # modified 11/18/2010
            He_path = 1.088   # from sample surface to Kapton cover/collimator
            air_path = WD - He_path            # 1.912cm between second aperture and Be of Vortex
            xKapton = xKapton+1         # one Kapton film used to fill He from sample surface to collimator
        else:
            air_path = WD0;     He_path=0
    air=Material('N1.56O0.48C0.03Ar0.01Kr0.000001Xe0.0000009', 0.0013, air_path)
    kapton=Material('C22H10O4N2', 1.42, 0.000762)       # 0.3mil thick
    HeGas=Material('He', 0.00009, He_path)
    AlFoil=Material('Al', 2.72, 0.00381)                # 1.5mil thick
    kaptonCollimator=Material('C22H10O4N2', 1.42, 0.000762) # 0.3mil thick
    #
    air.getLa(eV1, 1)   # 1 means number of layers here.
    kapton.getLa(eV1, xKapton)  #number of Kapton layers before collimator
    HeGas.getLa(eV1, xHe) # number of He gas layers, default=0
    AlFoil.getLa(eV1, xAl)  # number of Al foil,  default=0
    kaptonCollimator.getLa(eV1, kapton_inside)   # without collimator, no addition kapton inside
    #
    net=air.trans*HeGas.trans
    net=net*kapton.trans*AlFoil.trans*kaptonCollimator.trans
    if print2screen:
            print('%.3f eV: air.trans=%.3e, HeGas.trans=%.3e, kapton.trans=%.3e, AlFoil.trans=%.3e, kaptonCollimator.trans=%.3e, net=%.3e' %  \
            (eV1, air.trans, HeGas.trans, kapton.trans, AlFoil.trans,kaptonCollimator.trans, net))
    #print('%.3f eV: air.la=%.3e, HeGas.la=%.3e, kapton.la=%.3e' % (eV1, air.la, HeGas.la, kapton.la))
    return net

def Assemble_Detector(eV1, xHe=1, xAl=0,xKapton=0, WD=6.0, xsw=0):
    det_efficiency=Assemble_QuadVortex(eV1)
    trans2det=Assemble_Collimator(eV1, xHe, xAl,xKapton, WD, xsw)
    net=det_efficiency*trans2det
    return net

'''
----------------------------------------------------------------------------------------------------------
Combine detector, transmission, sample self-absorption, and fluorescing element distribution/concentration.
----------------------------------------------------------------------------------------------------------
'''
# this function is for XSW/XRM.
def cal_NetYield2(eV0, Atoms, xHe=0, xAl=0, xKapton=0, WD=6.0, xsw=0, WriteFile='Y' , xsect_resonant=0.0, sample=''):
    #   incident energy, list of elements, experimental conditions
    #   this one tries Ka, Kb, Lg, Lb, La, Lb
    angle0=45.; textOut=''
    if xsw!=0 and WriteFile=='Y':  print( 'XSW measurements')
    if sample=='':
        Include_SelfAbsorption='No'
    else:
        Include_SelfAbsorption='Yes'
        angle0=sample.angle/math.pi*180.
        textOut=sample.txt          # substrate concentration
    NetYield=[];    NetTrans=[]; Net=[];    out2='';	net=0.0
    text0=''
    edges      =['K',  'K',  'L1', 'L1', 'L2', 'L2', 'L2', 'L3', 'L3', 'L3']
    Fluo_lines =['Ka', 'Kb', 'Lb', 'Lg', 'Ln', 'Lb', 'Lg', 'Ll', 'La', 'Lb']
    outputfile='Elemental_Sensitivity.txt'
    if WriteFile=='Y':
        fo=open(outputfile, 'w')
        desc=' '
        if xHe==0:   desc=' not '
        if xsw==-1:
            out1='# 13IDC XRM/XSW using QuadVortex, incident x-ray energy at '+str(eV0)+' eV at '+str(angle0)+' degrees \n'
            out1+='# Helium path'+desc+'used, '+str(xAl)+' Al attenuators, '+str(xKapton+1)+' Kapton attenuators, '\
              +str(WD)+' cm working distance. \n'
        else:
            out1='# 13IDC XRM/XSW using QuadVortex + collimator,  incident x-ray energy at '+str(eV0)+' eV at '+str(angle0)+' degrees \n'
            out1+='# Helium path'+desc+'used, '+str(xAl)+' Al attenuators, '+str(xKapton)+' Kapton attenuators, '\
              +str(WD)+' cm working distance. \n'
        print( out1)
        fo.write(out1)
        if sample!='':
            for stuff in Atoms:
                text1='%6.3e %s/cm^3' % (stuff.Conc*sample.Nat1, stuff.AtSym)
                text0=text0+' '+text1
        textOut=textOut+' '+text0       # substrate concentration + other concentrations
        out1='# '+text0+'\n'
        if print2screen:
            print( out1)
        fo.write(out1)
        out1='%s\t%s\t%s   \t%s   \t%s   \t%s   \t%s\n' % ('atom', 'emit', 'emit_energy', 'yield', 'transmission', 'net_sensitivity', 'sensitivity*concentration')
        if print2screen:
            print( out1)
        fo.write(out1)
    for (ii, atom) in enumerate(Atoms):
        # MN replace:
        atnum=f1f2.AtSym2AtNum(atom.AtSym)
        # MN replace:
        temp=f1f2.AtNum2f1f2Xsect(atnum, eV0)
        xsect=temp[2]                           # photo-electric cross-section for fluorescence yield, temp[4] is total for attenuation
        con=atom.Conc
        for (nn, edge) in enumerate(edges):
            emit=Fluo_lines[nn]
            fy, emit_eV, emit_prob = fluo_yield(atom.AtSym, edge, emit, eV0)
            print(emit)
            if fy==0.0 or emit_prob==0:
                continue                        # try next item if FY=0
            else:
                if xsect_resonant!=0.0:         # use input value near edge
                    xsect=xsect_resonant	# for cross-section near absoprtion edge
                #print(xsect,fy,emit_prob)
                net_yield=xsect*fy*emit_prob    # net_yield --> cross-section, yield, emission_probability
                # net transmission --> transmission from surface through detector
## ------------------   [self-absorption]     ------------------------------
                trans_SelfAbsorp = 1.0           # for self-absorption.
                if Include_SelfAbsorption=='Yes':
                    sample.getPenetration(eV0)   # incident x-ray attenuation
                    if fy*emit_eV*emit_prob==0:  # any one of three is zero
                        break                    # skip
                    sample.getLa(emit_eV)
                    # account for incident x-ray attenuation, emitted x-ray attenuation
                    trans_SelfAbsorp = sample.scale         # for elements that are not part of substrate
                    if (atom in sample.ElemListFY):
                        if atom.tag=='substrate1':
                            trans_SelfAbsorp = sample.scale1  # for elements that are part of top substrate
                        if atom.tag=='substrate2':
                            trans_SelfAbsorp = sample.scale2  # for elements that are part of bottom substrate
## ----------------------------------------------------------------------------
            net_trans = Assemble_Detector(emit_eV, xHe, xAl,xKapton, WD, xsw)
            net = net_yield*net_trans*trans_SelfAbsorp  # elemental sensitivity
            inten = net*con  # sensitivity * concentration
            if WriteFile=='Y':
                out1='%s\t%s\t%6.1f   \t%.3e   \t%.3e   \t%.3e   \t%.3e\n' % (atom.AtSym+'_'+edge, emit, emit_eV, net_yield, net_trans, net, inten)
                fo.write(out1)
            if print2screen:
                print('%s %s %6.1f net_yield=%.3e net_trans=%.3e net=%.3e\t' % (atom.AtSym+'_'+edge, emit, emit_eV, net_yield, net_trans, net))
            #print(out1+'  %s, depth-dependent factor= %6.4f' % (atom.tag, trans_SelfAbsorp))
            if emit=='Kb' and fy!=0:        # if above K edge, don't bother trying L edges
                break
    return textOut

#def sim_spectra(eV0, Atoms, Conc, xHe=0, xAl=0, xKapton=0, WD=6.0, xsw=0, sample=''):
def sim_spectra(eV0, Atoms, xHe=0, xAl=0, xKapton=0, WD=6.0, xsw=0, sample=''):
    # sample=sample matrix with object attribues to add self-absorption effect
    # Atoms is a list with elements that have attributes AtSym, Conc, tag
    if xsw==-1:     xKapton=xKapton-1   # no collimator
    Include_SelfAbsorption='Yes'
    Print2Screen='No'
    if sample=='': Include_SelfAbsorption='No'
    if xsw!=0:  print('XSW measurements')
    xx=[]; yy=[]; tag=[];   intensity_max=-10.0; LoLimit=1e-10
    angle0=''; text1=''
    if sample!='':      # sample matrix option is used
        angle0=str(sample.angle*180./math.pi)
        angle0=angle0[:5]
        for (ii, item) in enumerate(sample.ElemListFY):  # add matrix elements to the lists
            Atoms.append(item)
    outputfile='simSpectrum_table.txt'
    fo=open(outputfile, 'w')
    out1='#incident x-ray at '+str(eV0)+' eV and '+angle0+' Deg.\n'
    if Print2Screen=='Yes': print(out1)
    fo.write(out1)
    out1='#Emission\tenergy(eV)\tintensity \n'
    if Print2Screen=='Yes': print(out1)
    fo.write(out1)
    out2='#'
    for (ix,atom) in enumerate(Atoms):
        # MN replace:
        atnum=f1f2.AtSym2AtNum(atom.AtSym)
        # con=Conc[ix]
        con=atom.Conc
        out2=out2+atom.AtSym+'['+str(con)+']   '
        for edge in ['K', 'L1', 'L2', 'L3']:
            # MN replace:
            temp=f1f2.AtNum2f1f2Xsect(atnum, eV0)
            xsect=temp[2]           # photoelectric crosssection for each element at incident x-ray, fluorescence yield
            # MN replace:
            temp=elam.use_ElamFY(atom.AtSym, edge)
            edge_eV=float(temp[2])   # absorption edge
            if eV0>edge_eV:
                fy=float(temp[3])
                for EmitLine in temp[4]:
                    EmitName=EmitLine[1]
                    emit_eV=float(EmitLine[2])
                    emit_prob=float(EmitLine[3])
                    if emit_eV<fluo_emit_min:  continue        # ignore fluorescence below fluo_emit_min (global variable)
                    name = atom.AtSym + '_'+ EmitName
                    # net transmission --> transmission from surface through detector
## ------------------   [self-absorption]     ------------------------------
                    trans_SelfAbsorp = 1.0              # for self-absorption.
                    if Include_SelfAbsorption=='Yes':
                        sample.getPenetration(eV0)      # incident x-ray attenuation
                        eV0str=str(eV0)
                        text1=' absorption_length1(%seV)= %2.2e%s \
                                absorption_length2(%seV)= %2.2e%s' % (eV0str, sample.la1*1.e4, 'microns', eV0str, sample.la2*1.e4, 'microns')
                        # text1 is added to sample.txt later
                        sample.getLa(emit_eV)           # emitted fluorescence attenuation
                        trans_SelfAbsorp = sample.scale # for elements that are not part of substrate
                        if (atom in sample.ElemListFY):
                            if atom.tag=='substrate1':
                                trans_SelfAbsorp = sample.scale1  # for elements that are part of top substrate
                            if atom.tag=='substrate2':
                                trans_SelfAbsorp = sample.scale2  # for elements that are part of bottom substrate
## ----------------------------------------------------------------------------
                    trans = Assemble_Detector(emit_eV, xHe, xAl,xKapton, WD, xsw)
                    intensity = con * fy * emit_prob * xsect * trans * trans_SelfAbsorp
                    if intensity<LoLimit: continue            # skip weak emission, arbtraray limit =1e-10.
                    if intensity>intensity_max: intensity_max=intensity
                    xx.append(emit_eV);    yy.append(intensity);   tag.append(name)
    for ix in range(len(yy)):
        yy[ix]=yy[ix]/intensity_max*100.00          # makes the strongest line to 100.0
        out1='%s\t%f\t%f \n' % (tag[ix], xx[ix], yy[ix])
        if Print2Screen=='Yes': print(out1)
        fo.write(out1)
    if Print2Screen=='Yes': print(out2)
    fo.write(out2)
    fo.close()
    out1=sim_GaussPeaks(xx, yy, det_res, eV0)        # det_res: detector resoultion for Gaussian width (global variable)
    if Include_SelfAbsorption=='Yes':
        sample.txt=sample.txt+text1                 # sample.txt is combined to output of cal_NetYield
    text=cal_NetYield2(eV0, Atoms, xHe, xAl, xKapton, WD, xsw, sample=sample)  # calculate net yield with weight-averaged emission
    print(out2)
    return text  # cal_NetYield2 output is str with number densities of elements

def sim_GaussPeaks(xx, yy, width, eV0):  #xx, yy: lists, width: a peak width, eV0: incident energy
    xline=[]; yline=[]; dX=10.0;    minX=fluo_emit_min #10eV steps, lowest-->fluo_emit_min
    amp=100                             # arbitrary multiplier to shift up spectrum
    NumOfSteps = int((eV0-minX)/dX)
    NumOfPeaks = len(xx)
    for ix in range(NumOfSteps+1):
        xline.append(0);  yline.append(0)
    for (iy,peak) in enumerate(xx):
        X0=float(peak)
        Y0=float(yy[iy])
        for ix in range(NumOfSteps+1):
            energy = minX + ix*dX
            inten = Y0*math.exp(-((energy-X0)/width)**2)*amp
            xline[ix]=energy
            yline[ix]+=inten
    outputfile='simSpectrum_plot.txt'
    fo=open(outputfile, 'w')
#    DetectorLimit=1e5                   # upperlimit for total counts to 1e5 (1e5 CPS)
    LoLimit=0.001                         # low limit for each channel
#    total=0.0
    factor=1.0
    for ix in range(NumOfSteps+1):
        if yline[ix]<LoLimit: yline[ix]=LoLimit
#        total+=yline[ix]  # add counts
#    factor=total/DetectorLimit
    for ix in range(NumOfSteps+1):
        yline[ix]=yline[ix]/factor
        out1=str(xline[ix])+'\t'+str(yline[ix])+'\n'
        fo.write(out1)
    fo.close()
    #return xline, yline

class input_param:
    def __init__(self, eV0=14000,
                 Atoms=[],
                 xHe=0.0,
                 xAl=0,
                 xKap=0,
                 WD=6.0,
                 xsw=0):
        if Atoms==[]:
            atom1=ElemFY()
            Atoms.append(atom1)
        self.eV0=eV0
        #incident x-ray energy in eV
        list1=[]; list2=[]
        for item in Atoms:
            list1.append(item.AtSym)
            list2.append(item.Conc)
        self.Atoms=list1
        #elements list
        self.Conc=list2
        #relative concentrations list
        self.xHe=xHe
        #He gas path before collimator?
        self.xAl=xAl
        #number of Al foils as attenuator
        self.xKap=xKap
        #number of Kapton foils as attenuator
        self.WD=WD
        #working distance in cm
        self.xsw=xsw
        #x-ray standing wave setup?  WD=3.0 for xsw=1

# ----------------------------------------------------------------

if __name__=='__main__':
    testing = 0
    if testing:
        #atom0='Fe'
        #emission0='Ka'
        #eV0=10000.
        #edge0='K'
        #eV1=6400.
        #mat0=Material('SiO2', 2.2, 0.001)
        #mat0.getLa(8000)
        #print(mat0.la, mat0.trans)
        #print(Assemble_QuadVortex(eV1))
        #print(Assemble_Collimator(eV1, xHe=1, xAl=0,xKapton=0, WD=6.0, xsw=0))
        # MN replace:
        matrix=SampleMatrix2('CaCO3', f1f2.nominal_density('CaCO3'), 0.001,'Fe2O3', f1f2.nominal_density('Fe2O3'), 0.001, 45.,  'all')
        eV0=7500.
        Atoms=[]
        atom=ElemFY('La', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Ce', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Nd', 10e-6);     Atoms.append(atom)
        for (ii, item) in enumerate(matrix.ElemListFY):
            pass
        sim_spectra(eV0, Atoms, sample=matrix)  # xKapton=-1 to remove WD60, -2 for WD30
    else:
        # March 2012 for Nov2011 beamtime
        eV0 = 10000.
        print2screen=0
        Atoms = []
        atom=ElemFY('Al', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Si', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Ca', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Cr', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Mn', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Fe', 10e-6);     Atoms.append(atom)
        atom=ElemFY('Ni', 10e-6);     Atoms.append(atom)
        #sim_spectra(eV0, Atoms, xHe=1, xAl=0, xKapton=0, WD=3.0, xsw=1, sample='')
        #Assemble_QuadVortex(1486.56)
        #Assemble_Collimator(1486.56, 1, 0, 0, 3., 1)
        #
        air_path = 1.088
        He_path = 1.912
        air=Material('N1.56O0.48C0.03Ar0.01Kr0.000001Xe0.0000009', 0.0013, air_path)
        kapton=Material('C22H10O4N2', 1.42, 0.000762)       # 0.3mil thick
        HeGas=Material('He', 0.00009, He_path)
        AlFoil=Material('Al', 2.72, 0.00381)                # 1.5mil thick
        #
        from math import *
        emitE = [
            1486.4, 1557.0, 1739.6, 1837.0, 3691.1, 4013.1, 5411.6,
            5947.0, 5896.5, 6492.0, 6400.8, 7059.6, 7474.4, 8266.6
            ]
        for eV1 in emitE:
            air.getLa(eV1, 1)   # 1 means number of layers here.
            print( eV1, air.la, exp(-1./air.la))

'''
def AtSym2FY(AtSym, Shell):  #returns fluorescence yield using atomic symbol and edge(K,L1,L2,L3,M)
    out=elam.use_ElamFY(AtSym, Shell)
    #AtSym, edge, edge-energy, FY, [ [transition, emission, energy, probability],[]...]
    FY=out[3]
    return FY

class Element:
    def __init__(self, AtNum, AtWt, f1, f2, Xsection):
        self.AtNum=AtNum
        self.AtWt=AtWt              #atomic weight g/mol
        self.f1=f1                  #real part of scattering factor
        self.f2=f2                  #imaginary part of scattering factor
        self.Xsection=Xsection      #atomic crossection Barns/atom
        # f1, f2, Xsection depends on incident x-ray energy

def setElement(AtSym, shell, eV0):
    AtNum=AtSym2AtNum(AtSym)
    AtWt=AtSym2AtWt(AtSym)
    f1f2=AtNum2f1f2Xsect(AtNum, eV0)
    AtSym=Element(AtNum, AtWt, f1f2[0], f1f2[1], f1f2[2])
    return AtSym

class input_param:
    def __init__(self, eV0=14000,
                 Atoms=[],
                 xHe=0.0,
                 xAl=0,
                 xKap=0,
                 WD=6.0,
                 xsw=0):
        if Atoms==[]:
            atom1=ElemFY()
            Atoms.append(atom1)
        self.eV0=eV0
        #incident x-ray energy in eV
        list1=[]; list2=[]
        for item in Atoms:
            list1.append(item.AtSym)
            list2.append(item.Conc)
        self.Atoms=list1
        #elements list
        self.Conc=list2
        #relative concentrations list
        self.xHe=xHe
        #He gas path before collimator?
        self.xAl=xAl
        #number of Al foils as attenuator
        self.xKap=xKap
        #number of Kapton foils as attenuator
        self.WD=WD
        #working distance in cm
        self.xsw=xsw
        #x-ray standing wave setup?  WD=3.0 for xsw=1
'''


# -*- coding: utf-8 -*-
#
# Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG),
# acting on behalf of its Max Planck Institute for Intelligent Systems and the
# Max Planck Institute for Biological Cybernetics. All rights reserved.
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights
# on this computer program. You can only use this computer program if you have closed a license agreement
# with MPG or you get the right to use the computer program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and liable to prosecution.
# Contact: ps-license@tuebingen.mpg.de
#
#
# If you use this code in a research publication please consider citing the following:
#
# Expressive Body Capture: 3D Hands, Face, and Body from a Single Image <https://arxiv.org/abs/1904.05866>
# AMASS: Archive of Motion Capture as Surface Shapes <https://arxiv.org/abs/1904.03278>
#
#
# Code Developed by:
# Nima Ghorbani <https://www.linkedin.com/in/nghorbani/>
# Vassilis Choutas <https://ps.is.tuebingen.mpg.de/employees/vchoutas> for ContinousRotReprDecoder

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from pytorch3d.transforms.rotation_conversions import matrix_to_axis_angle as matrot2aa

__all__ = [
    'VPoserV1',
    'VPoserV2',
    'create'
]

class BatchFlatten(nn.Module):
    def __init__(self):
        super(BatchFlatten, self).__init__()
        self._name = 'batch_flatten'

    def forward(self, x):
        return x.view(x.shape[0], -1)

class ContinousRotReprDecoder(nn.Module):
    def __init__(self):
        super(ContinousRotReprDecoder, self).__init__()

    def forward(self, module_input):
        reshaped_input = module_input.view(-1, 3, 2)

        b1 = F.normalize(reshaped_input[:, :, 0], dim=1)

        dot_prod = torch.sum(b1 * reshaped_input[:, :, 1], dim=1, keepdim=True)
        b2 = F.normalize(reshaped_input[:, :, 1] - dot_prod * b1, dim=-1)
        b3 = torch.cross(b1, b2, dim=1)

        return torch.stack([b1, b2, b3], dim=-1)

class NormalDistDecoder(nn.Module):
    def __init__(self, num_feat_in, latentD):
        super(NormalDistDecoder, self).__init__()

        self.mu = nn.Linear(num_feat_in, latentD)
        self.logvar = nn.Linear(num_feat_in, latentD)

    def forward(self, Xout):
        return torch.distributions.normal.Normal(self.mu(Xout), F.softplus(self.logvar(Xout)))

class VPoserV1(nn.Module):
    def __init__(self, num_neurons=512, latentD=32, data_shape=[1,21,3], use_cont_repr=True):
        super(VPoserV1, self).__init__()

        self.latentD = latentD
        self.use_cont_repr = use_cont_repr

        n_features = np.prod(data_shape)
        self.num_joints = data_shape[1]

        self.bodyprior_enc_bn1 = nn.BatchNorm1d(n_features)
        self.bodyprior_enc_fc1 = nn.Linear(n_features, num_neurons)
        self.bodyprior_enc_bn2 = nn.BatchNorm1d(num_neurons)
        self.bodyprior_enc_fc2 = nn.Linear(num_neurons, num_neurons)
        self.bodyprior_enc_mu = nn.Linear(num_neurons, latentD)
        self.bodyprior_enc_logvar = nn.Linear(num_neurons, latentD)
        self.dropout = nn.Dropout(p=.1, inplace=False)

        self.bodyprior_dec_fc1 = nn.Linear(latentD, num_neurons)
        self.bodyprior_dec_fc2 = nn.Linear(num_neurons, num_neurons)

        if self.use_cont_repr:
            self.rot_decoder = ContinousRotReprDecoder()

        self.bodyprior_dec_out = nn.Linear(num_neurons, self.num_joints* 6)

    def encode(self, Pin):
        '''

        :param Pin: Nx(numjoints*3)
        :return:
        '''
        Xout = Pin.view(Pin.size(0), -1)  # flatten input
        Xout = self.bodyprior_enc_bn1(Xout)

        Xout = F.leaky_relu(self.bodyprior_enc_fc1(Xout), negative_slope=.2)
        Xout = self.bodyprior_enc_bn2(Xout)
        Xout = self.dropout(Xout)
        Xout = F.leaky_relu(self.bodyprior_enc_fc2(Xout), negative_slope=.2)
        return torch.distributions.normal.Normal(self.bodyprior_enc_mu(Xout), F.softplus(self.bodyprior_enc_logvar(Xout)))

    def decode(self, Zin):
        bs = Zin.shape[0]

        Xout = F.leaky_relu(self.bodyprior_dec_fc1(Zin), negative_slope=.2)
        Xout = self.dropout(Xout)
        Xout = F.leaky_relu(self.bodyprior_dec_fc2(Xout), negative_slope=.2)
        Xout = self.bodyprior_dec_out(Xout)
        if self.use_cont_repr:
            Xout = self.rot_decoder(Xout)
        else:
            Xout = torch.tanh(Xout)

        return {
            'pose_body': matrot2aa(Xout.view(-1, 3, 3)).view(bs, -1, 3),
            'pose_body_matrot': Xout.view(bs, -1, 9)
        }

    def forward(self, Pin):
        '''

        :param Pin: aa: Nx1xnum_jointsx3
        :return:
        '''
        q_z = self.encode(Pin)
        q_z_sample = q_z.rsample()
        decode_results = self.decode(q_z_sample)
        decode_results.update({'poZ_body_mean': q_z.mean, 'poZ_body_std': q_z.scale, 'q_z': q_z})
        return decode_results

    def sample_poses(self, num_poses, seed=None):
        np.random.seed(seed)
        dtype = self.bodyprior_dec_fc1.weight.dtype
        device = self.bodyprior_dec_fc1.weight.device
        self.eval()
        with torch.no_grad():
            Zgen = torch.tensor(np.random.normal(0., 1., size=(num_poses, self.latentD)), dtype=dtype).to(device)
        return self.decode(Zgen)

class VPoserV2(nn.Module):
    def __init__(self, num_neurons=512,latentD=32):
        super(VPoserV2, self).__init__()

        num_neurons, self.latentD = num_neurons, latentD

        self.num_joints = 21
        n_features = self.num_joints * 3

        self.encoder_net = nn.Sequential(
            BatchFlatten(),
            nn.BatchNorm1d(n_features),
            nn.Linear(n_features, num_neurons),
            nn.LeakyReLU(),
            nn.BatchNorm1d(num_neurons),
            nn.Dropout(0.1),
            nn.Linear(num_neurons, num_neurons),
            nn.Linear(num_neurons, num_neurons),
            NormalDistDecoder(num_neurons, self.latentD)
        )

        self.decoder_net = nn.Sequential(
            nn.Linear(self.latentD, num_neurons),
            nn.LeakyReLU(),
            nn.Dropout(0.1),
            nn.Linear(num_neurons, num_neurons),
            nn.LeakyReLU(),
            nn.Linear(num_neurons, self.num_joints * 6),
            ContinousRotReprDecoder(),
        )

    def encode(self, pose_body):
        '''
        :param Pin: Nx(numjoints*3)
        :param rep_type: 'matrot'/'aa' for matrix rotations or axis-angle
        :return:
        '''
        return self.encoder_net(pose_body)

    def decode(self, Zin):
        bs = Zin.shape[0]

        prec = self.decoder_net(Zin)

        return {
            'pose_body': matrot2aa(prec.view(-1, 3, 3)).view(bs, -1, 3),
            'pose_body_matrot': prec.view(bs, -1, 9)
        }

    def forward(self, pose_body):
        '''
        :param Pin: aa: Nx1xnum_jointsx3 / matrot: Nx1xnum_jointsx9
        :param input_type: matrot / aa for matrix rotations or axis angles
        :param output_type: matrot / aa
        :return:
        '''

        q_z = self.encode(pose_body)
        q_z_sample = q_z.rsample()
        decode_results = self.decode(q_z_sample)
        decode_results.update({'poZ_body_mean': q_z.mean, 'poZ_body_std': q_z.scale, 'q_z': q_z})
        return decode_results

    def sample_poses(self, num_poses, seed=None):
        np.random.seed(seed)

        some_weight = [a for a in self.parameters()][0]
        dtype = some_weight.dtype
        device = some_weight.device
        self.eval()
        with torch.no_grad():
            Zgen = torch.tensor(np.random.normal(0., 1., size=(num_poses, self.latentD)), dtype=dtype, device=device)

        return self.decode(Zgen)

def create(f_state_dict,version):
    if version==1:
        vposer = VPoserV1()
    elif version==2:
        vposer = VPoserV2()
    state_dict = torch.load(f_state_dict,map_location='cpu')
    vposer.load_state_dict(state_dict)
    vposer.eval()
    return vposer