import math


def get_photon_density_gaussian(
    elec_x, elec_y, elec_z, ct, photon_n_lab_max, inv_laser_waist2,
    inv_laser_ctau2, laser_initial_z0, gamma_boost, beta_boost ):
    """
    Get the photon density in the scattering Gaussian laser pulse,
    at the position of a given electron, and at the current time.

    Parameters
    ----------
    elec_x, elec_y, elec_z: floats
        The position of the given electron (in the frame of the simulation)
    ct: float
        Current time in the simulation frame (multiplied by c)
    photon_n_lab_max: float
        Peak photon density (in the lab frame)
        (i.e. at the peak of the Gaussian pulse)
    inv_laser_waist2, inv_laser_ctau2, laser_initial_z0: floats
        Properties of the Gaussian laser pulse (in the lab frame)
    gamma_boost, beta_boost: floats
        Properties of the Lorentz boost between the lab and simulation frame.

    Returns
    -------
    photon_n_sim: float
        The photon density in the frame of the simulation
    """
    # Transform electrons coordinates from simulation frame to lab frame
    elec_zlab = gamma_boost*( elec_z + beta_boost*ct )
    elec_ctlab = gamma_boost*( ct + beta_boost*elec_z )

    # Get photon density *in the lab frame*
    photon_n_lab = photon_n_lab_max * math.exp(
        - 2*inv_laser_waist2*( elec_x**2 + elec_y**2 ) \
        - 2*inv_laser_ctau2*(elec_zlab - laser_initial_z0 + elec_ctlab)**2 )

    # Get photon density *in the simulation frame*
    photon_n_sim = gamma_boost*photon_n_lab*( 1 + beta_boost)

    return( photon_n_sim )