import numpy
import matplotlib.pyplot as plt
import matplotlib.cm as cmap
import at
from at.collective import BeamLoadingElement, add_beamloading
from at.collective import FeedbackMode as FBMode
from mpi4py import MPI
from at.tracking.utils import get_bunches_std_mean
from at.constants import qe
from at.physics import harmonic_analysis


def analytical_qs(ring, I, Qfactor, Rshunt, beta):

    E0 = ring.energy * qe
    alpha_c = at.get_mcf(ring.disable_6d(copy=True))
    omega_rf = 2 * numpy.pi * ring.get_rf_frequency()
    Vc = ring.get_rf_voltage()

    envel = at.envelope_parameters(ring.enable_6d(copy=True))
    U0 = envel.U0 * qe

    Q0 = Qfactor * (1 + beta)
    Rsh = Rshunt * (1 + beta)  # Difference in definition

    phi_s = numpy.pi - numpy.arcsin(U0 / (qe * Vc))  # sync. phase in radian
    psi = numpy.arctan(
        -2 * Rsh * I * numpy.cos(numpy.pi - phi_s) / (Vc * (1 + beta))
    )  # tuning angle

    omega_res = omega_rf / (
        1 - 2 * Rsh * I * (numpy.cos(numpy.pi - phi_s)) / (2 * Q0 * Vc)
    )  # resonant freq.

    h = ring.harmonic_number  # harmonic number
    omega0 = omega_rf / h  # rev. freq.
    T0 = 2 * numpy.pi / omega0  # rev. time

    omega_s0 = numpy.sqrt(
        qe * Vc * omega_rf * alpha_c * numpy.cos(numpy.pi - phi_s) / (E0 * T0)
    )  # Synch. freq. for a single particle

    K = alpha_c * qe * I / (E0 * T0) * Rsh / (1 + beta)

    # Synch. freq. with beam loading
    tilde_omega_s = numpy.sqrt(omega_s0**2 + K * omega_rf * numpy.sin(2 * psi) + 1j * 0)
    return tilde_omega_s / omega0, omega_s0 / omega0


comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

print(size, rank)

ring = at.load_lattice("../../../machine_data/esrf.m")

ring.enable_6d()
ring.set_rf_voltage(8.0e6)

Nbunches = 32

current = 150e-3


# In fact, the following resonator parameters are for the EBS machine
# Not the old ESRF SR, however for benchmarking this doesn't matter
beta = 2.8
q0 = 37500  # unloaded
qfactor = q0 / (1 + beta)  # loaded
rshunt = 145 * qfactor * 13  # loaded


tunes = ring.get_tune()
envel = ring.envelope_parameters()
mcf = at.get_mcf(ring.disable_6d(copy=True))
tauz = envel.Tau[2] * ring.revolution_frequency

simple_ring = at.simple_ring(
    ring.energy,
    ring.circumference,
    ring.harmonic_number,
    tunes[0],
    tunes[1],
    ring.get_rf_voltage(),
    mcf,
    U0=envel.U0,
    tauz=tauz,
    espread=envel.sigma_e,
)

simple_ring.set_beam_current(current)
simple_ring.set_fillpattern(Nbunches)
simple_ring.pop(2)  # drop diffusion element

simple_ring.set_cavity_phase()


# Npart must be at least Nbunches per core
Npart = Nbunches

# Now we give the fring and convert it
# into a beam loaded cavity.

# Be careful with the gains. A voltgain of 1.0
# might be stable with FBMode.ONETURN, but you
# will lose the reduction of the coherent tune
# due to the feedback working too quickly.

add_beamloading(
    simple_ring,
    qfactor,
    rshunt,
    Nslice=1,
    VoltGain=0.001,
    PhaseGain=0.001,
    buffersize=1000,
    windowlength=1000,
    FBMode=FBMode.WINDOW,
)


# Specify some simulation parameters
kickTurn = 0
Nturns = 2**14 + kickTurn

part = numpy.zeros((6, Npart))

if rank == 0:

    print("\n")
    print("Current:", current)

    z_all = numpy.zeros((Nturns, Nbunches))
    dp_all = numpy.zeros((Nturns, Nbunches))


# Here it should be considered that there are 2 ways to run
# simulations in pyat. Either with ring.track(part, nturns=Nturns)
# or with for i in np.arange(Nturns): ring.track(part, nturns=1).
# With the former, you should use the BeamMoments element to acquire
# the means and stds of each bunch turn by turn when using MPI.
# However, when you want to kick after a certain turn, you have 1
# of 2 choices. Either you use a BeamMoments element, and split the
# tracking into 2 (before and after). You can then concatenate the
# the results. Or you can use the code below to gather all particles
# yourself with the MPICOMM after each turn.

for i in numpy.arange(Nturns):
    # Apply a kick to ensure the coherent tune is large
    if i == kickTurn:
        part[4, :] += 1e-3

    simple_ring.track(part, nturns=1, refpts=None, in_place=True)

    # Gather particles over all cores (compatible with MPI on or off)
    allPartsg = comm.gather(part)
    if rank == 0:
        allPartsg = numpy.concatenate(allPartsg, axis=1)
        stds, means = get_bunches_std_mean(allPartsg, Nbunches)

        dp_all[i] = numpy.array([x[4] for x in means])
        z_all[i] = numpy.array([x[5] for x in means])
        if i % 1000 == 0:
            # Print the turn number and number of lost particles
            print(i, numpy.sum(numpy.isnan(allPartsg)))


# Now we compute the tune of each bunch, and compare with the analytical
if rank == 0:
    qscoh = numpy.zeros(Nbunches)
    for ib in numpy.arange(Nbunches):
        dp_dat = dp_all[kickTurn:, ib] - numpy.mean(dp_all[kickTurn:, ib])
        qs = harmonic_analysis.get_tunes_harmonic(
            dp_dat, num_harmonics=20, fmin=1e-5, fmax=0.1
        )
        qscoh[ib] = qs[0]

    qs_mn, qs_std = numpy.array([numpy.mean(qscoh), numpy.std(qscoh)])

    qs_theory, qs_zerocurrent = analytical_qs(ring, current, qfactor, rshunt, beta)

    print("Analytical:", numpy.real(qs_theory))
    print("Simulated:", qs_mn, "pm", qs_std)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 8))

    freq = numpy.fft.rfftfreq(Nturns - kickTurn)
    for i in numpy.arange(Nbunches):
        fftdat = numpy.abs(numpy.fft.rfft(dp_all[kickTurn:, i]))
        if i == 0:
            ax2.semilogy(
                freq,
                fftdat,
                color=cmap.jet(float(i + 1) / Nbunches),
                label="Bunch by Bunch FFT",
            )
        else:
            ax2.semilogy(freq, fftdat, color=cmap.jet(float(i + 1) / Nbunches))

    ax1.plot(1e3 * numpy.mean(z_all, axis=1), color="r")
    ax1.set_xlabel("Turn")
    ax1.set_ylabel("<z> [mm]")
    ax1.set_xlim(0, 1000)

    ax2.set_xlim([0, 8e-3])
    ax2.set_xlabel("Tune")
    ax2.set_ylabel("FFT Amp [A.U.]")
    ax2.axvline(numpy.mean(qscoh), label="Coherent Tune", linestyle="dashed", color="k")
    ax2.axvline(
        numpy.real(qs_theory),
        label="Analytical Beam Loaded Tune",
        linestyle="dashed",
        color="r",
    )
    ax2.axvline(
        qs_zerocurrent, label="Zero Current Tune", linestyle="dashed", color="b"
    )
    ax2.legend()
    plt.suptitle("VGain=0.001, PGain=0.001")
    plt.show()
