% Author: D. Quartullo
% The code below has been readapted from the Fortran code developed by M. Migliorati in 1993 for his PhD studies on the DAFNE accelerator.
% Some references:
% Migliorati M. "Simulazioni numeriche della dinamica multipacchetto", PhD thesis, University of Rome 'La Sapienza', Rome, 1996.
% Bassetti M. et al. "A time domain simulation code of the longitudinal multibunch instabilities", DAFNE technical note G-19, INFN-LNF Accelerator Division, Frascati, 1993.
% Gallo A. et al. "Simulations of the bunch-by-bunch feedback operation with a broadband RF cavity as longitudinal kicker", DAFNE technical note G-31, INFN-LNF Accelerator Division, Frascati, 1995.
% There are four scripts in total, named Parts 1, 2, 3 and 4. Each Part is indipendent of the others, and no external libraries are needed to run it. Nevertheless, it is advisable to study the four Parts in order, since each builds on top of the previous ones.
% The characters ??? indicate code portions that the reader has to correctly fill in to be able to procede with the script. It is advisible to run the script section after section.
%% PART 3: adding the induced voltage from one HOM
%
%
% We set the accelerator parameters.

% Clear variables and figures
clc
clf
clear all
close all force

% Insert the parameters of the table
h = ???;
Cr = ???; % [m]
E0 = ???; % [eV]
Vrf = ???; % [V]
alpha0 = ???;
U0 = ???; % [eV]

% Compute periods and frequencies
c = physconst('LightSpeed'); % [m/s]
T0 = ???; % [s]
Trf = ???; % [s]
f0 = ???; % [Hz]
frf = ???; % [Hz]
omegar_rf = ???; % [rad/s]

%
% We consider four equispaced bunches in the ring. The initial phase of each bunch is the synchronous phase. The energy deviation is zero.

filled_buckets = [1,31,61,91];
n_bunches = length(filled_buckets);
phi_s = acos(???); % [rad]
init_phase = phi_s;
phase_array = init_phase*ones(1,n_bunches); 
delta_array = zeros(1,n_bunches);

%
% DAFNE usually operates with 105 contiguous bunches per beam, with average currents per beam up to 1.7 A. The HOMs of the RF accelerating cavity were measured in the 90s, and the parameters are:
% f_r [MHz]
% R/Q [\Omega]
% Undamped Q
% Waveguide-loaded Q
% 745.7
% 16
% 24000
% 70
% 796.8
% 0.5
% 40000
% 210
% 1023.6
% 0.9
% 28000
% 90
% 1121.1
% 0.3
% 12000
% 300
% 1175.9
% 0.6
% 5000
% 90
% 1201.5
% 0.2
% 9000
% 180
% 1369.0
% 2.0
% 5000
% 170
% 1431.7
% 1.0
% 4000
% 550
% In the 90s, the accelerating cavity was loaded with waveguides to significantly damp the quality factors of the HOMs (column 4).
% In the following, we consider the undamped mode at 796.8 MHz. In fact, we change (slightly) its resonant frequency so to excite a particular mode of coupled-bunch oscillations. The critical frequencies leading to coupled-bunch instabilities are given by
% f_{l,\mu } =\left(\mu -N_b l\right)f_0 +f_s
% where l is a negative integer,\mu is the coupled-bunch mode and N_b is the number of bunches. In our example we set l=-66 and \mu =1.
% In principle, f_s is the small-amplitude synchrotron frequency including collective effects, but here, as an approximation, we consider f_s without collective effects, which us known analytically (see Part 2).

% We define the resonator parameters
Q_factor = 40000;
Rs = 20000; % [Ohm]
mu = 1;
l = -66;
alpha_damp_AN = ???;
fs_AN = sqrt(???*Vrf*alpha0*sin(phi_s)/(2*pi*E0*T0*Trf)-???); 
fr = (???-n_bunches*l)*f0+fs_AN; % [Hz]
% We define some derived parameters
omega_r = ???;
Gamma = omega_r/(2*Q_factor);
omega_n = sqrt(omega_r^2-Gamma^2);
Ga_omn = Gamma/omega_n;
omn_omr = omega_n/omega_r;
A_cost = omn_omr*Q_factor/Rs;
B_cost = omn_omr*Rs/Q_factor;

%
% Including the collective effects from the HOMs, the longitudinal equations of motion for the bunch k are
% \begin{array}{l}
% {\Delta \varphi }_k^{\left(n+1\right)} ={\Delta \varphi }_k^{\left(n\right)} +2\pi h\alpha_0 \delta_k^{\left(n\right)} \\
% \delta_k^{\left(n+1\right)} =\delta_k^{\left(n\right)} +\frac{eV_{\textrm{rf}} \;\cos {\Delta \varphi }_k^{\left(n+1\right)} }{E_0 }-\frac{U_0 }{E_0 }\left(1+2\delta_k^{\left(n\right)} \right)+\frac{e}{E_0 }\sum_{j=1}^{N_{\textrm{HOM}} } \left\lbrack V_{k,j,\textrm{res}}^{\left(n\right)} +\frac{\Delta V_{k,j} }{2}\right\rbrack
% \end{array}
% where V_{k,j,\textrm{res}}^{\left(n\right)} is the residual induced voltage in the HOMj which the bunch k sees when it traverses the RF cavity at turn n. \Delta V_{k,j} =-R_{s,j} {\;\omega }_{r,j} \;Q_{b,k} \;/\;Q_j is the voltage induced by the bunch k on the HOM j. \Delta V_{k,j} \;/\;2 is the beam-loading term, i.e. the bunch sees half of \Delta V_{k,j} as soon as it goes through the cavity.
% After defining the bunch current, we can directly compute \Delta V_{k,j}.

current_per_bunch = 15e-3; % [A]
charge_per_bunch = current_per_bunch*T0; % [C]
delta_V = ???; % [V]

%
% As usual, let's track the four bunches and save their coordinates in a matrix.
% We assume that the HOM is initially unloaded, i.e.V_{\textrm{res}} =I_{\textrm{res}} =0 at t=0. As in Part 2, at a given turn n each bunch k performs a drift along the ring, it loses energy due to synchrotron radiation and then it receves an energy kick from the accelerating cavity. Now the difference is that the HOM gives another energy kick to the bunch, equal to
% V_{k,\textrm{res}}^{\left(n\right)} +\frac{\Delta V_k }{2}
% The bunch k contributes with \Delta V_k to the HOM residual voltage, while it doesn't contribute to the HOM residual current. After having computed the time distance ('time') between the bunches k and k+1, we can evaluate the HOM residual voltage and current seen by the bunch k+1:
% \left\lbrack \begin{array}{c}
% V_{k+1,\textrm{res}}^{\left(n\right)} \\
% i_{k+1,\textrm{res}}^{\left(n\right)}
% \end{array}\right\rbrack =e^{-\frac{\omega_r }{2Q}\textrm{time}} \left\lbrack \begin{array}{cc}
% \cos \left(\omega_n \;\textrm{time}\right)-\frac{\omega_r }{2Q\omega_n }\sin \left(\omega_n \;\textrm{time}\right) & -\frac{R_s \omega_r }{Q\omega_n }\sin \left(\omega_n \;\textrm{time}\right)\\
% \frac{\omega_r Q}{R_s \omega_n }\sin \left(\omega_n \;\textrm{time}\right) & \cos \left(\omega_n \;\textrm{time}\right)+\frac{\omega_r }{2Q\omega_n }\sin \left(\omega_n \;\textrm{time}\right)
% \end{array}\right\rbrack \left\lbrack \begin{array}{c}
% V_{k,\textrm{res}}^{\left(n\right)} +\Delta V_k \\
% i_{k,\textrm{res}}^{\left(n\right)}
% \end{array}\right\rbrack

n_turns = 100000; 
phase_matrix = zeros(n_turns, n_bunches);
delta_matrix = zeros(n_turns, n_bunches);
V_out = 0;
i_out = 0;
% Tracking
for i = 1:n_turns
    % kick and drift equations of motion (as in Part 2)
    phase_array = phase_array + 2*pi*h*alpha0*delta_array;
    delta_array = delta_array + Vrf*cos(phase_array)/E0 - ???;
    
    % induced voltage contribution, for each bunch k
    for k = 1:n_bunches
        %%% VOLTAGE FROM THE HOM TO THE BUNCH K 
        % the voltage seen by the bunch k is the residual voltage
        % in the HOM plus the beam-loading voltage
        vind = V_out + delta_V/???;
        delta_array(k) = delta_array(k) + vind/E0;
        %%% VOLTAGE FROM THE BUNCH K TO THE HOM
        % the bunch k contributes with delta_V
        % to the voltage of the HOM
        V_inp = V_out + delta_V;
        i_inp = i_out;
        %%% COMPUTATION OF THE RESIDUAL VOLTAGE FOR THE BUNCH K+1
        % time distance between bunch k+1 and k
        if k == n_bunches 
            new_phase_bunch1 = phase_array(1) + 2*pi*h*alpha0*delta_array(1);
            time = (new_phase_bunch1-phase_array(n_bunches))/omegar_rf + ...
                    Trf*(h-filled_buckets(n_bunches)+filled_buckets(1));
        else
            time = (phase_array(k+1)-phase_array(k))/omegar_rf + ...
                    Trf*(filled_buckets(k+1)-filled_buckets(k));
        end
        % computation of the residual voltage for bunch k+1
        omn_time = omega_n*time;
        cos_omn_time = cos(omn_time);
        sin_omn_time = sin(omn_time);
        a11 = cos_omn_time - Ga_omn*sin_omn_time;
        a12 = -sin_omn_time/A_cost;
        a21 = sin_omn_time/B_cost;
        a22 = cos_omn_time + Ga_omn*sin_omn_time;
        cost = exp(-Gamma*time);
        V_out = cost*(V_inp*a11+i_inp*a12);
        i_out = cost*(V_inp*a21+i_inp*a22);
    end
    phase_matrix(i,:) = phase_array;
    delta_matrix(i,:) = delta_array;
    if mod(i,10000)==0
        fprintf('turn %f\n', i);
    end
end

%
% Let's now plot the phase evolutions of the four bunches. Note that consecutive bunches are misaligned in phase by 90°. Indeed, we know from theory that the phase misalignment is
% \Delta \varphi_{\mu } =\frac{2\pi }{N_b }\mu, with N_b =4 and \mu =1 in our case.

fig1 = figure(1);
plot(phase_matrix(:,1)*180/pi, color = 'b');
hold on
plot(phase_matrix(:,2)*180/pi, color = 'g');
plot(phase_matrix(:,3)*180/pi, color = 'r');
plot(phase_matrix(:,4)*180/pi, color = 'y');
xlabel('Turn');
ylabel('Phase [deg]');
grid on
hold off

%
% Now we plot the four bunches in phase space. Only for easiness of visualization, the bunches are plotted in the same RF bucket between -90° and 270°. Verify that the phase misalignment between consecutive bunches is indeed 90°.

fig2 = figure(2);
plot_every = 5;
plot_up_to_turn = 5000;
PLOT = plot(nan,nan, nan,nan, nan,nan, nan,nan);
bunch_phases_1 = phase_matrix(:,1)*180/pi;
bunch_deltas_1 = delta_matrix(:,1);
bunch_phases_2 = phase_matrix(:,2)*180/pi;
bunch_deltas_2 = delta_matrix(:,2);
bunch_phases_3 = phase_matrix(:,3)*180/pi;
bunch_deltas_3 = delta_matrix(:,3);
bunch_phases_4 = phase_matrix(:,4)*180/pi;
bunch_deltas_4 = delta_matrix(:,4);
axis([min(bunch_phases_1(1:plot_up_to_turn)) ...
      max(bunch_phases_1(1:plot_up_to_turn)) ...
      min(bunch_deltas_1(1:plot_up_to_turn)) ... 
      max(bunch_deltas_1(1:plot_up_to_turn))]);
xlabel('Phase [deg]');
ylabel('delta [1]');
grid
for ii = 1:plot_every:plot_up_to_turn
    pause(0.04)
    set(PLOT(1), 'XData', bunch_phases_1(ii), ...
              'YData', bunch_deltas_1(ii), ...
              marker ='.', color = 'b', MarkerSize = 30);
    set(PLOT(2), 'XData', bunch_phases_2(ii), ...
              'YData', bunch_deltas_2(ii), ...
              marker ='.', color = 'g', MarkerSize = 30);
    set(PLOT(3), 'XData', bunch_phases_3(ii), ...
              'YData', bunch_deltas_3(ii), ...
              marker ='.', color = 'r', MarkerSize = 30);
    set(PLOT(4), 'XData', bunch_phases_4(ii), ...
              'YData', bunch_deltas_4(ii), ...
              marker ='.', color = 'y', MarkerSize = 30);
    title(strcat('Turn =  ', num2str(ii)))
end

%
% We consider the phase oscillations of bunch 1 and we perform an exponential fit of the envelope. From this we obtain the numerical grow rate.

% compute the maxima
time_array = T0: T0: n_turns*T0;
fit_from_to_turn = (300:7000);
time_array_fit = time_array(fit_from_to_turn);
bunch_phases_fit = bunch_phases_1(fit_from_to_turn);
[pks_phase,locs_phase] = findpeaks(bunch_phases_fit);
% isolate the maxima
time_peaks = time_array_fit(locs_phase);
phase_peaks = bunch_phases_fit(locs_phase);
% compute the initial guesses
c0 = mean(???);
b0 = log((phase_peaks(1)-c0)/(phase_peaks(2)-c0))/(time_peaks(1)-time_peaks(2));
a0 = (phase_peaks(1)-c0)*exp(-b0*time_peaks(1));
% compute the fit and print and the numerical growth rate
g = fittype('a*exp(b*x)+c','indep','x');
fit_object = fit(time_peaks', phase_peaks, ...
                 g, 'StartPoint',[a0 b0 c0]);
coefficients = coeffvalues(fit_object);
fprintf('alpha_grow_FIT = %f\n', coefficients(2))
% plot of the fit
fig3 = figure(3);
plot(bunch_phases_1, color = 'b')
hold on
plot(fit_from_to_turn,bunch_phases_fit, color = 'r')
phase_fit = coefficients(1)*exp(coefficients(2)*time_array_fit)+coefficients(3);
plot(fit_from_to_turn,phase_fit, 'g', 'LineWidth',3)
xlabel('Turn');
ylabel('Phase [deg]');
grid
hold off
fig4 = figure(4);
plot(fit_from_to_turn,bunch_phases_fit, color = 'r')
hold on
plot(fit_from_to_turn,phase_fit, 'g', 'LineWidth',3)
xlabel('Turn');
ylabel('Phase [deg]');
grid
hold off
title('Zoom')

%
% To derive the growth rate from the analytical formula, first we define the longitudinal impedance of a resonator.
% Z\left(f\right)=\frac{R_s }{1+\textrm{jQ}\left(\frac{f}{f_r }-\frac{f_r }{f}\right)}

impedance = @(Q_factor, Rs, fr, freq) Rs*fr*freq./(fr*freq+1i*Q_factor*(???));

%
% We can plot for instance the real part of the impedance, which is an even function of frequency.

frequency_array = 0:3e3:1e9; % [Hz]
impedance_array = impedance(Q_factor, Rs, fr, frequency_array); % [Ohm]
fig5 = figure(5);
plot(frequency_array/1e6, real(impedance_array)/1e3)
xlabel('Frequency [MHz]');
ylabel('Re Z [kOhm]');
grid

%
% We can evaluate the analytical growth rate with the following expression
% \alpha_{r,\textrm{SR},\textrm{HOM}} =-\frac{U_0 }{E_0 T_0 }-\frac{c\alpha_0 eN_b Q_b }{2C_r E_0 T_0 f_s }\sum_{l=-\infty }^{+\infty } \left\lbrack \left({l\;N}_b -\mu \right)f_0 -f_s \right\rbrack \;\textrm{Re}\left\lbrace Z\left\lbrack \left({l\;N}_b -\mu \right)f_0 -f_s \right\rbrack \right\rbrace
% Note that f_s in the formula should be computed including collective effects, but this doesn't usually impact significantly the evaluation of \alpha_{r,\textrm{SR},\textrm{HOM}}.

sum_fZexp = 0;
lmax = 1000;
for l = -lmax : lmax+1
    freq = (l*n_bunches-mu)*f0-fs_AN;
    sum_fZexp = sum_fZexp + freq*???;                    
end
growth_rate = ??? -c*alpha0*???*n_bunches*charge_per_bunch*sum_fZexp/(2*Cr*E0*T0*fs_AN);
fprintf('alpha_grow_AN = %f\n', growth_rate)

% Verify that the analytical and numerical estimations of the grow rate are close to each other.