from copy import deepcopy

import numpy as np
import scipy.constants as ct

from wake_t import GaussianPulse, PlasmaStage, Beamline
from wake_t.utilities.bunch_generation import get_matched_bunch


def test_variable_parabolic_coefficient():
    """
    Checks that a z-dependent parabolic coefficient works as expected.
    A piecewise function for the coefficient is defined, which should be
    equivalent to tracking multiple plasma stages each having a different
    parabolic coefficient.
    """
    # Plasma density.
    n_p = 1e23

    # Grid and field parameters.
    xi_max = 0
    xi_min = -100e-6
    r_max = 100e-6
    Nxi = 100
    Nr = 100
    dz_fields = 1e-3

    # Laser parameters.
    a0 = 1
    w0 = 30e-6
    tau = 25e-15
    l0 = 0.8e-6

    # Guiding plasma channel
    r_e = ct.e**2 / (4.0 * np.pi * ct.epsilon_0 * ct.m_e * ct.c**2)
    rel_delta_n_over_w2 = 1.0 / (np.pi * r_e * w0**4 * n_p)

    # Length and parabolic coefficient of each section (stretch).
    L_stretches = [1e-2, 1e-2, 1e-2, 1e-2, 1e-2]
    pc_stretches = rel_delta_n_over_w2 * np.array([1, 2, 1, 0.5, 2])

    # Define z-dependent parabolic coefficient.
    def parabolic_coefficient(z):
        z_stretch_end = np.cumsum(L_stretches)
        i = np.sum(np.float32(z_stretch_end) <= np.float32(z))
        if i == len(L_stretches):
            i -= 1
        return pc_stretches[i]

    # Create identical laser pulses for each case.
    laser = GaussianPulse(-50e-6, a0, w0, tau, z_foc=0.0, l_0=l0)
    laser_1 = deepcopy(laser)
    laser_2 = deepcopy(laser)

    # Create identical bunches for each case.
    bunch = get_matched_bunch(
        1e-6, 1e-6, 200, 1, 3, laser.xi_c - 30e-6, 1e-6, 1e4, n_p=n_p
    )
    bunch_1 = bunch.copy()
    bunch_2 = bunch.copy()

    # Create single plasma stage (containing all sections).
    plasma_single = PlasmaStage(
        np.sum(L_stretches),
        n_p,
        wakefield_model="quasistatic_2d",
        n_out=10,
        xi_min=xi_min,
        xi_max=xi_max,
        r_max=r_max,
        r_max_plasma=r_max,
        laser=laser_1,
        laser_evolution=True,
        n_r=Nr,
        n_xi=Nxi,
        ppc=2,
        dz_fields=dz_fields,
        parabolic_coefficient=parabolic_coefficient,
    )

    # Track single plasma.
    plasma_single.track(bunch_1)

    # Create set of plasma stages, one per section.
    sub_stages = []
    for i, (length, pc) in enumerate(zip(L_stretches, pc_stretches)):
        stage = PlasmaStage(
            length,
            n_p,
            wakefield_model="quasistatic_2d",
            n_out=3,
            xi_min=xi_min,
            xi_max=xi_max,
            r_max=r_max,
            r_max_plasma=r_max,
            laser=laser_2,
            laser_evolution=True,
            n_r=Nr,
            n_xi=Nxi,
            ppc=2,
            dz_fields=dz_fields,
            parabolic_coefficient=pc,
        )
        sub_stages.append(stage)
    plasma_multi = Beamline(sub_stages)

    # Track through set of stages.
    plasma_multi.track(bunch_2)

    # Get final envelope of both lasers.
    a_env_1 = laser_1.get_envelope()
    a_mod_1 = np.abs(a_env_1)
    a_phase_1 = np.angle(a_env_1)
    a_env_2 = laser_2.get_envelope()
    a_mod_2 = np.abs(a_env_2)
    a_phase_2 = np.angle(a_env_2)

    # Check that both envelopes are equal.
    np.testing.assert_almost_equal(a_mod_2, a_mod_1)
    np.testing.assert_almost_equal(a_phase_2, a_phase_1, decimal=5)

    # Check final spot size value
    dr = r_max / Nr
    w0_final = calculate_spot_size(a_env_1, dr)
    np.testing.assert_almost_equal(w0_final, 2.03471741366461e-05)


def calculate_spot_size(a_env, dr):
    # Project envelope to r
    a_proj = np.sum(np.abs(a_env), axis=0)

    # Maximum is on axis
    a_max = a_proj[0]

    # Get first index of value below a_max / e
    i_first = np.where(a_proj <= a_max / np.e)[0][0]

    # Do linear interpolation to get more accurate value of w.
    # We build a line y = a + b*x, where:
    #     b = (y_2 - y_1) / (x_2 - x_1)
    #     a = y_1 - b*x_1
    #
    #     y_1 is the value of a_proj at i_first - 1
    #     y_2 is the value of a_proj at i_first
    #     x_1 and x_2 are the radial positions of y_1 and y_2
    #
    # We can then determine the spot size by interpolating between y_1 and y_2,
    # that is, do x = (y - a) / b, where y = a_max/e
    y_1 = a_proj[i_first - 1]
    y_2 = a_proj[i_first]
    x_1 = (i_first - 1) * dr + dr / 2
    x_2 = i_first * dr + dr / 2
    b = (y_2 - y_1) / (x_2 - x_1)
    a = y_1 - b * x_1
    w = (a_max / np.e - a) / b
    return w


if __name__ == "__main__":
    test_variable_parabolic_coefficient()
