import numpy as np
from scipy import constants as cst

import xobjects as xo
import xtrack as xt
import xwakes as xw

def test_fakemultibunch():

    context = xo.ContextCpu()

    energy = 7E3
    qx = 62.31
    qy = 60.32
    basic_slot_time = 25E-9
    n_bunch_max = 3000
    n_bunch = 100
    bunch_intensity = 5E12
    emit_norm = 2.5E-6
    sigma_z = 0.08
    selected_bunch = n_bunch-1
    r_over_q = 1806
    q = 5E3
    coupled_bunch_number_x = 16
    coupled_bunch_number_y = 73
    n_turns_wake = 1
    num_slices = 11

    assert n_bunch_max%n_bunch == 0
    circumference = n_bunch_max * basic_slot_time * cst.c
    bunch_spacing_zeta = circumference/n_bunch
    beta_x = circumference/(2.0*np.pi*qx)
    beta_y = circumference/(2.0*np.pi*qy)
    mass0 = cst.value('proton mass energy equivalent in MeV')*1E-3
    gamma = energy/mass0
    RF_freq = 10.0/basic_slot_time
    sigma_x = np.sqrt(emit_norm*beta_x/gamma)
    sigma_y = np.sqrt(emit_norm*beta_y/gamma)
    betar = np.sqrt(1 - 1 / gamma ** 2)
    zeta_range=(-3*sigma_z, 3*sigma_z)
    filling_scheme = np.ones(n_bunch,dtype=int)

    n_part = num_slices
    slice_intensity = bunch_intensity/num_slices
    x0 = np.ones(n_part)
    z_edges = np.linspace(zeta_range[0],zeta_range[1],num_slices+1)
    z0 = z_edges[:-1]+(z_edges[1]-z_edges[0])/2
    zeta0 = z0-1*(n_bunch-1)*bunch_spacing_zeta
    particles = xt.Particles(
            _context = context, 
            q0       = 1,
            p0c      = energy*1E9,
            mass0    = mass0*1E9,
            x        = sigma_x*x0,
            y        = sigma_y*x0,
            zeta     = zeta0,
            weight   = bunch_intensity / n_part
            )

    wfx = xw.wit.ComponentResonator(kind = 'dipole_x',
                                    r=r_over_q*q, q=q, f_r=RF_freq+3E5)
    wfy = xw.wit.ComponentResonator(kind = 'dipole_y',
                                    r=r_over_q*q, q=q, f_r=RF_freq-3E5)

                                        
    coupled_bunch_phase_x = 2*np.pi*coupled_bunch_number_x/n_bunch
    wfx.configure_for_tracking(zeta_range=zeta_range,
                        num_slices=num_slices,
                        bunch_spacing_zeta=bunch_spacing_zeta,
                        num_turns=n_turns_wake,
                        circumference=circumference,
                        filling_scheme = np.ones(n_bunch,dtype=int),
                        bunch_selection = [selected_bunch],
                        fake_coupled_bunch_phase_x = coupled_bunch_phase_x,
                        beta_x = beta_x,
                        )
    coupled_bunch_phase_y = 2*np.pi*coupled_bunch_number_y/n_bunch
    wfy.configure_for_tracking(zeta_range=zeta_range,
                        num_slices=num_slices,
                        bunch_spacing_zeta=bunch_spacing_zeta,
                        num_turns=n_turns_wake,
                        circumference=circumference,
                        filling_scheme = np.ones(n_bunch,dtype=int),
                        bunch_selection = [selected_bunch],
                        fake_coupled_bunch_phase_y = coupled_bunch_phase_y,
                        beta_y = beta_y,
                        )


    px0 = np.copy(particles.px)
    wfx.track(particles)
    kicks_x_from_track = particles.px-px0
    py0 = np.copy(particles.py)
    wfy.track(particles)
    kicks_y_from_track = particles.py-py0

    moments_data = wfx._xfields_wf.moments_data
    z,x = moments_data.get_moment_profile('x',0)
    moments_data = wfy._xfields_wf.moments_data
    z,y = moments_data.get_moment_profile('y',0)
    zetas = np.array([])
    positions_x = np.array([])
    positions_y = np.array([])
    for slot in np.arange(n_bunch):
        zetas = np.hstack([zetas,z0-bunch_spacing_zeta*slot])
        positions_x = np.hstack([positions_x,x0*np.cos(coupled_bunch_phase_x*(selected_bunch-slot))])
        positions_y = np.hstack([positions_y,x0*np.cos(coupled_bunch_phase_y*(selected_bunch-slot))])
    indices = np.argsort(zetas)
    zetas = zetas[indices]
    positions_x = positions_x[indices]
    positions_y = positions_y[indices]

    assert np.allclose(z/sigma_z,zetas/sigma_z)
    assert np.allclose(x/sigma_x,positions_x)
    assert np.allclose(y/sigma_y,positions_y)

    scaling_constant = particles.q0**2 * cst.e**2 / (particles.p0c[0] * particles.beta0[0] * cst.e)
    for i_slice in range(num_slices):
        zetas_slice = zeta0[i_slice]-zetas
        kicks = positions_x*sigma_x*scaling_constant*slice_intensity*wfx.function_vs_zeta(zetas_slice,beta0=betar,dzeta=moments_data.dz)
        assert np.isclose(np.sum(kicks),kicks_x_from_track[i_slice],atol=1E-15)
        kicks = positions_y*sigma_y*scaling_constant*slice_intensity*wfy.function_vs_zeta(zetas_slice,beta0=betar,dzeta=moments_data.dz)
        assert np.isclose(np.sum(kicks),kicks_y_from_track[i_slice],atol=1E-15)
