# encoding: utf-8
#
#Copyright (C) 2017-2024, P. R. Wiecha, A. Arbouet
#
#This program is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
"""
fast electron simulations

Contributed by A. Arbouet, CEMES-CNRS, 2020
email: arbouet@cemes.fr

"""
from __future__ import print_function
from __future__ import absolute_import

import warnings
# import math

import numpy as np
import copy
import numba

import time

# =============================================================================
# additional visualization
# =============================================================================
def visu_structure_electron(struct, x0=None, y0=None, 
                            scale_e_pos=10, marker_e_pos='x', color_e_pos='r',
                            show=True, **kwargs):
    """plot 2D structure projection together with electron beam position
    
    all further kwargs are passed to `pyGDM2.visu.structure`
    
    Parameters
    ----------
    struct : list or :class:`.core.simulation`
          either list of 3d coordinate tuples or simulation description object 
          
    x0, y0 : float, default: None
        position of the electron beam in the (OXY) plane. If `None` and 
        simulation instance is passed as structure, try to infer electron 
        beam position(s) from incident field.
    
    scale_e_pos : float, default: 15
        scale of electron-position marker
    
    marker_e_pos : str, default: 'x'
        maker symbol for electron-position
    
    color_e_pos : str, default: 'r'
        color of electron-position marker
    
    Returns
    -------
    result returned by matplotlib's `scatter`
    
    """
    import matplotlib.pyplot as plt
    from pyGDM2 import visu
    from pyGDM2 import tools
    from pyGDM2 import core_py
    from pyGDM2 import fields_py
    step = tools.get_step_from_geometry(struct)
    
    ## if not given, try to infer electron positions defined in simulation
    if x0 is None and y0 is None and type(struct) == core_py.simulation:
        if struct.efield.field_generator.__name__ == fields_py.fast_electron.__name__:
            elpos = []
            for kw in struct.efield.kwargs_permutations:
                elpos.append([kw['x0'], kw['y0']])
            x0, y0 = np.unique(elpos, axis=0).T
    
    im = visu.structure(struct, show=False, **kwargs)
    plt.scatter(x0, y0, color=color_e_pos, marker=marker_e_pos, s=scale_e_pos*np.int32(step), lw=3)    
    plt.gca().set_aspect('equal')
    
    if show:
        plt.show()
    
    return im




# =============================================================================
# helper
# =============================================================================
def electron_speed(electron_kinetic_energy):
    """ Velocity of electron with kinetic energy electron_kinetic_energy
       
    Parameters
    ----------   
    electron_kinetic_energy : float
        electron kinetic energy (keV)
    
    Returns
    -------
    velec : speed of electron in units of the speed of light
    """        
    Eo_el = 511.0   # electron rest mass (keV)
    gamma = 1.0 + electron_kinetic_energy/Eo_el 
    beta = np.sqrt(1.0 - 1.0/gamma**2)
    return beta





# =============================================================================
# evaluation functions
# =============================================================================
def CL(sim, field_index, r=10000., tetamin=0, tetamax=np.pi, Nteta=30, 
       phimin=0, phimax=2*np.pi, Nphi=36, polarizerangle='none', 
       return_value='int_Es'):
    """Cathodoluminescence emission generated by a fast electron moving along (OZ)
       
    The fundamental E-field created by the fast electron is given by the 
    function :func:`pyGDM2.fields.fast_electron`.
    Electron kinetic energy, beam position and propagation direction are parameters 
    of fields.fast_electron.
    
    Using :func:`pyGDM2.linear_py.farfield` for calculation of farfield scattering.
    
    For further documentation on the possible kwargs, their meaning and 
    returned values, see doc of :func:`pyGDM2.linear_py.farfield`
    
    Parameters
    ----------
    sim : :class:`.core.simulation`
        simulation description
        
    field_index : int
        index of evaluated self-consistent field to use for calculation. Can be
        obtained for specific parameter-set using :func:`.tools.get_closest_field_index`
           
    Returns
    -------    
    CL_signal : Cathodoluminescence emission intensity 
        number of emitted photons per eV per electron
        
    Notes
    -----
    For details of the EELS computation:
    Arnaud Arbouet et al. *Electron energy losses and cathodoluminescence 
    from complex plasmonic nanostructures: spectra, maps and radiation patterns 
    from a generalized field propagator*, New J. Phys. 16, 113012 (2014)
    """
    from pyGDM2 import linear_py as linear
        
    if sim.E is None: 
        raise ValueError("Error: Scattering field inside the structure not yet evaluated. Run `core.scatter` simulation first.")
    
    field_params = sim.E[field_index][0]
    wavelength   = field_params['wavelength'] 
    MULT = 2.31E20
    factor = MULT * wavelength
    
    ff_scat = linear.farfield(sim, field_index, r_probe=None, r=r, 
                              tetamin=tetamin, tetamax=tetamax, Nteta=Nteta, 
                              phimin=phimin, phimax=phimax, Nphi=Nphi, 
                              polarizerangle=polarizerangle, return_value=return_value)
    
    if return_value.lower() in ["map", "efield"]:        
        CL_signal = [ff_scat[0], ff_scat[1], 
                     ff_scat[2]*factor, ff_scat[3]*factor, ff_scat[4]*factor]
    else:
        CL_signal = factor * ff_scat
 
    return CL_signal



def EELS(sim, field_index):       
    """Electron Energy Loss probability of a fast electron moving along (OZ)
       
    The E-field created by the fast electron is given by the function fields.fast_electron.
    Electron kinetic energy, beam position and propagation direction are parameters 
    of fields.fast_electron.       
    
    Parameters
    ----------
    sim : :class:`.core.simulation`
        simulation description
        
    field_index : int
        index of evaluated self-consistent field to use for calculation. Can be
        obtained for specific parameter-set using :func:`.tools.get_closest_field_index`
           
    Returns
    -------    
    eels_signal : electron energy loss probability per eV per electron
        
    Notes
    -----
    For details of the EELS computation:
    Arnaud Arbouet et al 2014 New J. Phys. 16 113012     
    """
    from pyGDM2 import tools
    
    if sim.E is None: 
        raise ValueError("Error: Scattering field inside the structure not yet evaluated. " +
                         "Run `core.scatter` simulation first.")
    
    ## --- constants
    MULT =  929448800             # Multiplication constant 
    
    step = tools.get_step_from_geometry(sim)
    field_params = sim.E[field_index][0]
    wavelength   = field_params['wavelength']    
    electron_kinetic_energy  = field_params['electron_kinetic_energy'] 
    x0    = field_params['x0']   
    y0    = field_params['y0']
    
    env_dict = sim.dyads.getConfigDictG(wavelength, sim.struct, sim.efield)
    eps2 = env_dict['eps2']
    n2 = eps2**0.5
    
    k0 =  n2*(2.*np.pi/wavelength)  
    velec = electron_speed(electron_kinetic_energy)
    krho = k0*np.sqrt(1-velec**2.) / velec
    
# =============================================================================
#     Computation of E-field along electron trajectory
# =============================================================================
    ## compute the dipole moments at each discretization cell
    Eint = sim.E[field_index][1]
    alpha_tensor = sim.dyads.getPolarizabilityTensor(wavelength, sim.struct)
    P = np.matmul(alpha_tensor, Eint[...,None])[...,0]
       
    ## compute the loss probability
    from scipy import special
    res = 1j*0.
    for idip, _pos in enumerate(sim.struct.geometry):   
        RT = np.sqrt( (_pos[0] - x0)**2. + (_pos[1] - y0)**2. )
        if (RT > 1.5*step):
            U = krho * RT
            contrib1 = - krho**2.*special.kv(0, U) * P[idip][2]
            contrib2 = 1j*k0*krho*((_pos[0]-x0) * P[idip][0] + 
                                   (_pos[1]-y0) * P[idip][1]) * \
                                              special.kv(1, U) / (velec*RT)
            contrib = (contrib1 + contrib2) * np.exp(-1j*k0*_pos[2]/velec)
        else: 
            contrib = 0
        res = res + contrib
    
    ## electron energy loss
    eels_ret = np.real(MULT*res/k0)
        
    return eels_ret




