# -*- coding: utf-8 -*-
r"""
.. _synchrotron-sources:

Synchrotron sources
-------------------

The images below are produced by
``\tests\raycing\test_sources.py`` and by
``\examples\withRaycing\01_SynchrotronSources\synchrotronSources.py``.

Bending magnet
~~~~~~~~~~~~~~

On a transversal screen the image is unlimited horizontally (practically
limited by the front end). The energy distribution is red-shifted for off-plane
photons. The polarization is primarily horizontal. The off-plane radiation has
non-zero projection to the vertical polarization plane.

+----------+---------------------+---------------------+---------------------+
| source   |     total flux      |   horiz. pol. flux  |   vert. pol. flux   |
+==========+=====================+=====================+=====================+
| using WS |    |bmTotalWS|      |    |bmHorizWS|      |     |bmVertWS|      |
+----------+---------------------+---------------------+---------------------+
| internal |                     |                     |                     |
| xrt      |    |bmTotalXRT|     |    |bmHorizXRT|     |     |bmVertXRT|     |
+----------+---------------------+---------------------+---------------------+

.. |bmTotalWS| imagezoom:: _images/3bm_ws1-n-wideE-1TotalFlux.png
.. |bmHorizWS| imagezoom:: _images/3bm_ws1-n-wideE-2horizFlux.png
.. |bmVertWS| imagezoom:: _images/3bm_ws1-n-wideE-3vertFlux.png
   :loc: upper-right-corner
.. |bmTotalXRT| imagezoom:: _images/3bm_xrt1-n-wideE-1TotalFlux.png
.. |bmHorizXRT| imagezoom:: _images/3bm_xrt1-n-wideE-2horizFlux.png
.. |bmVertXRT| imagezoom:: _images/3bm_xrt1-n-wideE-3vertFlux.png
   :loc: upper-right-corner

The off-plane radiation is in fact left and right polarized:

+----------+-----------------------------+
| source   | circular polarization rate  |
+==========+=============================+
| using    |                             |
| WS       |     |bmCircPolRateWS|       |
+----------+-----------------------------+
| internal |                             |
| xrt      |     |bmCircPolRateXRT|      |
+----------+-----------------------------+

.. |bmCircPolRateWS| imagezoom:: _images/3bm_ws1-n-wideE-5CircPolRate.png
.. |bmCircPolRateXRT| imagezoom:: _images/3bm_xrt1-n-wideE-5CircPolRate.png

The horizontal phase space projected to a transversal plane at the origin is
parabolic:

+-------------------------+---------------------+
| zero electron beam size | σ\ :sub:`x` = 49 µm |
+=========================+=====================+
|         |bmPhaseSp0|    |     |bmPhaseSpN|    |
+-------------------------+---------------------+

.. |bmPhaseSp0| imagezoom:: _images/3bm_xrt1-0-wideE-horPhaseSpace.png
.. |bmPhaseSpN| imagezoom:: _images/3bm_xrt1-n-wideE-horPhaseSpace.png

Multipole wiggler
~~~~~~~~~~~~~~~~~

The horizontal image size is determined by the parameter K. The energy
distribution is red-shifted for off-plane photons. The polarization is
primarily horizontal. The off-plane radiation has non-zero projection to the
vertical polarization plane.

+----------+------------------+------------------+------------------+
| source   | total flux       | horiz. pol. flux | vert. pol. flux  |
+==========+==================+==================+==================+
| using    |                  |                  |                  |
| WS       |    |wTotalWS|    |    |wHorizWS|    |    |wVertWS|     |
+----------+------------------+------------------+------------------+
| internal |                  |                  |                  |
| xrt      |    |wTotalXRT|   |    |wHorizXRT|   |    |wVertXRT|    |
+----------+------------------+------------------+------------------+

.. |wTotalWS| imagezoom:: _images/2w_ws1-n-wideE-1TotalFlux.png
.. |wHorizWS| imagezoom:: _images/2w_ws1-n-wideE-2horizFlux.png
.. |wVertWS| imagezoom:: _images/2w_ws1-n-wideE-3vertFlux.png
   :loc: upper-right-corner
.. |wTotalXRT| imagezoom:: _images/2w_xrt1-n-wideE-1TotalFlux.png
.. |wHorizXRT| imagezoom:: _images/2w_xrt1-n-wideE-2horizFlux.png
.. |wVertXRT| imagezoom:: _images/2w_xrt1-n-wideE-3vertFlux.png
   :loc: upper-right-corner

The horizontal longitudinal cross-section reveals a sinusoidal shape of the
source. The horizontal phase space projected to the transversal plane at the
origin has individual branches for each pole.

+-------------------------+-------------------------+
| zero electron beam size |   σ\ :sub:`x` = 49 µm   |
+=========================+=========================+
|          |wYX0|         |          |wYXN|         |
+-------------------------+-------------------------+
|       |wPhaseSp0|       |       |wPhaseSpN|       |
+-------------------------+-------------------------+

.. |wYX0| imagezoom:: _images/2w_xrt1-0-wideE-crossectionYX.png
.. |wYXN| imagezoom:: _images/2w_xrt1-n-wideE-crossectionYX.png
   :loc: upper-right-corner
.. |wPhaseSp0| imagezoom:: _images/2w_xrt1-0-wideE-horPhaseSpace.png
.. |wPhaseSpN| imagezoom:: _images/2w_xrt1-n-wideE-horPhaseSpace.png
   :loc: upper-right-corner

Undulator
~~~~~~~~~

The module :mod:`~tests.raycing.test_sources` has functions for
visualization of the angular and energy distributions of the implemented
sources in 2D and 3D. This is especially useful for undulators because they
have sharp peaks, which requires a proper selection of angular and energy
meshes.

|I0_x'E|  |I0_z'E|  |IpPol|

.. |I0_x'E| imagezoom:: _images/I0_x'E-1-und-xrt.png
.. |I0_z'E| imagezoom:: _images/I0_z'E-1-und-xrt.png
.. |IpPol| imagezoom:: _images/Itot

The ray traced images of an undulator source (produced by
``\examples\withRaycing\01_SynchrotronSources\synchrotronSources.py``)
are feature-rich. The polarization is primarily horizontal. The off-plane
radiation has non-zero projection to the vertical polarization plane.

+---------+----------------+----------------+----------------+----------------+
| source  |   total flux   | hor. pol. flux | ver. pol. flux |  deg. of pol.  |
+=========+================+================+================+================+
| using   |                |                |                |                |
| Urgent  | |uTotalUr|     |   |uHorizUr|   |    |uVertUr|   | |uDegPolUr|    |
+---------+----------------+----------------+----------------+----------------+
| internal|                |                |                |                |
| xrt     | |uTotalXRT|    |  |uHorizXRT|   |    |uVertXRT|  | |uDegPolXRT|   |
+---------+----------------+----------------+----------------+----------------+

.. |uTotalUr| imagezoom:: _images/1u_urgent3-n-monoE-1TotalFlux.png
.. |uHorizUr| imagezoom:: _images/1u_urgent3-n-monoE-2horizFlux.png
.. |uVertUr| imagezoom:: _images/1u_urgent3-n-monoE-3vertFlux.png
   :loc: upper-right-corner
.. |uDegPolUr| imagezoom:: _images/1u_urgent3-n-monoE-4DegPol.png
   :loc: upper-right-corner
.. |uTotalXRT| imagezoom:: _images/1u_xrt3-n-monoE-1TotalFlux-Espread.png
.. |uHorizXRT| imagezoom:: _images/1u_xrt3-n-monoE-2horizFlux-Espread.png
.. |uVertXRT| imagezoom:: _images/1u_xrt3-n-monoE-3vertFlux-Espread.png
   :loc: upper-right-corner
.. |uDegPolXRT| imagezoom:: _images/1u_xrt3-n-monoE-4DegPol-Espread.png
   :loc: upper-right-corner

Elliptical undulator
~~~~~~~~~~~~~~~~~~~~

An elliptical undulator gives circular images with a higher circular
polarization rate in the inner rings:

+----------+--------------------+--------------------+--------------------+
| source   |    total flux      |   hor. pol. flux   |   ver. pol. flux   |
+==========+====================+====================+====================+
| using    |                    |                    |                    |
| Urgent   |     |euTotalUr|    |     |euHorizUr|    |     |euVertUr|     |
+----------+--------------------+--------------------+--------------------+
| internal |                    |                    |                    |
| xrt      |    |euTotalXRT|    |    |euHorizXRT|    |     |euVertXRT|    |
+----------+--------------------+--------------------+--------------------+

.. |euTotalUr| imagezoom:: _images/4eu_urgent3-n-monoE-1TotalFlux.png
.. |euHorizUr| imagezoom:: _images/4eu_urgent3-n-monoE-2horizFlux.png
.. |euVertUr| imagezoom:: _images/4eu_urgent3-n-monoE-3vertFlux.png
   :loc: upper-right-corner
.. |euTotalXRT| imagezoom:: _images/4eu_xrt3-n-monoE-1TotalFlux.png
.. |euHorizXRT| imagezoom:: _images/4eu_xrt3-n-monoE-2horizFlux.png
.. |euVertXRT| imagezoom:: _images/4eu_xrt3-n-monoE-3vertFlux.png
   :loc: upper-right-corner

+----------+----------------------------+----------------------------+
| source   |     deg. of pol.           | circular polarization rate |
+==========+============================+============================+
| using    |                            |                            |
| Urgent   |       |euDegPolUr|         |     |euCircPolRateUr|      |
+----------+----------------------------+----------------------------+
| internal |                            |                            |
| xrt      |        |euDegPolXRT|       |    |euCircPolRateXRT|      |
+----------+----------------------------+----------------------------+

.. |euDegPolUr| imagezoom:: _images/4eu_urgent3-n-monoE-4DegPol.png
.. |euCircPolRateUr| imagezoom:: _images/4eu_urgent3-n-monoE-5CircPolRate.png
   :loc: upper-right-corner
.. |euDegPolXRT| imagezoom:: _images/4eu_xrt3-n-monoE-4DegPol.png
.. |euCircPolRateXRT| imagezoom:: _images/4eu_xrt3-n-monoE-5CircPolRate.png
   :loc: upper-right-corner

.. _undulator_custom:

Custom field undulator
~~~~~~~~~~~~~~~~~~~~~~

A custom magnetic field can be specified by an Excel file or a column text
file. The example below is based on a table supplied by Hamed Tarawneh
[Tarawneh]_. The idea of introducing quasi-periodicity is to shift the n-th
harmonics down in energy relative to the exact n-fold multiple of the 1st
harmonic energy. This trick eliminates higher *monochromator* harmonics that
are situated at the exact n-fold energies, which is a safer solution compared
to a gas absorption filter.

Compare the harmonic energies (half-maximum position at the higher energy side)
of the 3rd harmonic with the triple energy of the 1st harmonic.

.. [Tarawneh] Quasi-periodic undulator field for ARPES beamline at
   MAX IV 1.5 GeV ring, (2016) unpublished.

.. note::

    The definition of xyz coordinate system differs for the tabulated field and
    for xrt screens: z is along the beam direction in the tabulation and as a
    vertical axis in xrt.

+--------------------+--------------------------+--------------------------+
|                    |         periodic         |       quasi-periodic     |
+====================+==========================+==========================+
| tabulated field    |       |EPU_field|        |      |QEPU_field|        |
+--------------------+--------------------------+--------------------------+
| trajectory         |                          |                          |
| top view           |        |EPU_traj|        |       |QEPU_traj|        |
+--------------------+--------------------------+--------------------------+
| wide band          |                          |                          |
| image and spectrum |        |EPU_wide|        |       |QEPU_wide|        |
+--------------------+--------------------------+--------------------------+
| 1st harmonic       |                          |                          |
| image and spectrum |        |EPU_1sth|        |       |QEPU_1sth|        |
+--------------------+--------------------------+--------------------------+
| 3rd harmonic       |                          |                          |
| image and spectrum |        |EPU_3rdh|        |       |QEPU_3rdh|        |
+--------------------+--------------------------+--------------------------+

.. |EPU_field| imagezoom:: _images/1-EPU_HP_field.png
.. |EPU_traj| imagezoom:: _images/1-EPU_HP_mode-x_average.png
.. |EPU_wide| imagezoom:: _images/1EPU_HP_mode-1-band1totalFlux.png
.. |EPU_1sth| imagezoom:: _images/1EPU_HP_mode-2-1stHarmonic1totalFlux.png
.. |EPU_3rdh| imagezoom:: _images/1EPU_HP_mode-6-3rdHarmonic1totalFlux.png
.. |QEPU_field| imagezoom:: _images/2-QEPU_HP_field.png
   :loc: upper-right-corner
.. |QEPU_traj| imagezoom:: _images/2-QEPU_HP_mode-x_average.png
   :loc: upper-right-corner
.. |QEPU_wide| imagezoom:: _images/2QEPU_HP_mode-1-band1totalFlux.png
   :loc: upper-right-corner
.. |QEPU_1sth| imagezoom:: _images/2QEPU_HP_mode-2-1stHarmonic1totalFlux.png
   :loc: upper-right-corner
.. |QEPU_3rdh| imagezoom:: _images/2QEPU_HP_mode-6-3rdHarmonic1totalFlux.png
   :loc: upper-right-corner

For validation, our calculations are compared here with those by Spectra for a
particular case — quasi-periodic undulator defined by the same tabulated field,
the 3rd harmonic, at E=20.5 eV. Notice again that Spectra provides either a
spectrum *or* a transverse image while xrt can combine both by using colors and
brightness. Notice also that on the following pictures the p-polarized flux is
only ~3% of the total flux.

+------+------------+------------+
|      |  SPECTRA   |    xrt     |
+======+============+============+
| |tf| |  |CFspeT|  |  |CFxrtT|  |
+------+------------+------------+
| |pf| |  |CFspeP|  |  |CFxrtP|  |
+------+------------+------------+

.. |tf| replace:: total flux
.. |pf| replace:: p-pol flux
.. |CFspeT| imagezoom:: _images/spectra-custom.png
.. |CFspeP| imagezoom:: _images/spectra-custom_p.png
   :loc: lower-left-corner
.. |CFxrtT| imagezoom:: _images/2QEPU_mono3rdHarmonic.png
.. |CFxrtP| imagezoom:: _images/2QEPU_mono3rdHarmonic_p.png
   :loc: lower-left-corner

"""

"""
.. _THz:

Terahertz undulator
~~~~~~~~~~~~~~~~~~~

This is a computationally difficult example that is hardly possible to
calculate with other commonly used codes. To achieve good convergence in a
reasonable calculation time, we have introduced a new undulator parameter
*gIntervals* that controls the underlying Gauss-Legendre mesh. The user may
want to modify it in extreme cases (wigglers, near field, wide angles etc.).


+------------------+-----------------+
| horiz. pol. flux | vert. pol. flux |
+==================+=================+
|     |uTHzH|      |     |uTHzV|     |
+------------------+-----------------+

.. |uTHzH| image:: _images/2THzU_xrt1-n-narrowE-2horizFlux-U_code.png
   :scale: 50 %
.. |uTHzV| image:: _images/2THzU_xrt1-n-narrowE-3vertFlux-U_code.png
   :scale: 50 %
"""
__author__ = "Konstantin Klementiev", "Roman Chernikov"
__date__ = "08 Mar 2016"

import os, sys; 
#newCode = False
newCode = True

codePath = r"c:\github\xrt" if newCode else os.path.join('..', '..', '..')
sys.path.append(codePath)
import time
import numpy as np
import matplotlib
matplotlib.use("Agg")
from matplotlib import pyplot as plt
import xrt.backends.raycing as raycing
from xrt.backends.raycing.physconsts import K2B
raycing._VERBOSITY_ = 100
import xrt.backends.raycing.sources as rs
import xrt.backends.raycing.screens as rsc
import xrt.backends.raycing.run as rr
import xrt.plotter as xrtp
import xrt.runner as xrtr

showIn3D = False

isInternalSource = True  # xrt source or (Urgent or WS)
bins = 256  # Number of bins in the plot histogram
ppb = 1  # Number of pixels per histogram bin

def build_beamline(nrays=1e5,
                   sourceType='u',
                   energyRange='1harmonic',
                   what='rays',
                   targetOpenCL='auto',
                   E0=6900):

    beamLine = raycing.BeamLine(alignE=E0)

    # one of 'u', 'w', 'bm', 'eu', 'wu':
#    sourceType = 'u'
    # one of 'mono', '1harmonic', 'smaller', 'wide'
#    energyRange = '1harmonic'
#    energyRange = 'mono'
    #energyRange = 'wide'
#    print(what)
    beamLine.what = what

    #what = 'wave'  # only for rs.Undulator
    #suffix = '_fullLength_emittance_NF'
    beamLine.limitsFSM0X = 'symmetric'
    beamLine.limitsFSM0Z = 'symmetric'
    beamLine.suffix = '_new' if newCode else '_old'
#    E0 = 6900  # eV
    R0 = 25000.  # Distance to the screen [mm]

    uniformRayDensity=True if what=='wave' else False
    filamentBeam=True if what=='wave' else False
    eEpsilonX=0.0 if what=='wave' else 0.263
    eEpsilonZ=0.0 if what=='wave' else 0.008
    eEspread=0.0 if what=='wave' else 0.001
    eSigmaX=0.0 if what=='wave' else 48.65
    eSigmaZ=0.0 if what=='wave' else 6.197
    
    #for sourceType in ['u', 'w', 'bm', 'eu', 'wu']:
    
    if sourceType == 'u':
        whose = '_xrt'
        pprefix = '1'+sourceType+whose
        Source = rs.Undulator
        kwargs = dict(
            eE=3., eI=0.5,  # Parameters of the synchrotron ring [GeV], [Ampere]
            eEspread=eEspread,  # Energy spread of the electrons in the ring
            period=30., n=40,  # Parameters of the undulator, period in [mm]
            K=1.45,  # Deflection parameter (ignored if targetE is not None)
            eSigmaX=eSigmaX, eSigmaZ=eSigmaZ,  # Size of the electron beam [mkm]
            uniformRayDensity=uniformRayDensity,
            filamentBeam=filamentBeam,  # Single wavefront
            targetOpenCL=targetOpenCL,
            # R0=R0,   # Near Field.
            eEpsilonX=eEpsilonX, eEpsilonZ=eEpsilonZ)  # Emittance [nmrad]
        beamLine.xlimits = [-2, 2]  # Horizontal limits of the plot [mm]
        beamLine.zlimits = [-2, 2]  # Vertical limits of the plot [mm]
        beamLine.xlimitsZoom = [-2, 2]
        beamLine.zlimitsZoom = [-2, 2]
        beamLine.xPrimelimits = [-0.3, 0.3]  # Angular limits of the plot [mrad]
        kwargs['xPrimeMaxAutoReduce'] = False
        kwargs['zPrimeMaxAutoReduce'] = False
    elif sourceType == 'w':
        whose = '_xrt' if isInternalSource else '_ws'
        pprefix = '3'+sourceType+whose
        Source = rs.Wiggler if isInternalSource else rs.WigglerWS
        kwargs = dict(period=80., K=13., n=10, eE=3., eI=0.5)
        beamLine.xlimits = [-40, 40]
        beamLine.zlimits = [-20, 20]
        beamLine.xlimitsZoom = [-8, 8]
        beamLine.zlimitsZoom = [-4, 4]
        beamLine.xPrimelimits = [-2.3, 2.3]
        beamLine.limitsFSM0X = [-1000, 1000]
        beamLine.limitsFSM0Z = [-100, 100]
    elif sourceType == 'bm':
        whose = '_xrt' if isInternalSource else '_ws'
        pprefix = '4'+sourceType+whose
        Source = rs.BendingMagnet if isInternalSource else rs.BendingMagnetWS
        kwargs = dict(B0=1.7, eI=0.1, eE=3.)
    #    kwargs['uniformRayDensity'] = True
        beamLine.xlimits = [-40, 40]
        beamLine.zlimits = [-20, 20]
        beamLine.xlimitsZoom = [-8, 8]
        beamLine.zlimitsZoom = [-4, 4]
        beamLine.xPrimelimits = 'symmetric'
        beamLine.limitsFSM0X = [-500, 500]
        beamLine.limitsFSM0Z = [-20, 20]
    elif sourceType == 'eu':
        whose = '_xrt' if isInternalSource else '_urgent'
        pprefix = '5'+sourceType+whose
        Source = rs.Undulator if isInternalSource else rs.UndulatorUrgent
        kwargs = dict(
            period=30., Ky=1.45, Kx=1.45, n=40,
            eE=3., eI=0.5,
            eSigmaX=48.65, eSigmaZ=6.197,
            targetOpenCL=targetOpenCL,
            eEpsilonX=0.263, eEpsilonZ=0.008)
        if isInternalSource:
            kwargs['phaseDeg'] = 90
            kwargs['xPrimeMaxAutoReduce'] = False
            kwargs['zPrimeMaxAutoReduce'] = False
        beamLine.xlimits = [-12, 12]
        beamLine.zlimits = [-12, 12]
        beamLine.xlimitsZoom = [-5, 5]
        beamLine.zlimitsZoom = [-5, 5]
        beamLine.xPrimelimits = [-0.55, 0.55]
    elif sourceType == 'wu':  # wiggler by undulator code
        whose = '_xrt' if isInternalSource else '_ws'
        pprefix = '6'+sourceType+whose
        Source = rs.Undulator if isInternalSource else rs.WigglerWS
        kwargs = dict(period=80., K=13., n=10, eE=3., eI=0.5)
        kwargs['targetOpenCL'] = targetOpenCL
        beamLine.xlimits = [-40, 40]
        beamLine.zlimits = [-20, 20]
        beamLine.xlimitsZoom = [-5, 5]
        beamLine.zlimitsZoom = [-5, 5]
    #    kwargs['uniformRayDensity'] = True
    #    kwargs['gIntervals'] = 99
        kwargs['zPrimeMaxAutoReduce'] = False
        beamLine.xPrimelimits = [-2.3, 2.3]
        beamLine.limitsFSM0X = [-1000, 1000]
        beamLine.limitsFSM0Z = [-100, 100]
    elif sourceType == 'tu':  # tapered undulator
        whose = '_xrt'
        pprefix = '7'+sourceType+whose
        Source = rs.Undulator
        kwargs = dict(
            eE=3., eI=0.5,  # Parameters of the synchrotron ring [GeV], [Ampere]
            eEspread=eEspread,  # Energy spread of the electrons in the ring
            period=30., n=40,  # Parameters of the undulator, period in [mm]
            K=1.45,  # Deflection parameter (ignored if targetE is not None)
            eSigmaX=eSigmaX, eSigmaZ=eSigmaZ,  # Size of the electron beam [mkm]
            uniformRayDensity=uniformRayDensity,
            filamentBeam=filamentBeam,  # Single wavefront
            taper=[2, 10],
            targetOpenCL=targetOpenCL,
            # R0=R0,   # Near Field.
            eEpsilonX=eEpsilonX, eEpsilonZ=eEpsilonZ)  # Emittance [nmrad]
        beamLine.xlimits = [-2, 2]  # Horizontal limits of the plot [mm]
        beamLine.zlimits = [-2, 2]  # Vertical limits of the plot [mm]
        beamLine.xlimitsZoom = [-0.2, 0.2]
        beamLine.zlimitsZoom = [-0.2, 0.2]
        beamLine.xPrimelimits = [-0.3, 0.3]  # Angular limits of the plot [mrad]
        kwargs['xPrimeMaxAutoReduce'] = False
        kwargs['zPrimeMaxAutoReduce'] = False
    elif sourceType == 'custom':  # tapered undulator from custom field
        whose = '_xrt'
        Np = 40
        L0 = 30
        Kmax = 1.45
        alphaZ = 0.2
        zMin, zMax = -0.5*Np*L0, 0.5*Np*L0
        fieldZ = np.linspace(zMin, zMax, int(Np*L0)+1)
        
        #tapered
        fieldK = Kmax*(1.-alphaZ*(0.5-(fieldZ-zMin)/(Np*L0)))*np.sin(2*np.pi*fieldZ/L0)
        fieldB = K2B * fieldK / L0
        plt.figure("B")
        plt.plot(fieldZ, fieldB)
        plt.xlabel('z, mm', fontsize=12)
        plt.ylabel('B, T', fontsize=12)
        plt.savefig("B vs Z.png")
        #plt.show()
        fieldArray = np.vstack((fieldZ, fieldB)).T

        pprefix = '8'+sourceType+whose
        Source = rs.SourceFromField
        kwargs = dict(
            eE=3., eI=0.5,  # Parameters of the synchrotron ring [GeV], [Ampere]
            eEspread=eEspread,  # Energy spread of the electrons in the ring
#            period=30., n=40,  # Parameters of the undulator, period in [mm]
#            K=1.45,  # Deflection parameter (ignored if targetE is not None)
            eSigmaX=eSigmaX, eSigmaZ=eSigmaZ,  # Size of the electron beam [mkm]
            uniformRayDensity=uniformRayDensity,
            filamentBeam=filamentBeam,  # Single wavefront
            customField = fieldArray,
            targetOpenCL=targetOpenCL,
            # R0=R0,   # Near Field.
            eEpsilonX=eEpsilonX, eEpsilonZ=eEpsilonZ)  # Emittance [nmrad]
        beamLine.xlimits = [-2, 2]  # Horizontal limits of the plot [mm]
        beamLine.zlimits = [-2, 2]  # Vertical limits of the plot [mm]
        beamLine.xlimitsZoom = [-0.2, 0.2]
        beamLine.zlimitsZoom = [-0.2, 0.2]
        beamLine.xPrimelimits = [-0.3, 0.3]  # Angular limits of the plot [mrad]
    else:
        raise ValueError('Unknown source type!')
    
    if False:  # force zero source size:
        kwargs['eSigmaX'] = 0
        kwargs['eSigmaZ'] = 0
        kwargs['eEpsilonX'] = 0
        kwargs['eEpsilonZ'] = 0
        kwargs['eEspread'] = 0
    
        eEpsilonC = '0'
    else:
        eEpsilonC = 'n'
    
    kwargs['xPrimeMax'] = beamLine.xlimits[-1] / R0 * 1e3
    kwargs['zPrimeMax'] = beamLine.zlimits[-1] / R0 * 1e3
    
    beamLine.prefix = pprefix+'-{0}-{1}E-'.format(eEpsilonC, energyRange)
    if energyRange == 'mono':
        beamLine.eMinRays, beamLine.eMaxRays = E0-1, E0+1
        beamLine.eUnit = 'eV'
    elif energyRange == '1harmonic':
        beamLine.eMinRays, beamLine.eMaxRays = E0-300, E0+300
        beamLine.eUnit = 'eV'
    elif energyRange == 'smaller':
        beamLine.eMinRays, beamLine.eMaxRays = 1500, 7500
        beamLine.eUnit = 'keV'
    elif energyRange == 'wide':
        beamLine.eMinRays, beamLine.eMaxRays = 1500, 37500
        beamLine.eUnit = 'keV'
    
    kwargs['eMin'] = beamLine.eMinRays
    kwargs['eMax'] = beamLine.eMaxRays
    if Source == rs.UndulatorUrgent:
        kwargs['processes'] = 'half'

    beamLine.source = Source(
        beamLine, eN=1000, nx=40, nz=20, nrays=nrays, **kwargs)
    beamLine.fsm0 = rsc.Screen(beamLine, 'FSM0', (0, 0, 0))
    beamLine.fsm1 = rsc.Screen(beamLine, 'FSM1', (0, R0, 0))
    return beamLine


def run_process(beamLine):
    startTime = time.time()
    kw = {}
    if 'mono' in beamLine.prefix and isInternalSource:
        kw['fixedEnergy'] = beamLine.alignE
#    print(beamLine.what, isinstance(beamLine.source, rs.Undulator))

    if 'wave' in beamLine.what: # and isinstance(beamLine.source, rs.Undulator):
        wave1 = beamLine.fsm1.prepare_wave(
            beamLine.source, beamLine.fsmExpX, beamLine.fsmExpZ)
        wave1zoom = beamLine.fsm1.prepare_wave(
            beamLine.source, beamLine.fsmExpXzoom, beamLine.fsmExpZzoom)

        kw['wave'] = wave1
        beamSource = beamLine.source.shine(**kw)
        beamFSM0 = beamLine.fsm0.expose(beamSource)
        kw['wave'] = wave1zoom
        beamSourceZ = beamLine.source.shine(**kw)
        outDict = {'beamSource': beamSourceZ,
                   'beamFSM0': beamFSM0,
                   'beamFSM1': wave1, 'beamFSM1zoom': wave1zoom}
    else:
        beamSource = beamLine.source.shine(**kw)
        beamFSM0 = beamLine.fsm0.expose(beamSource)
        beamFSM1 = beamLine.fsm1.expose(beamSource)
        outDict = {'beamSource': beamSource,
                   'beamFSM0': beamFSM0,
                   'beamFSM1': beamFSM1, 'beamFSM1zoom': beamFSM1}
    print('shine time = {0}s'.format(time.time() - startTime))
    if showIn3D:
        beamLine.prepare_flow()
    return outDict

rr.run_process = run_process


def define_plots(beamLine):
    plots = []
    plotsE = []

    xaxis = xrtp.XYCAxis(r'$y$', 'mm', bins=256)
    yaxis = xrtp.XYCAxis(r'$x$', '$\mu$m', limits='symmetric')
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None)
    plot = xrtp.XYCPlot(
        'beamSource', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='YX source cross-section')
    plot.saveName = beamLine.prefix + 'crossectionYX' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$y$', 'mm', bins=256)
    yaxis = xrtp.XYCAxis(r'$z$', '$\mu$m', limits='symmetric')
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None)
    plot = xrtp.XYCPlot(
        'beamSource', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='YZ source cross-section')
    plot.saveName = beamLine.prefix + 'crossectionYZ' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', '$\mu$m', limits=beamLine.limitsFSM0X,
                         bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', '$\mu$m', limits=beamLine.limitsFSM0Z,
                         bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamSource', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='image at 0')
    plot.saveName = beamLine.prefix + 'fsm0' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)
    plotsE.append(plot)

    beam = 'beamFSM0'
    xaxis = xrtp.XYCAxis(r'$x$', '$\mu$m', limits=beamLine.limitsFSM0X)
    yaxis = xrtp.XYCAxis(r"$x'$", 'mrad', limits=beamLine.xPrimelimits)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None)
    plot = xrtp.XYCPlot(
        beam, (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='horizontal phase space')
    plot.saveName = beamLine.prefix + 'horPhaseSpace' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', '$\mu$m', limits=[-80, 80])
    yaxis = xrtp.XYCAxis(r"$x'$", 'mrad', limits=[-0.15, 0.15])
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None)
    plot = xrtp.XYCPlot(
        beam, (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='horizontal phase space zoomed')
    plot.saveName = beamLine.prefix + 'horPhaseSpaceZoom' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimits,
                         bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimits,
                         bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='total flux')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '1TotalFlux' + beamLine.suffix + '.png'
    #plot.persistentName = prefix + '1TotalFlux' + suffix + '.mat'
    plots.append(plot)
    plotsE.append(plot)
    ax = plot.xaxis
    edges = np.linspace(ax.limits[0], ax.limits[1], ax.bins+1)
    beamLine.fsmExpX = (edges[:-1] + edges[1:]) * 0.5 / ax.factor
    ax = plot.yaxis
    edges = np.linspace(ax.limits[0], ax.limits[1], ax.bins+1)
    beamLine.fsmExpZ = (edges[:-1] + edges[1:]) * 0.5 / ax.factor

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimitsZoom, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimitsZoom, bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1zoom', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        aspect='auto', title='total flux zoom')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '1TotalFluxZoom' + beamLine.suffix + '.png'
    plots.append(plot)
    plotsE.append(plot)
    ax = plot.xaxis
    edges = np.linspace(ax.limits[0], ax.limits[1], ax.bins+1)
    beamLine.fsmExpXzoom = (edges[:-1] + edges[1:]) * 0.5 / ax.factor
    ax = plot.yaxis
    edges = np.linspace(ax.limits[0], ax.limits[1], ax.bins+1)
    beamLine.fsmExpZzoom = (edges[:-1] + edges[1:]) * 0.5 / ax.factor

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimits, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimits, bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        fluxKind='s', aspect='auto', title='horizontal polarization flux')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '2horizFlux' + beamLine.suffix + '.png'
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimitsZoom,
                         bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimitsZoom,
                         bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1zoom', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        fluxKind='s', aspect='auto', title='horizontal polarization flux zoom')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '2horizFluxZoom' + beamLine.suffix + '.png'
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimits,
                         bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimits,
                         bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        fluxKind='p', aspect='auto', title='vertical polarization flux')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '3vertFlux' + beamLine.suffix + '.png'
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimitsZoom, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimitsZoom, bins=bins, ppb=ppb)
    caxis = xrtp.XYCAxis('energy', beamLine.eUnit, fwhmFormatStr=None,
                         bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1zoom', (1,), xaxis=xaxis, yaxis=yaxis, caxis=caxis,
        fluxKind='p', aspect='auto', title='vertical polarization flux zoom')
    plot.caxis.fwhmFormatStr = None
    plot.saveName = beamLine.prefix + '3vertFluxZoom' + beamLine.suffix + '.png'
    plots.append(plot)
    plotsE.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimits, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimits, bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1', (1,), xaxis=xaxis, yaxis=yaxis,
        caxis=xrtp.XYCAxis('degree of polarization', '',
                           data=raycing.get_polarization_degree,
                           limits=[0.9, 1.01], bins=bins, ppb=ppb),
        aspect='auto', title='degree of polarization')
    plot.saveName = beamLine.prefix + '4DegPol' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimitsZoom, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimitsZoom, bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1zoom', (1,), xaxis=xaxis, yaxis=yaxis,
        caxis=xrtp.XYCAxis('degree of polarization', '',
                           data=raycing.get_polarization_degree,
                           limits=[0.9, 1.01], bins=bins, ppb=ppb),
        aspect='auto', title='degree of polarization zoom')
    plot.saveName = beamLine.prefix + '4DegPolZoom' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimits, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimits, bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1', (1,), xaxis=xaxis, yaxis=yaxis,
        caxis=xrtp.XYCAxis('circular polarization rate', '',
                           data=raycing.get_circular_polarization_rate,
                           limits=[-1, 1], bins=bins, ppb=ppb),
        aspect='auto', title='circular polarization rate')
    plot.saveName = beamLine.prefix + '5CircPolRate' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)

    xaxis = xrtp.XYCAxis(r'$x$', 'mm', limits=beamLine.xlimitsZoom, bins=bins, ppb=ppb)
    yaxis = xrtp.XYCAxis(r'$z$', 'mm', limits=beamLine.zlimitsZoom, bins=bins, ppb=ppb)
    plot = xrtp.XYCPlot(
        'beamFSM1zoom', (1,), xaxis=xaxis, yaxis=yaxis,
        caxis=xrtp.XYCAxis('circular polarization rate', '',
                           data=raycing.get_circular_polarization_rate,
                           limits=[-1, 1], bins=bins, ppb=ppb),
        aspect='auto', title='circular polarization rate zoom')
    plot.saveName = beamLine.prefix + '5CircPolRateZoom' + beamLine.suffix + '.png'
    plot.caxis.fwhmFormatStr = None
    plots.append(plot)

    for plot in plotsE:
        f = plot.caxis.factor
        plot.caxis.limits = beamLine.eMinRays*f, beamLine.eMaxRays*f
    for plot in plots:
        plot.xaxis.fwhmFormatStr = '%.2f'
        plot.yaxis.fwhmFormatStr = '%.2f'
        plot.fluxFormatStr = '%.2p'
    return plots, plotsE


def afterScript(plots, beamLine):
    import os
    import pickle
    plot = plots[-1]
    flux = [plot.intensity, plot.nRaysAll, plot.nRaysSeeded]
    cwd = os.getcwd()
    pickleName = os.path.join(cwd, beamLine.prefix+'.pickle')
    with open(pickleName, 'wb') as f:
        pickle.dump((flux, plot.caxis.binEdges, plot.caxis.total1D), f,
                    protocol=2)


def main(**mainKWargs):
    beamLine = build_beamline(**mainKWargs)
    if showIn3D:
        beamLine.glow()
    else:
        plots, plotsE = define_plots(beamLine)
        xrtr.run_ray_tracing(plots, repeats=1,
#                             afterScript=afterScript, 
#                             afterScriptArgs=[plots, beamLine],
                             beamLine=beamLine)
        

if __name__ == '__main__':
    # ['custom','u', 'bm', 'eu', 'wu', 'tu']
    kwargs = {}
    for sourceType in ['eu', 'wu']:
        kwargs['sourceType'] = sourceType
        kwargs['what'] = 'wave'
        kwargs['energyRange'] = 'mono'
        kwargs['nrays'] = 1e6
        kwargs['targetOpenCL'] = None
        main(**kwargs)
        print(sourceType, "done")
