# -*- coding: utf-8 -*-
"""
qspec.physics
=============

Created on 01.04.2020

@author: Patrick Mueller

Module for physical functions useful for CLS.
"""

import string
import numpy as np
import scipy.constants as sc
import scipy.stats as st
import scipy.special as sp

from qspec._types import *
from qspec import tools

__all__ = ['L_LABEL', 'E_NORM', 'LEMNISCATE', 'mu_N', 'mu_B', 'g_s', 'me_u', 'me_u_d', 'gp_s', 'gn_s',
           'inv_cm_to_freq', 'freq_to_inv_cm', 'wavelength_to_freq', 'freq_to_wavelength', 'inv_cm_to_wavelength',
           'wavelength_to_inv_cm', 'beta', 'gamma', 'gamma_e', 'gamma_e_kin', 'e_rest', 'e_kin', 'e_total', 'e_el',
           'v_e', 'v_e_d1', 'v_el', 'v_el_d1', 'p_v', 'p_e', 'p_el', 'doppler', 'doppler_d1', 'doppler_e_d1',
           'doppler_el_d1', 'inverse_doppler', 'inverse_doppler_d1', 'alpha_atom', 'v_rec', 'photon_recoil',
           'photon_recoil_v', 'get_f', 'get_m', 'hyperfine', 'lande_n', 'lande_j', 'lande_f', 'zeeman', 'hyper_zeeman',
           'a_hyper_mu', 'saturation_intensity', 'saturation', 'rabi', 'scattering_rate', 'mass_factor', 'delta_r2',
           'delta_r4', 'delta_r6', 'lambda_r', 'lambda_rn', 'schmidt_line', 'sellmeier', 'gamma_3d', 'boost',
           'doppler_3d', 'gaussian_beam_3d', 'gaussian_doppler_3d', 'thermal_v_pdf', 'thermal_v_rvs', 'thermal_e_pdf',
           'thermal_e_rvs', 'convolved_boltzmann_norm_pdf', 'convolved_thermal_norm_v_pdf',
           'convolved_thermal_norm_f_pdf', 'convolved_thermal_norm_f_lin_pdf', 'source_energy_pdf']


L_LABEL = ['S', 'P', 'D', ] + list(string.ascii_uppercase[5:])
E_NORM = sc.e
LEMNISCATE = 2.6220575543
mu_N = sc.physical_constants['nuclear magneton'][0]
mu_B = sc.physical_constants['Bohr magneton'][0]
g_s = sc.physical_constants['electron g factor'][0]
me_u = sc.physical_constants['electron mass in u'][0]
me_u_d = sc.physical_constants['electron mass in u'][2]
gp_s = sc.physical_constants['proton g factor'][0]
gn_s = sc.physical_constants['neutron g factor'][0]


""" Units """


def inv_cm_to_freq(k: array_like):
    """
    :param k: The wavenumber k of a transition (cm ** -1)
    :returns: The frequency corresponding to the wavenumber k (MHz).
    """
    return k * sc.c * 1e-4


def freq_to_inv_cm(freq: array_like):
    """
    :param freq: The frequency f of a transition (MHz)
    :returns: The wavenumber k corresponding to the frequency freq (cm ** -1).
    """
    return freq / sc.c * 1e4


def wavelength_to_freq(lam: array_like):
    """
    :param lam: The wavelength lambda of a transition (um)
    :returns: The frequency corresponding to the wavelength lam (MHz).
    """
    return sc.c / lam


def freq_to_wavelength(freq: array_like):
    """
    :param freq: The frequency f of a transition (MHz)
    :returns: The wavelength corresponding to the frequency freq (um).
    """
    return sc.c / freq


def inv_cm_to_wavelength(k: array_like):
    """
    :param k: The wavenumber k of a transition (cm ** -1)
    :returns: The wavelength corresponding to the wavenumber k (um).
    """
    return 1e4 / k


def wavelength_to_inv_cm(lam: array_like):
    """
    :param lam: The wavelength lambda of a transition (um)
    :returns: The wavenumber k corresponding to the wavelength lam (cm ** -1).
    """
    return 1e4 / lam


""" 1-D kinematics """


def beta(v: array_like) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :returns: The velocity v relative to light speed.
    """
    v = np.asarray(v)
    return v / sc.c


def gamma(v: array_like) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :returns: The time-dilation/Lorentz factor corresponding to the velocity v.
    """
    return 1. / np.sqrt(1. - beta(v) ** 2)


def gamma_e(e: array_like, m: array_like) -> array_like:
    """
    :param e: The total energy of a body, including the energy of the rest mass (eV).
    :param m: The mass of the body (amu).
    :returns: The time-dilation/Lorentz factor corresponding to the total energy e of a body with mass m.
    """
    e = np.asarray(e)
    return e / e_rest(m)


def gamma_e_kin(e: array_like, m: array_like) -> array_like:
    """
    :param e: The kinetic energy of a body (eV).
    :param m: The mass of the body (amu).
    :returns: The time-dilation/Lorentz factor corresponding to the kinetic energy e of a body with mass m.
    """
    return 1. + gamma_e(e, m)


def e_rest(m: array_like) -> array_like:
    """
    :param m: The mass of a body (amu).
    :returns: The resting energy of the body with mass m (eV).
    """
    m = np.asarray(m)
    return m * sc.atomic_mass * sc.c ** 2 / E_NORM


def e_kin(v: array_like, m: array_like, relativistic=True) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :param m: The mass of the body (amu).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The kinetic energy of a body with velocity v and mass m (eV).
    """
    if relativistic:
        return (gamma(v) - 1.) * e_rest(m)
    else:
        v, m = np.asarray(v), np.asarray(m)
        return m * sc.atomic_mass * v ** 2 / 2. / E_NORM


def e_total(v: array_like, m: array_like) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :param m: The mass of the body (amu).
    :returns: The total energy of a body with velocity v and mass m (eV). """
    return gamma(v) * e_rest(m)


def e_el(u: array_like, q: array_like) -> array_like:
    """
    :param u: The electric potential difference (V).
    :param q: The charge of a body (e).
    :returns: The potential energy difference of a body with charge q inside an electric potential with voltage u (eV).
    """
    q, u = np.asarray(q), np.asarray(u)
    return q * u


def v_e(e: array_like, m: array_like, v0: array_like = 0, relativistic=True) -> array_like:
    """
    :param e: Energy which is added to the kinetic energy of a body with velocity v0 (eV).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The velocity of a body with mass m and velocity v0 after the addition of the kinetic energy e (m/s).
    """
    if relativistic:
        return sc.c * np.sqrt(1. - (1. / (gamma(v0) + gamma_e(e, m))) ** 2)
    else:
        v0, e, m = np.asarray(v0), np.asarray(e), np.asarray(m)
        return np.sqrt(v0 ** 2 + 2. * e * E_NORM / (m * sc.atomic_mass))


def v_e_d1(e: array_like, m: array_like, v0: array_like = 0, relativistic=True) -> array_like:
    """
    :param e: Energy which is added to the kinetic energy of a body with velocity v0 (eV).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The first derivative of 'v_e' regarding 'e' (m/(s eV)).
    """
    m = np.asarray(m)
    dv = 1. / (m * sc.atomic_mass * v_e(e, m, v0=v0, relativistic=relativistic))
    if relativistic:
        dv /= (gamma(v0) + gamma_e(e, m)) ** 3
    return dv * E_NORM


def v_el(u: array_like, q: array_like, m: array_like, v0: array_like = 0, relativistic=True) -> array_like:
    """
    :param u: The electric potential difference added to the kinetic energy of a body with velocity v0 (V).
    :param q: The charge of a body (e).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The velocity of a body with starting velocity v0, charge q and mass m
     after electrostatic acceleration with voltage u (m/s).
    """
    return v_e(e_el(u, q), m, v0=v0, relativistic=relativistic)


def v_el_d1(u: array_like, q: array_like, m: array_like, v0: array_like = 0, relativistic=True) -> array_like:
    """
    :param u: The electric potential difference added to the kinetic energy of a body with velocity v0 (V).
    :param q: The charge of a body (e).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The first derivative of 'v_el' regarding 'u' (m/(s V)).
    """
    return v_e_d1(e_el(u, q), m, v0=v0, relativistic=relativistic) * q


def p_v(v: array_like, m: array_like, relativistic=True) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :param m: The mass of the body (amu).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The momentum of a body with velocity v and mass m.
    """
    v, m = np.asarray(v), np.asarray(m)
    if relativistic:
        return gamma(v) * m * sc.atomic_mass * v
    else:
        return m * sc.atomic_mass * v


def p_e(e: array_like, m: array_like, p0: array_like = 0, relativistic=True) -> array_like:
    """
    :param e: Energy which is added to the kinetic energy of a body with velocity v0 (eV).
    :param m: The mass of the body (amu).
    :param p0: The initial momentum of the body (amu m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The momentum of a body with starting momentum p0 and mass m
     after the addition of the kinetic energy e (amu m/s).
    """
    e, p0 = np.asarray(e), np.asarray(p0)
    if relativistic:
        pc_square = (p0 * sc.c) ** 2 / E_NORM
        return np.sqrt(e ** 2 + pc_square + 2 * e * np.sqrt(pc_square + e_rest(m))) / sc.c
    else:
        m = np.asarray(m)
        return np.sqrt(p0 ** 2 + 2 * m * sc.atomic_mass * e * E_NORM)


def p_el(u: array_like, q: array_like, m: array_like, p0: array_like = 0, relativistic=True) -> array_like:
    """
    :param u: The electric potential difference added to the kinetic energy of a body with velocity v0 (V).
    :param q: The charge of a body (e).
    :param m: The mass of the body (amu).
    :param p0: The initial momentum of the body (amu m/s).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The momentum of a body with starting momentum p0, charge q and mass m
     after electrostatic acceleration with voltage u.
    """
    return p_e(e_el(u, q), m, p0, relativistic=relativistic)


def doppler(f: array_like, v: array_like, alpha: array_like, return_frame='atom') -> array_like:
    """
    :param f: The frequency of light (arb. units).
    :param v: The velocity of a body (m/s).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param return_frame: The coordinate system in which the frequency is returned. Can be either 'atom' or 'lab'.
    :returns: the Doppler-shifted frequency in either the rest frame of the atom or the laboratory frame ([f]).
    :raises ValueError: rest_frame must be either 'atom' or 'lab'.
    """
    f, alpha = np.asarray(f), np.asarray(alpha)
    if return_frame == 'atom':
        """ Return freq in the atomic system, alpha=0 == Col, alpha in laboratory system """
        return f * gamma(v) * (1. - beta(v) * np.cos(alpha))
    elif return_frame == 'lab':
        """ Return freq in the laboratory system, alpha=0 == Col, alpha in laboratory system """
        return f / (gamma(v) * (1. - beta(v) * np.cos(alpha)))
    else:
        raise ValueError('rest_frame must be either "atom" or "lab".')


def doppler_d1(f: array_like, v: array_like, alpha: array_like, return_frame='atom') -> array_like:
    """
    :param f: The frequency of light (arb. units).
    :param v: The velocity of a body (m/s).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param return_frame: The coordinate system in which the frequency is returned. Can be either 'atom' or 'lab'.
    :returns: the first derivative of the 'doppler' formula regarding 'v' ([f] s/m).
    :raises ValueError: rest_frame must be either 'atom' or 'lab'.
    """
    if return_frame == 'atom':
        """ Return df/dv in the atomic system, alpha=0 == Col, alpha in laboratory system. """
        f, alpha = np.asarray(f), np.asarray(alpha)
        return f * gamma(v) ** 3 * (beta(v) - np.cos(alpha)) / sc.c
    elif return_frame == 'lab':
        """ Return df/dv in the laboratory system, alpha=0 == Col, alpha in laboratory system. """
        f_lab = doppler(f, v, alpha, return_frame='lab')
        return -f_lab / f * doppler_d1(f_lab, v, alpha, return_frame='atom')
    else:
        raise ValueError('rest_frame must be either "atom" or "lab".')


def doppler_e_d1(f: array_like, alpha: array_like, e: array_like, m: array_like,
                 v0: array_like = 0, return_frame='atom', relativistic=True) -> array_like:
    """
    :param f: The frequency of light (arb. units).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param e: Energy which is added to the kinetic energy of a body with velocity v0 (eV).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param return_frame: The coordinate system in which the frequency is returned. Can be either 'atom' or 'lab'.
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: the first derivative of the 'doppler' formula regarding 'e' ([f]/eV).
    :raises ValueError: rest_frame must be either 'atom' or 'lab'.
    """
    v = v_e(e, m, v0=v0, relativistic=relativistic)
    return doppler_d1(f, v, alpha, return_frame=return_frame) * v_e_d1(e, m, v0=v0, relativistic=relativistic)


def doppler_el_d1(f: array_like, alpha: array_like, u: array_like, q: array_like, m: array_like,
                  v0: array_like = 0., return_frame='atom', relativistic=True) -> array_like:
    """
    :param f: The frequency of light (arb. units).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param u: The electric potential difference added to the kinetic energy of a body with velocity v0 (V).
    :param q: The charge of a body (e).
    :param m: The mass of the body (amu).
    :param v0: The initial velocity of the body (m/s).
    :param return_frame: The coordinate system in which the frequency is returned. Can be either 'atom' or 'lab'.
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: the first derivative of the 'doppler' formula regarding 'u' ([f]/V).
    :raises ValueError: rest_frame must be either 'atom' or 'lab'.
    """
    v = v_el(u, q, m, v0=v0, relativistic=relativistic)
    return doppler_d1(f, v, alpha, return_frame=return_frame) * v_el_d1(u, q, m, v0=v0, relativistic=relativistic)


def inverse_doppler(f_atom: array_like, f_lab: array_like, alpha: array_like,
                    mode='raise-raise', return_mask=False) -> array_like:
    """
    :param f_lab: The frequency of light in the laboratory frame (arb. units).
    :param f_atom: The frequency of light in the atoms rest frame ([f_lab]).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param mode: The mode how to handle numpy.nans and ambiguous velocities. Available options are:
        * 'raise-raise': Raise an error if there are numpy.nans and if the velocity is ambiguous.
        * 'raise-small': Raise an error if there are numpy.nans and return the smaller velocity.
        * 'raise-large': Raise an error if there are numpy.nans and return the larger velocity.
        * 'isnan-raise': Ignore numpy.nans and raise an error if the velocity is ambiguous.
        * 'isnan-small': Ignore numpy.nans and return the smaller velocity.
        * 'isnan-large': Ignore numpy.nans and return the larger velocity.
    :param return_mask: Whether the mask where the velocity is ambiguous is returned as a second argument.
    :returns: the velocity required to shift f_lab to f_atom (m/s)
     and optionally the mask where the velocity is ambiguous.
    """
    modes = ['raise-raise', 'raise-small', 'raise-large', 'isnan-raise', 'isnan-small', 'isnan-large']
    if mode not in modes:
        raise ValueError('mode must be in {}.'.format(modes))
    f_lab, f_atom, alpha = np.asarray(f_lab), np.asarray(f_atom), np.asarray(alpha)
    scalar_true = tools.check_shape((), f_lab, f_atom, alpha, return_mode=True)
    if scalar_true:
        alpha = np.array([alpha])
    cos = np.cos(alpha)
    square_sum = (f_atom / f_lab) ** 2 + cos ** 2
    nan = square_sum < 1.
    np.seterr(invalid='ignore')
    bet1 = (cos + f_atom / f_lab * np.sqrt(square_sum - 1.)) / square_sum
    bet2 = (cos - f_atom / f_lab * np.sqrt(square_sum - 1.)) / square_sum
    np.seterr(invalid='warn')
    bet1[nan] = 2.
    bet2[nan] = 2.
    mask1 = np.abs(0.5 - bet1) < 0.5
    mask1 += bet1 == 0.
    mask2 = np.abs(0.5 - bet2) < 0.5
    mask2 += bet2 == 0.
    ambiguous = ~(~mask1 + ~mask2)
    nan = ~(mask1 + mask2)
    bet = np.zeros_like(square_sum)
    bet[nan] = np.nan
    if mode[6:] in ['small', 'raise']:
        bet[mask1] = bet1[mask1]
        bet[mask2] = bet2[mask2]
    elif mode[6:] == 'large':
        bet[mask2] = bet2[mask2]
        bet[mask1] = bet1[mask1]
    if np.any(nan):
        if mode[:5] == 'raise':
            raise ValueError('Situation is physically impossible for at least one argument.')
        # print('WARNING: Situation is physically impossible for at least one argument.')
    if np.any(ambiguous):
        if mode[6:] == 'raise':
            raise ValueError('Situation allows two different velocities.')
        # print('WARNING: Situation allows two different velocities. Taking the {} velocity.'.format(mode[6:]))
    if return_mask:
        if scalar_true:
            return (bet * sc.c)[0], ambiguous
        return bet * sc.c, ambiguous
    if scalar_true:
        return (bet * sc.c)[0]
    return bet * sc.c


def inverse_doppler_d1(f_atom: array_like, f_lab: array_like, alpha: array_like,
                       mode='raise-raise', return_mask=False) -> array_like:
    """
    :param f_lab: The frequency of light in the laboratory frame (arb. units).
    :param f_atom: The frequency of light in the atoms rest frame ([f_lab]).
    :param alpha: The angle between the velocity- and the wave-vector in the laboratory frame (rad).
    :param mode: The mode how to handle numpy.nans and ambiguous velocities. Available options are:
        * 'raise-raise': Raise an error if there are numpy.nans and if the velocity is ambiguous.
        * 'raise-small': Raise an error if there are numpy.nans and return the smaller velocity.
        * 'raise-large': Raise an error if there are numpy.nans and return the larger velocity.
        * 'isnan-raise': Ignore numpy.nans and raise an error if the velocity is ambiguous.
        * 'isnan-small': Ignore numpy.nans and return the smaller velocity.
        * 'isnan-large': Ignore numpy.nans and return the larger velocity.
    :param return_mask: Whether the mask where the velocity is ambiguous is returned as a second argument.
    :returns: the first derivative of the 'inverse_doppler' formula regarding f_atom (m/(s MHz)).
    """
    f_lab, f_atom, alpha = np.asarray(f_lab), np.asarray(f_atom), np.asarray(alpha)
    scalar_true = tools.check_shape((), f_lab, f_atom, alpha, return_mode=True)
    if scalar_true:
        alpha = np.array([alpha])
    v, ambiguous = inverse_doppler(f_atom, f_lab, alpha, mode=mode, return_mask=True)
    cos = np.cos(alpha)
    square_sum = (f_atom / f_lab) ** 2 + cos ** 2
    np.seterr(invalid='ignore')
    bet = np.sqrt(square_sum - 1.) + (f_atom / f_lab) ** 2 / np.sqrt(square_sum - 1.)
    np.seterr(invalid='warn')
    if mode[6:] in ['small', 'raise']:
        bet[ambiguous] = -bet[ambiguous]
    bet += -2. * v * (f_atom / f_lab) / sc.c
    if return_mask:
        if scalar_true:
            return (bet * sc.c / (f_lab * square_sum))[0], ambiguous
        return bet * sc.c / (f_lab * square_sum), ambiguous
    if scalar_true:
        return (bet * sc.c / (f_lab * square_sum))[0]
    return bet * sc.c / (f_lab * square_sum)


def alpha_atom(alpha: array_like, v: array_like) -> array_like:
    """
    :param alpha: The angle between a velocity- and a wave-vector in the laboratory frame (rad).
    :param v: The velocity of a body (m/s).
    :returns: The angle between the velocity- and the wave-vector in the atoms rest frame (rad).
    """
    alpha = np.asarray(alpha)
    cos = np.cos(alpha)
    arg = (beta(v) + cos) / (1. + beta(v) * cos)
    return np.arccos(arg)


def v_rec(f, m) -> array_like:
    """
    :param f: The frequency of light in the atoms rest frame (MHz).
    :param m: The mass of a body (amu).
    :returns: The change of velocity of an atom at rest
     due to the absorption of a photon with frequency f (m/s).
    """
    f, m = np.asarray(f), np.asarray(m)
    return sc.h * f / (m * sc.atomic_mass * sc.c) * 1e6


def photon_recoil(f: array_like, m: array_like) -> array_like:
    """
    :param f: The frequency of light in the atoms rest frame (MHz).
    :param m: The mass of a body (amu).
    :returns: The change of a transition frequency of an atom at rest
     due to the absorption of a photon with frequency f (MHz).
    """
    f, m = np.asarray(f), np.asarray(m)
    return (sc.h * (f * 1e6) ** 2) / (2 * m * sc.atomic_mass * sc.c ** 2) * 1e-6


def photon_recoil_v(v: array_like, alpha: array_like, f_lab: array_like, m: array_like) -> array_like:
    """
    :param v: The velocity of a body (m/s).
    :param alpha: The angle between a velocity- and a wave-vector in the laboratory frame (rad).
    :param f_lab: The frequency of light in the laboratory frame (MHz).
    :param m: The mass of a body (amu).
    :returns: The change of a transition frequency of an atom moving with velocity v
     due to the absorption of a laser photon with frequency f (MHz).
    """
    f = doppler(f_lab, v, alpha, return_frame='atom')
    return photon_recoil(f, m)


""" Atomic physics """


def get_f(i: scalar, j: scalar):
    """
    :param i: The nuclear spin quantum number I.
    :param j: The electronic total angular momentum quantum number J.
    :returns: All possible f quantum numbers.
    """
    return [k + abs(i - j) for k in range(int(i + j + 1 - abs(i - j)))]


def get_m(f: scalar):
    """
    :param f: The total angular momentum quantum number F (= J if I == 0).
    :returns: All possible zeeman substates of the specified quantum number.
    """
    return [k - f for k in range(int(2 * f + 1))]


def hyperfine(i: float, j: float, f: float, a: array_like, b: array_like = None) -> array_like:
    """
    :param i: The nuclear spin quantum number I.
    :param j: The electronic total angular momentum quantum number J.
    :param f: The total angular momentum quantum number F.
    :param a: The magnetic dipole hyperfine constant A (arb. units).
    :param b: The electric quadrupole hyperfine constant B ([a]).
    :returns: The hyperfine shift of an atomic state (i, j, f) with hyperfine constants a and b ([a]).
    """
    a = np.asarray(a)
    if b is None:
        b = np.zeros_like(a)
    if i < 0. or j < 0. or f < 0.:
        print('Either i >= 0, j >= 0 or f >= 0 is not fulfilled.')
        raise ValueError
    if f < abs(i - j) or f > i + j:
        print('|i - j| <= f <= i + j must be fulfilled.')
        raise ValueError
    if i == 0. or j == 0.:
        return 0.
    k = f * (f + 1) - i * (i + 1) - j * (j + 1)
    shift = a * k / 2
    if i > 0.5 and j > 0.5:
        k_2 = 3 * k * (k + 1) / 2 - 2 * i * (i + 1) * j * (j + 1)
        k_2 /= i * (2 * i - 1) * j * (2 * j - 1)
        shift += b * k_2 / 4
    return shift


def lande_n(gyro: array_like) -> array_like:
    """
    :param gyro: The gyromagnetic ratio (MHz).
    :returns: The nuclear g-factor.
    """
    gyro = np.asarray(gyro)
    return gyro * sc.h / mu_N


def lande_j(s: float, ll: float, j: float, approx_g_s: bool = False) -> float:
    """
    :param s: The electron spin quantum number S.
    :param ll: The electronic angular momentum quantum number L.
    :param j: The electronic total angular momentum quantum number J.
    :param approx_g_s: Whether to use g_s = -2 or the QED result g_s = -2.0023... .
    :returns: The electronic g-factor.
    """
    if j == 0:
        return 0.
    g = -2 if approx_g_s else g_s
    jj = j * (j + 1)
    ls = ll * (ll + 1) - s * (s + 1)
    val = -(jj + ls) / (2 * jj)
    val += (jj - ls) / (2 * jj) * g
    return val


def lande_f(i: float, j: float, f: float, g_n: float, g_j: float) -> float:
    """
    :param i: The nuclear spin quantum number I.
    :param j: The electronic total angular momentum quantum number J.
    :param f: The total angular momentum quantum number F.
    :param g_n: The nuclear g-factor.
    :param g_j: The electronic g-factor.
    :returns: The hyperfine structure g-factor.
    """
    ff = f * (f + 1.)
    ji = j * (j + 1.) - i * (i + 1.)
    val = (ff + ji) / (2 * ff) * g_j
    val += (ff - ji) / (2 * ff) * g_n * mu_N / mu_B
    return val


def zeeman(m: float, b: array_like, g: float, as_freq=True) -> array_like:
    """
    :param m: The B-field-axis component quantum number m of the total angular momentum.
    :param b: The B-field (T).
    :param g: The g-factor.
    :param as_freq: The zeeman shift can be returned in energy or frequency units.
    :returns: The energy shift of an atomic state due to the zeeman effect in energy or frequency units (eV or MHz).
    """
    b = np.asarray(b)
    delta = -g * m * mu_B * b / E_NORM
    if as_freq:
        delta /= sc.h * 1e6 / E_NORM
    return delta


def hyper_zeeman(i: float, s: float, ll: float, j: float, f: float, m: float, g_n: float,
                 a_hyper: array_like, b_hyper: array_like, b: array_like,
                 g_n_as_gyro: bool = False, as_freq: bool = True) -> array_like:
    """
    :param i: The nuclear spin quantum number I.
    :param s: The electron spin quantum number S.
    :param ll: The electronic angular momentum quantum number L.
    :param j: The electronic total angular momentum quantum number J.
    :param f: The total angular momentum quantum number F.
    :param m: The B-field-axis component quantum number m of the total angular momentum.
    :param g_n: The nuclear g-factor or the gyromagnetic ratio if g_n_as_gyro == True.
    :param a_hyper: The magnetic dipole hyperfine constant A (eV or MHz).
    :param b_hyper: The electric quadrupole hyperfine constant B ([a]).
    :param b: The B-field (T).
    :param g_n_as_gyro: Whether g_n is the nuclear g-factor or the gyromagnetic ratio.
    :param as_freq: The shift can be returned in energy or frequency units.
    :returns: The total energy shift of an atomic state due to the hyperfine splitting
     and the zeeman effect in energy or frequency units (eV or MHz).
    """
    a_hyper, b_hyper, b = np.asarray(a_hyper), np.asarray(b_hyper), np.asarray(b)
    g_j = lande_j(s, ll, j)
    g_f = lande_f(i, j, f, lande_n(g_n) if g_n_as_gyro else g_n, g_j) if i != 0 or j != 0 else 0.
    shift = hyperfine(i, j, f, a_hyper, b_hyper)
    shift += zeeman(m, b, g_f, as_freq=as_freq)
    return shift


def a_hyper_mu(i: scalar, j: scalar, mu: array_like, b: array_like):
    """
    :param i: The nuclear spin quantum number I.
    :param j: The electronic total angular momentum quantum number J.
    :param mu: The magnetic moment of the nucleus in units of the nuclear magneton (mu_N).
    :param b: The B-field of the atomic electrons at the nucleus (T).
    :returns: The hyperfine structure constant A (MHz).
    """
    mu, b = np.asarray(mu), np.asarray(b)
    if i == 0 or j == 0:
        return np.zeros_like(mu * b)
    return mu * b / np.sqrt(i * (i + 1) * j * (j + 1)) / sc.h


def saturation_intensity(f: array_like, a: array_like, a_dipole: array_like):
    """
    :param f: The frequency of the transition (MHz).
    :param a: The Einstein A coefficient (MHz).
    :param a_dipole: The reduced dipole coefficient of the transition (see algebra.a_dipole).
    :returns: The saturation intensity.
    """
    f, a, a_dipole = np.asarray(f), np.asarray(a), np.asarray(a_dipole)
    return np.pi * (f * 1e6) ** 3 * sc.h * a * 1e6 / (3 * sc.c ** 2 * a_dipole)


def saturation(i: array_like, f: array_like, a: array_like, a_dipole: array_like):
    """
    :param i: The intensity of the laser (MHz).
    :param f: The frequency of the transition (MHz).
    :param a: The Einstein A coefficient (MHz).
    :param a_dipole: The reduced dipole coefficient of the transition (see algebra.a_dipole).
    :returns: The saturation parameter.
    """
    i = np.asarray(i)
    return i / saturation_intensity(f, a, a_dipole)


def rabi(a: array_like, s: array_like):
    """
    :param a: The Einstein A coefficient (MHz).
    :param s: The saturation parameter.
    :returns: The rabi frequency.
    """
    a, s = np.asarray(a), np.asarray(s)
    return a * np.sqrt(s / 2.)


def scattering_rate(df: array_like, a: array_like, s: array_like):
    """
    :param df: The detuning of to be scattered light from the transition.
     This must be differences of real frequencies, such that w = 2 pi * df (MHz).
    :param a: The Einstein A coefficient (MHz).
    :param s: The saturation parameter.
    :returns: The 2-state-equilibrium scattering-rate of an electronic transition.
    """
    df, a, s = np.asarray(df), np.asarray(a), np.asarray(s)
    return 0.125 * s * a ** 3 / (0.25 * (1 + s) * a ** 2 + (2 * np.pi * df) ** 2)


def mass_factor(m: array_like, m_ref: array_like, m_d: array_like = 0, m_ref_d: array_like = 0, k_inf: bool = True) \
        -> (ndarray, ndarray):
    """
    :param m: The mass of the isotope (amu).
    :param m_ref: The mass of the reference isotope (amu). Must be a scalar or have the same shape as 'm'.
    :param m_d: The uncertainty of the mass of the isotope (amu). Must be a scalar or have the same shape as 'm'.
    :param m_ref_d: The uncertainty of the mass of the reference isotope (amu).
     Must be a scalar or have the same shape as 'm'.
    :param k_inf: Whether the normal mass-shift factor K(NMS) is defined mass independently
     as m_e * T(inf) (= True) or as m_e * T(A_ref) (= False). Compare (6.4) with (3.17)
     in [W. H. King, Isotope shifts in atomic spectra (1984)].
    :returns: the mass factor and its uncertainty needed to calculate modified isotope shifts or charge radii.
    """
    m, m_d, m_ref, m_ref_d = np.asarray(m), np.asarray(m_d), np.asarray(m_ref), np.asarray(m_ref_d)
    if k_inf:
        mu = (m + me_u) * (m_ref + me_u) / (m - m_ref)
        if np.all(m_d) == 0 and np.all(m_ref_d) == 0:
            return mu, np.zeros_like(mu)
        mu_d = ((mu / (m + me_u) - mu / (m - m_ref)) * m_d) ** 2
        mu_d += ((mu / (m_ref + me_u) + mu / (m - m_ref)) * m_ref_d) ** 2
        mu_d += ((mu / (m + me_u) + mu / (m_ref + me_u)) * me_u_d) ** 2
    else:
        mu = (m + me_u) * m_ref / (m - m_ref)
        if np.all(m_d) == 0 and np.all(m_ref_d) == 0:
            return mu, np.zeros_like(mu)
        mu_d = (-m_ref * (m_ref + me_u) / ((m - m_ref) ** 2) * m_d) ** 2
        mu_d += (m * (m + me_u) / ((m - m_ref) ** 2) * m_ref_d) ** 2
        mu_d += (m_ref / (m - m_ref) * me_u_d) ** 2
    return mu, np.sqrt(mu_d)


def delta_r2(r: array_like, r_d: array_like, r_ref: array_like, r_ref_d: array_like,
             delta_r: array_like, delta_r_d: array_like, v2: array_like, v2_ref: array_like):
    """
    :param r: The Barrett radius of an isotope.
    :param r_d: The uncertainty of the Barrett radius.
    :param r_ref: The Barrett radius of a reference isotope.
    :param r_ref_d: The uncertainty of the Barrett radius of the reference isotope.
    :param delta_r: The difference between the Barrett radius of the isotope and the reference isotope.
    :param delta_r_d: The uncertainty of the difference between
     the Barrett radius of the isotope and the reference isotope.
    :param v2: The V2 factor of the isotope.
    :param v2_ref: The V2 factor of the reference isotope.
    :returns: The difference of the mean square nuclear charge radius between two isotopes and its uncertainty.
    """
    r, r_d = np.asarray(r, dtype=float), np.asarray(r_d, dtype=float)
    r_ref, r_ref_d = np.asarray(r_ref, dtype=float), np.asarray(r_ref_d, dtype=float)
    delta_r, delta_r_d = np.asarray(delta_r, dtype=float), np.asarray(delta_r_d, dtype=float)
    v2, v2_ref = np.asarray(v2, dtype=float), np.asarray(v2_ref, dtype=float)

    sum_term = (r / v2 + r_ref / v2_ref) / v2
    delta_term = delta_r + r_ref * (1. - v2 / v2_ref)
    val = sum_term * delta_term  # (r/v2)**2 - (r_ref/v2_ref)**2

    err = (sum_term * delta_r_d) ** 2
    err += (delta_term * r_d / (v2 ** 2)) ** 2
    err += ((delta_term / (v2 * v2_ref) + sum_term * (1. - v2 / v2_ref)) * r_ref_d) ** 2
    return val, np.sqrt(err)


def delta_r4(r: array_like, r_d: array_like, r_ref: array_like, r_ref_d: array_like,
             delta_r: array_like, delta_r_d: array_like, v4: array_like, v4_ref: array_like):
    """
    :param r: The Barrett radius of an isotope.
    :param r_d: The uncertainty of the Barrett radius.
    :param r_ref: The Barrett radius of a reference isotope.
    :param r_ref_d: The uncertainty of the Barrett radius of the reference isotope.
    :param delta_r: The difference between the Barrett radius of the isotope and the reference isotope.
    :param delta_r_d: The uncertainty of the difference between
     the Barrett radius of the isotope and the reference isotope.
    :param v4: The V4 factor of the isotope.
    :param v4_ref: The V4 factor of the reference isotope.
    :returns: The difference of the mean quartic nuclear charge radius between two isotopes and its uncertainty.
    """
    r, r_d = np.asarray(r, dtype=float), np.asarray(r_d, dtype=float)
    r_ref, r_ref_d = np.asarray(r_ref, dtype=float), np.asarray(r_ref_d, dtype=float)
    delta_r, delta_r_d = np.asarray(delta_r, dtype=float), np.asarray(delta_r_d, dtype=float)
    v4, v4_ref = np.asarray(v4, dtype=float), np.asarray(v4_ref, dtype=float)

    sum_term = (r / v4) ** 2 + (r_ref / v4_ref) ** 2
    delta_term = delta_r2(r, r_d, r_ref, r_ref_d, delta_r, delta_r_d, v4, v4_ref)
    val = sum_term * delta_term[0]  # (r/v4)**4 - (r_ref/v4_ref)**4

    err = (sum_term * delta_term[1]) ** 2
    err += (2. * delta_term[0] * r * r_d / (v4 ** 2)) ** 2
    err += (2. * delta_term[0] * r_ref * r_ref_d / (v4_ref ** 2)) ** 2
    return val, np.sqrt(err)


def delta_r6(r: array_like, r_d: array_like, r_ref: array_like, r_ref_d: array_like,
             delta_r: array_like, delta_r_d: array_like, v6: array_like, v6_ref: array_like):
    """
    :param r: The Barrett radius of an isotope.
    :param r_d: The uncertainty of the Barrett radius.
    :param r_ref: The Barrett radius of a reference isotope.
    :param r_ref_d: The uncertainty of the Barrett radius of the reference isotope.
    :param delta_r: The difference between the Barrett radius of the isotope and the reference isotope.
    :param delta_r_d: The uncertainty of the difference between
     the Barrett radius of the isotope and the reference isotope.
    :param v6: The V6 factor of the isotope.
    :param v6_ref: The V6 factor of the reference isotope.
    :returns: The difference of the mean sextic nuclear charge radius between two isotopes and its uncertainty.
    """
    r, r_d = np.asarray(r, dtype=float), np.asarray(r_d, dtype=float)
    r_ref, r_ref_d = np.asarray(r_ref, dtype=float), np.asarray(r_ref_d, dtype=float)
    delta_r, delta_r_d = np.asarray(delta_r, dtype=float), np.asarray(delta_r_d, dtype=float)
    v6, v6_ref = np.asarray(v6, dtype=float), np.asarray(v6_ref, dtype=float)

    sum_term = (v6 / r) * ((r / v6) ** 3 + (r_ref / v6_ref) ** 3)
    delta = delta_r4(r, r_d, r_ref, r_ref_d, delta_r, delta_r_d, v6, v6_ref)
    delta_term = delta[0] + (r_ref / v6_ref) ** 4 * (1. - (r / v6) * (v6_ref / r_ref))
    val = sum_term * delta_term  # (r/v6)**6 - (r_ref/v6_ref)**6

    err = (sum_term * delta[1]) ** 2
    err += ((-(r_ref / v6_ref) ** 3 * sum_term / v6
             + delta_term * (-sum_term / r + 3. * r / (v6 ** 2))) * r_d) ** 2
    err += (((4 * r_ref ** 3 / (v6_ref ** 4) * (1. - (r / v6) * (v6_ref / r_ref))
              + (r / v6) * r_ref ** 2 / (v6_ref ** 3)) * sum_term
             + delta_term * 3. * (v6 / r) * r_ref ** 2 / (v6_ref ** 3)) * r_ref_d) ** 2
    return val, np.sqrt(err)


def lambda_r(r: float, r_d: float, r_ref: float, r_ref_d: float, delta_r: float, delta_r_d: float,
             v2: float, v2_ref: float, v4: float, v4_ref: float, v6: float, v6_ref: float, c2c1: float, c3c1: float):
    """
    :param r: The Barrett radius of an isotope.
    :param r_d: The uncertainty of the Barrett radius.
    :param r_ref: The Barrett radius of a reference isotope.
    :param r_ref_d: The uncertainty of the Barrett radius of the reference isotope.
    :param delta_r: The difference between the Barrett radius of the isotope and the reference isotope.
    :param delta_r_d: The uncertainty of the difference between
     the Barrett radius of the isotope and the reference isotope.
    :param v2: The V2 factor of the isotope.
    :param v2_ref: The V2 factor of the reference isotope.
    :param v4: The V4 factor of the isotope.
    :param v4_ref: The V4 factor of the reference isotope.
    :param v6: The V6 factor of the isotope.
    :param v6_ref: The V6 factor of the reference isotope.
    :param c2c1: Seltzer's coefficient for the quartic moment.
    :param c3c1: Seltzer's coefficient for the sextic moment.
    :returns: The difference of the mean sextic nuclear charge radius between two isotopes and its uncertainty.
    """
    r2 = delta_r2(r, r_d, r_ref, r_ref_d, delta_r, delta_r_d, v2, v2_ref)
    r4 = delta_r4(r, r_d, r_ref, r_ref_d, delta_r, delta_r_d, v4, v4_ref)
    r6 = delta_r6(r, r_d, r_ref, r_ref_d, delta_r, delta_r_d, v6, v6_ref)
    return lambda_rn(r2[0], r2[1], r4[0], r4[1], r6[0], r6[1], c2c1, c3c1)


def lambda_rn(r_2: float, r_2_d: float, r_4: float, r_4_d: float, r_6: float, r_6_d: float, c2c1: float, c3c1: float):
    """
    :param r_2: The difference of the mean square nuclear charge radius between two isotopes.
    :param r_2_d: The uncertainty of the difference of the mean square nuclear charge radius.
    :param r_4: The difference of the mean quartic nuclear charge radius between two isotopes.
    :param r_4_d: The uncertainty of the difference of the mean quartic nuclear charge radius.
    :param r_6: The difference of the mean sextic nuclear charge radius between two isotopes.
    :param r_6_d: The uncertainty of the difference of the mean sextic nuclear charge radius.
    :param c2c1: Seltzer's coefficient for the quartic moment.
    :param c3c1: Seltzer's coefficient for the sextic moment.
    :returns: the Lambda observable for the given differences in mean square, quartic and sextic nuclear charge radii
     and its uncertainty.
    """
    val = r_2 + c2c1 * r_4 + c3c1 * r_6
    err = r_2_d ** 2
    err += (c2c1 * r_4_d) ** 2
    err += (c3c1 * r_6_d) ** 2
    return val, np.sqrt(err)


def schmidt_line(ll, j, is_proton):
    _g_s = gp_s if is_proton else gn_s
    _g_l = 1 if is_proton else 0
    if j < ll:
        return j / (j + 1) * ((ll + 1) * _g_l - 0.5 * _g_s)
    return ll * _g_l + 0.5 * _g_s


""" Optics """


def sellmeier(w: array_like, a: array_iter, b: array_iter):
    """
    :param w: The wavelength in µm.
    :param a: The a coefficients.
    :param b: The b coefficients.
    :return: The index of refraction for the wavelength w and the given material.
    """
    a, b = np.asarray(a), np.asarray(b)
    tools.check_dimension(a.shape[0], 0, b)
    sum_term = np.sum([a_i * w ** 2 / (w ** 2 - b_i) for a_i, b_i in zip(a, b)], axis=0)
    return np.sqrt(1 + sum_term)


""" 3-D kinematics """


def gamma_3d(v: array_like, axis=-1) -> array_like:
    """
    :param v: The velocity 3-vector (m/s).
    :param axis: The axis along which the vector components are aligned.
    :returns: The time-dilation/Lorentz factor corresponding to the velocity vector v.
    :raises ValueError: v must have 3 components along the specified axis.
    """
    tools.check_dimension(3, axis, v)
    return gamma(tools.absolute(v, axis=axis))


def boost(x: array_like, v: array_like, axis=-1) -> array_like:
    """
    :param x: The 4-vector x in the current rest frame (arb. units).
    :param v: The velocity 3-vector (m/s).
    :param axis: The axis along which the vector components are aligned.
    :returns: The 4-vector x in the coordinate system moving with velocity v relative to the current rest frame ([x]).
    :raises ValueError: x and v must have 4 and 3 components along the specified axis, respectively.
     The shapes of x and v must be compatible.
    """
    x, v = np.asarray(x), np.asarray(v)
    tools.check_dimension(4, axis, x)
    tools.check_dimension(3, axis, v)
    bet = beta(v)
    bet_abs = beta(tools.absolute(v, axis=axis))
    tools.check_shape_like(np.sum(x, axis=axis), bet_abs, allow_scalar=False)
    bet_abs[bet_abs == 0] = 1

    gam = gamma_3d(v, axis=axis)
    b_xyz = np.array([[1. + (gam - 1.) * np.take(bet, i, axis=axis) * np.take(bet, j, axis=axis) / (bet_abs ** 2)
                       if i == j else (gam - 1.) * np.take(bet, i, axis=axis) * np.take(bet, j, axis=axis)
                       / (bet_abs ** 2) for j in range(3)] for i in range(3)])
    b = np.array([[gam, -gam * np.take(bet, 0, axis=axis), -gam * np.take(bet, 1, axis=axis),
                   -gam * np.take(bet, 2, axis=axis)],
                  [-gam * np.take(bet, 0, axis=axis), b_xyz[0, 0], b_xyz[0, 1], b_xyz[0, 2]],
                  [-gam * np.take(bet, 1, axis=axis), b_xyz[1, 0], b_xyz[1, 1], b_xyz[1, 2]],
                  [-gam * np.take(bet, 2, axis=axis), b_xyz[2, 0], b_xyz[2, 1], b_xyz[2, 2]]])
    axes = list(range(len(v.shape)))
    axes.insert(0, axes.pop(axis))
    x = np.transpose(x, axes=axes)
    y = np.array([np.sum(b[i] * x, axis=0) for i in range(4)])
    axes = list(range(1, len(axes)))
    axes.insert(axis, 0) if axis != -1 else axes.append(0)
    return np.transpose(y, axes=axes)


def doppler_3d(k: array_like, v: array_like, return_frame='atom', axis=-1) -> array_like:
    """
    :param k: The k-wave-3-vector of light (arb. units).
    :param v: The velocity 3-vector (m/s).
    :param return_frame: The coordinate system in which the frequency is returned. Can be either 'atom' or 'lab'.
    :param axis: The axis along which the vector components are aligned.
    :returns: the Doppler-shifted k-wave-4-vector in either the rest frame of the atom or the laboratory frame ([k]).
    :raises ValueError: rest_frame must be either 'atom' or 'lab'. The shapes of k and v must be compatible.
    """
    k, v = np.asarray(k), np.asarray(v)
    tools.check_dimension(3, axis, k, v)
    k_0 = tools.absolute(k, axis=axis)
    k_4 = np.concatenate([np.expand_dims(k_0, axis=axis), k], axis=axis)
    if return_frame == 'atom':
        """ Return k in the atomic system. """
        return boost(k_4, v)
    elif return_frame == 'lab':
        """ Return k in the laboratory system. """
        return boost(k_4, -v)
    else:
        raise ValueError('rest_frame must be either "atom" or "lab".')


def gaussian_beam_3d(r: array_like, k: array_like, w0: array_like,
                     r0: array_like = None, p0: array_like = None, axis: int = -1) -> array_like:
    """
    :param r: The position 3-vector where to calculate the beam intensity (m).
    :param k: The k-wave-3-vector of light (rad / m).
    :param w0: The beam waist (m).
    :param r0: The position 3-vector of the beam waist. Is (0m, 0m, 0m) if r0 is not specified (m).
    :param p0: The total power propagated by the gaussian beam. Is 1W if p0 is not specified (W).
    :param axis: The axis along which the vector components are aligned.
    :returns: The intensity of a gaussian beam with k-wave-vector k at the position r - r0 (W/m**2 == µW/mm**2).
    :raises ValueError: r, k and r0 must have 3 components along the specified axis.
     The shapes of r, k, w0, r0 and p0 must be compatible.
    """
    if r0 is None:
        r0 = np.zeros_like(r)
    if p0 is None:
        p0 = np.ones_like(w0)
    r, r0, k = np.asarray(r, dtype=float), np.asarray(r0, dtype=float), np.asarray(k, dtype=float)
    tools.check_dimension(3, axis, r, r0, k)
    # tools.check_shape_like(np.sum(r, axis=axis), np.sum(k, axis=axis), w0, np.sum(r0, axis=axis), p0)
    k_abs = tools.absolute(k, axis=axis)
    e_r, e_theta, e_phi = tools.orthonormal(k)
    rho = np.sqrt(np.sum((r - r0) * e_theta, axis=axis) ** 2 + np.sum((r - r0) * e_phi, axis=axis) ** 2)
    z = np.sum((r - r0) * e_r, axis=axis)
    z0 = 0.5 * w0 ** 2 * k_abs
    w_z = w0 * np.sqrt(1. + (z / z0) ** 2)
    return 2. * p0 / (np.pi * w_z ** 2) * np.exp(-2. * (rho / w_z) ** 2)


def gaussian_doppler_3d(r: array_like, k: array_like, w0: array_like, v: array_like, r0=None, axis=-1) -> array_like:
    """
    :param r: The position 3-vector relative to 'r0' where to calculate the doppler-shifted wave number (m).
    :param k: The k-wave-3-vector of light (rad / m).
    :param w0: The beam waist (m).
    :param v: The velocity 3-vector (m/s).
    :param r0: The position 3-vector of the beam waist. Is (0m, 0m, 0m) if r0 is not specified (m).
    :param axis: The axis along which the vector components are aligned.
    :returns: The length of the k-wave-3-vector in the atoms rest frame (rad / m).
    :raises ValueError: r, k, v and r0 must have 3 components along the specified axis.
     The shapes of r, k, w0, v and r0 must be compatible.
    """
    if r0 is None:
        r0 = np.zeros_like(r)
    r, r0, k, v = np.asarray(r), np.asarray(r0), np.asarray(k), np.asarray(v)
    tools.check_dimension(3, axis, r, r0, k, v)
    tools.check_shape_like(np.sum(r, axis=axis), np.sum(k, axis=axis), np.array(w0),
                           np.sum(v, axis=axis), np.sum(r0, axis=axis))
    k_abs = tools.absolute(k, axis=axis)
    e_r, e_theta, e_phi = tools.orthonormal(k)
    rho = np.sqrt(np.sum((r - r0) * e_theta, axis=axis) ** 2 + np.sum((r - r0) * e_phi, axis=axis) ** 2)
    z = np.sum((r - r0) * e_r, axis=axis)
    z_0 = 0.5 * w0 ** 2 * k_abs
    z_plus = z ** 2 + z_0 ** 2
    z_minus = z ** 2 - z_0 ** 2
    alpha = tools.angle(v, k, axis=axis)
    bet_abs = beta(tools.absolute(v, axis=axis))
    return k_abs * gamma_3d(v) * (1. - bet_abs * np.cos(alpha) * (1. - w0 ** 2 / 2. / z_plus
                                                                  - rho ** 2 / 2. * z_minus / (z_plus ** 2))
                                  - bet_abs * np.sin(alpha) * rho * z / z_plus)


""" Probability distributions """


def thermal_v_pdf(v: array_like, m: array_like, t: array_like) -> array_like:
    """
    :param v: velocity quantiles (m/s).
    :param m: The mass of the ensembles bodies (amu).
    :param t: The temperature of the ensemble (K).
    :returns: The probability density in thermal equilibrium at the velocity v (s/m).
    """
    v, m, t = np.asarray(v), np.asarray(m), np.asarray(t)
    scale = np.sqrt(sc.k * t / (m * sc.atomic_mass))
    return st.norm.pdf(v, scale=scale)


def thermal_v_rvs(m: array_like, t: array_like, size: Union[int, tuple] = 1) -> array_like:
    """
    :param m: The mass of the ensembles bodies (amu).
    :param t: The temperature of the ensemble (K).
    :param size: Either the size (int) or shape (tuple) of the returned velocity array.
     If 'm' or 't' is an iterable/array, their common shape must be appended to the desired shape of the random samples.
    :returns: Random velocities according to the thermal equilibrium distribution (m/s).
    """
    m, t = np.asarray(m), np.asarray(t)
    scale = np.sqrt(sc.k * t / (m * sc.atomic_mass))
    return st.norm.rvs(scale=scale, size=size)


def thermal_e_pdf(e: array_like, t: array_like) -> array_like:
    """
    :param e: energy quantiles (eV).
    :param t: The temperature of the ensemble (K).
    :returns: The probability density at the energy e, distributed according to a boltzmann distribution (1/eV).
    """
    e, t = np.asarray(e), np.asarray(t)
    scale = sc.k * t / 2. / E_NORM
    return st.chi2.pdf(e, 1, scale=scale)


def thermal_e_rvs(t: array_like, size: Union[int, tuple] = 1) -> array_like:
    """
    :param t: The temperature of the ensemble (K).
    :param size: Either the size (int) or shape (tuple) of the returned energy array.
     If 't' is an iterable/array, its shape must be appended to the desired shape of the random samples.
    :returns: Random energies according to the boltzmann distribution (m/s).
    """
    t = np.asarray(t)
    scale = sc.k * t / 2. / E_NORM
    return st.chi2.rvs(1, scale=scale, size=size)


def convolved_boltzmann_norm_pdf(e: array_like, t: array_like, scale_e: array_like, e0: array_like = 0) -> array_like:
    """
    :param e: energy quantiles (eV).
    :param t: The temperature of the ensemble (K).
    :param scale_e: The standard deviation of the normal distribution (eV).
    :param e0: The mean energy of the normal distribution (eV).
    :returns: The probability density at the energy e, distributed according
     to a convolution of the boltzmann and a normal distribution (1/eV).
    """
    e, t, scale_e, e0 = np.asarray(e), np.asarray(t), np.asarray(scale_e), np.asarray(e0)
    t /= E_NORM
    scale = scale_e / (sc.k * t)
    loc = (e - e0) / (sc.k * t) - scale ** 2
    nonzero = loc.astype(bool)
    loc = loc[nonzero]
    norm = np.exp(-0.5 * scale ** 2) \
        / (np.sqrt(2.) * np.pi * scale * sc.k * t)
    x = (loc / (2. * scale)) ** 2
    main = np.full(e.shape, np.sqrt(LEMNISCATE * np.sqrt(np.pi) * scale))
    main_nonzero = np.empty_like(e[nonzero], dtype=float)
    mask = loc < 0.
    main_nonzero[mask] = np.sqrt(-loc[mask] / 2.) * np.exp(-loc[mask]) \
        * sp.kv(0.25, x[mask]) * np.exp(-x[mask])
    main_nonzero[~mask] = np.pi / 2. * np.sqrt(loc[~mask]) * np.exp(-loc[~mask]) \
        * (sp.ive(0.25, x[~mask]) + sp.ive(-0.25, x[~mask]))
    main[nonzero] = main_nonzero * norm
    return main


def convolved_thermal_norm_v_pdf(v: array_like, m: array_like, t: array_like,
                                 scale_e: array_like, e0: array_like = 0, relativistic=True) -> array_like:
    """
    :param v: velocity quantiles. All values must have the same sign (m/s).
    :param m: The mass of the ensembles bodies (amu).
    :param t: The temperature of the ensemble (K).
    :param scale_e: The standard deviation of the normal distribution (eV).
    :param e0: The mean energy of the normal distribution (eV).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The probability density at the velocity v, corresponding to the kinetic energy, distributed according
     to a convolution of the boltzmann and a normal distribution (s/m).
    """
    v, m, t, scale_e, e0 = np.asarray(v), np.asarray(m), np.asarray(t), np.asarray(scale_e), np.asarray(e0)
    if np.any(v < 0.) and np.any(v > 0.):
        raise ValueError('This pdf can only describe the case where all velocities have the same sign.')
    energy = e_kin(v, m, relativistic)
    tr = m * sc.atomic_mass * np.abs(v)
    if relativistic:
        tr *= gamma(v) ** 3
    return convolved_boltzmann_norm_pdf(energy, t, scale_e, e0=e0) * tr / E_NORM


def convolved_thermal_norm_f_pdf(f: array_like, f_lab: array_like, alpha: array_like, m: array_like, t: array_like,
                                 scale_e: array_like, e0: array_like = 0, relativistic=True) -> array_like:
    """
    :param f: Frequency quantiles (arb. units).
    :param f_lab: Laser frequency in the laboratory frame ([f]).
    :param alpha: Angle between the laser and the atoms velocity direction (rad).
    :param m: The mass of the ensembles bodies (amu).
    :param t: The temperature of the ensemble (K).
    :param scale_e: The standard deviation of the normal distribution (eV).
    :param e0: The mean energy of the normal distribution (eV).
    :param relativistic: The calculation is performed either relativistically or classically.
    :returns: The probability density at the frequency 'f' in the atoms rest frame,
     related to the kinetic energy via the laser frequency 'f_lab' and the Doppler effect.
     The kinetic energies are distributed according to a convolution of the boltzmann and a normal distribution (1/MHz).
    """
    f, f_lab = np.asarray(f), np.asarray(f_lab)
    m, t, scale_e, e0 = np.asarray(m), np.asarray(t), np.asarray(scale_e), np.asarray(e0)

    v = inverse_doppler(f, f_lab, alpha, mode='isnan-small')
    tr = np.abs(inverse_doppler_d1(f, f_lab, alpha, mode='isnan-small'))
    mask = np.isnan(v)
    ret = np.zeros(f.shape)
    ret[~mask] = convolved_thermal_norm_v_pdf(v[~mask], m, t, scale_e, e0=e0, relativistic=relativistic) * tr[~mask]
    return ret


def convolved_thermal_norm_f_lin_pdf(f: array_like, xi: array_like, sigma: array_like, col=True) -> array_like:
    """
    :param f: Frequency quantiles (arb. units).
    :param xi: The proportionality constant between kinetic energy differences and frequency differences ([f]).
    :param sigma: The standard deviation of the underlying normal distribution in frequency units ([f]).
    :param col: Col/Acol alignment of the laser relative to the atom beam.
    :returns: The probability density at the frequency 'f' in the atoms rest frame,
     related to differences in kinetic energy via the proportionality constant 'xi'.
     The kinetic energies are distributed according to a convolution of the boltzmann and a normal distribution (1/[f]).
    """
    pm = 1. if col else -1.
    f, xi, sigma = np.asarray(f), np.asarray(xi), np.asarray(sigma)
    scalar_true = tools.check_shape((), f, xi, sigma, return_mode=True)
    if scalar_true:
        f = np.array([f])
    sig = (0.5 * sigma / xi) ** 2
    norm = np.exp(-0.5 * sig) / (np.sqrt(2.) * np.pi * sigma)
    mu = -0.5 * pm * f / xi - sig
    b_arg = 0.25 * mu ** 2 / sig

    nonzero = mu.astype(bool)
    mu = mu[nonzero]
    b_arg = b_arg[nonzero]
    main = np.full(f.shape, np.sqrt(LEMNISCATE * np.sqrt(sig * np.pi)))
    main_nonzero = np.empty_like(f[nonzero], dtype=float)
    mask = mu < 0.

    main_nonzero[mask] = np.sqrt(-0.5 * mu[mask]) * np.exp(-mu[mask]) \
        * np.exp(-b_arg[mask]) * sp.kv(0.25, b_arg[mask])
    main_nonzero[~mask] = 0.5 * np.pi * np.sqrt(mu[~mask]) * np.exp(-mu[~mask]) \
        * (sp.ive(0.25, b_arg[~mask]) + sp.ive(-0.25, b_arg[~mask]))
    main[nonzero] = main_nonzero
    if scalar_true:
        return main[0] * norm
    return main * norm


def source_energy_pdf(f, f0, sigma, xi, collinear=True):
    """
    :param f: Frequency quantiles (arb. units).
    :param f0: Frequency offset (arb. units).
    :param sigma: The standard deviation of the underlying normal distribution in frequency units ([f]).
    :param xi: The proportionality constant between kinetic energy differences and frequency differences ([f]).
    :param collinear:
    :returns: PDF of rest frame frequencies after acceleration of thermally and normally distributed kinetic energies.
    """
    pm = 1. if collinear else -1.
    f = np.asarray(f)
    sig = (sigma / (2. * xi)) ** 2
    _norm = np.exp(-0.5 * sig) / (sigma * np.sqrt(2. * np.pi))
    mu = -pm * (f - f0) / (2. * xi) - sig
    nonzero = mu.astype(bool)
    mu = mu[nonzero]
    b_arg = mu ** 2 / (4. * sig)
    main = np.full(f.shape, np.sqrt(LEMNISCATE * np.sqrt(sig / np.pi)))
    main_nonzero = np.empty_like(f[nonzero], dtype=float)
    mask = mu < 0.

    main_nonzero[mask] = np.sqrt(-0.5 * mu[mask] / np.pi) * np.exp(-mu[mask]) \
        * np.exp(-b_arg[mask]) * sp.kv(0.25, b_arg[mask])
    main_nonzero[~mask] = 0.5 * np.sqrt(mu[~mask] * np.pi) * np.exp(-mu[~mask]) \
        * (sp.ive(0.25, b_arg[~mask]) + sp.ive(-0.25, b_arg[~mask]))
    main[nonzero] = main_nonzero
    return main * _norm
