%% Script to fit the envelope of phase oscillations with an exponential function.
%% Author: D. Quartullo

clear all 
close all
clc
hold on

%%% inputs

function_type = 'measurements'; % analytical, measurements
if strcmp(function_type, 'analytical')
    % timings
    t_start_grow = 3e-3;
    t_end_grow = 7e-3;
    t_start_damp = 7e-3;
    t_end_damp = 10e-3;
    % function parameters
    a_grow = 5.05;
    a_damp = 60.15;
    b_grow = 1.2*1e3;
    b_damp = -0.5*1e3;
    c_grow_damp = 101.1;
    omega1 = 1e6;
    omega2 = 1e6;
    n_points = 100000;
    % function construction
    t_load = linspace(0,t_end_damp,n_points);
    x_load = zeros(1,n_points);
    x_load(t_load>=t_start_grow) = a_grow*exp(b_grow*(t_load(t_load>=t_start_grow)-t_start_grow)).*sin(omega1*(t(t>=t_start_grow)-t_start_grow));
    x_load(t_load>=t_start_damp) = a_damp*exp(b_damp*(t_load(t_load>=t_start_damp)-t_start_damp)).*sin(omega2*(t(t>=t_start_damp)-t_start_damp));
    x_load = x_load+c_grow_damp;
    a_grow_eff = a_grow*exp(-b_grow*t_start_grow);
    a_damp_eff = a_damp*exp(-b_damp*t_start_damp);
elseif strcmp(function_type, 'measurements')
    just_plot_data = false;
    % timings
    t_start_grow = 5e-3;
    t_end_grow = 7e-3;
    t_start_damp = 10.5e-3;
    t_end_damp = 10.9e-3;
    % load row data
    bunch_number = 1;
    filename_measurements = './meas_growDamp_dataRaw.mat' ;
    S = load([filename_measurements]);
    DownS = str2num(S.dimtelData{5,2});
    x_load = S.dimtelData{5,10}(:,bunch_number)';
    % create time array
    frf = 368.667*1e6;
    frev = frf/120;
    Trev = 1/frev;
    tSample_step = Trev*DownS;
    tSample_max =  tSample_step*size(x_load,2) ;
    t_load = tSample_step : tSample_step : tSample_max ;
    % plot
    if just_plot_data
        plot(t_load*1e3,x_load,'color','blue')
        grid on
        xline(3,'color','red')
        error('stop')
    end
end

plot_envelope = true;
apply_bounds = true;

%%% computations

% compute grow rate
[f0, coefficients, t_analysis, x_analysis, t_env, x_env_up] = compute_rate(t_start_grow, t_end_grow, t_load, x_load, apply_bounds);
if strcmp(function_type, 'analytical')
    fprintf('a = %f, b = %f, c = %f\n', a_grow_eff, b_grow, c_grow_damp)
    error_rate = abs(coefficients(2)-b_grow)*100/b_grow;
    fprintf('error grow rate [%%] = %f\n', error_rate)
end
plot(t_load*1e3,x_load,'color','black')
plot(t_analysis*1e3,x_analysis,'color','blue')
grid
if plot_envelope
    plot(t_env*1e3,x_env_up,'color','red','LineWidth',3)
end
fit_on_t_env = coefficients(1)*exp(coefficients(2)*t_env)+coefficients(3);
plot(t_env*1e3,fit_on_t_env,'color','green','LineWidth',3);
yline(mean(x_load));
fprintf('Difference std = %f\n\n', std(fit_on_t_env-x_env_up)) 

if t_start_damp ~= -1
    % compute damp rate
    [f0_d, coefficients_d, t_analysis_d, x_analysis_d, t_env_d, x_env_up_d] = compute_rate(t_start_damp, t_end_damp, t_load, x_load, apply_bounds);
    if strcmp(function_type, 'analytical')
        fprintf('a = %f, b = %f, c = %f\n', a_damp_eff, b_damp, c_grow_damp)
        error_rate_d = abs(coefficients_d(2)-b_damp)*100/b_damp;
        fprintf('error damp rate [%%] = %f\n', error_rate_d)
    end
    % plot damp
    plot(t_analysis_d*1e3,x_analysis_d,'color','blue')
    if plot_envelope
        plot(t_env_d*1e3,x_env_up_d,'color','red','LineWidth',3)
    end
    fit_on_t_env_d = coefficients_d(1)*exp(coefficients_d(2)*t_env_d)+coefficients_d(3);
    plot(t_env_d*1e3,fit_on_t_env_d,'color','green','LineWidth',3);
    fprintf('Difference std = %f\n\n', std(fit_on_t_env_d-x_env_up_d))
end

function [f0, coefficients, t_analysis, x_analysis, t_env, x_env_up] = compute_rate(t_start, t_end, t_load, x_load, apply_bounds)
    
    % parameters
    method_envelope = 'loc_max'; % loc_max, envel
    method_initial_guess = 'avg_2points'; % avg_2points, lreg_3points (less reliable)
    iniGuess_lreg_inv = '\'; % inv, \
    
    % compute mean and slice array
    c0 = mean(x_load);
    t_analysis = t_load(t_load>=t_start & t_load<=t_end);
    x_analysis = x_load(t_load>=t_start & t_load<=t_end);

    % find local maxima
    [pks,locs] = findpeaks(x_analysis,'MinPeakHeight',c0);
    
    % compute the envelope
    if strcmp(method_envelope, 'envel')
        % e.g. 1e5, 'analytic' or 5, 'peaks'
        [x_env_up,x_env_down] = envelope(x_analysis, 1e5, 'analytic');
        t_env = t_analysis;
    elseif strcmp(method_envelope, 'loc_max')
        t_env = t_analysis(locs);
        x_env_up = x_analysis(locs);
    end
    
    % compute the initial guesses
    if strcmp(method_initial_guess, 'avg_2points')
        % % avg_2points
        % they can be whatever indices
        first_index = round(length(t_env)/4);
        second_index = round(length(t_env)*3/4);
        % compute the guesses
        b0 = log((x_env_up(first_index)-c0)/(x_env_up(second_index)-c0))/(t_env(first_index)-t_env(second_index));
        a0 = (x_env_up(first_index)-c0)*exp(-b0*t_env(first_index));
    elseif strcmp(method_initial_guess, 'lreg_3points')
        % % it seems worse than avg_2points
        % they must be consecutive indices
        second_index = round(length(t_env)*2/4);
        first_index = second_index-1;
        third_index = second_index+1;
        % contruct D
        S1 = 0;
        S2 = S1 + 0.5* (x_env_up(second_index)+x_env_up(first_index)) * (t_env(second_index)-t_env(first_index));
        S3 = S2 + 0.5* (x_env_up(third_index)+x_env_up(second_index)) * (t_env(third_index)-t_env(second_index));
        D_11 = (t_env(second_index)-t_env(first_index))^2 + (t_env(third_index)-t_env(first_index))^2;
        D_12 = (t_env(second_index)-t_env(first_index))*S2 + (t_env(third_index)-t_env(first_index))*S3;
        D_22 = S1^2 + S2^2 + S3^2;
        D_matrix = [D_11 D_12; D_12 D_22];
        % contruct the c vector
        c1 = ((x_env_up(second_index)-x_env_up(first_index)) * (t_env(second_index)-t_env(first_index))+ ...
              (x_env_up(third_index)-x_env_up(first_index)) * (t_env(third_index)-t_env(first_index)));
        c2 = ((x_env_up(second_index)-x_env_up(first_index))*S2 + (x_env_up(third_index)-x_env_up(first_index))*S3);
        if strcmp(iniGuess_lreg_inv, 'inv')
            % method 1 to invert D and compute b0
            D_inverse = inv(D_matrix);
            b0 = D_inverse(2,1)*c1 + D_inverse(2,2)*c2;
        elseif strcmp(iniGuess_lreg_inv, '\')
            % method 2 to invert D and compute b0
            b0T = D_matrix\[c1; c2];
            b0 = b0T(2);
        end
        % compute a0 and c0
        a0 = (x_env_up(first_index)-x_env_up(second_index))/exp(b0*(t_env(first_index)-t_env(second_index)));
        c0 = x_env_up(first_index)-a0*exp(b0*t_env(first_index));
    end
    fprintf('a0 = %f, b0 = %f, c0 = %f\n', a0, b0, c0)
    
    % compute the fit of the local maxima
    if apply_bounds
        Lower = [0 -Inf 0.9*c0];
        Upper = [+Inf +Inf 1.1*c0];
    else
        Lower = [];
        Upper = [];
    end
    fo = fitoptions('Method','NonlinearLeastSquares','StartPoint',[a0 b0 c0],...
                    'MaxIter',5000,'MaxFunEvals',5000,...
                    'TolFun',1e-12,'TolX',1e-12,...
                    'DiffMinChange',1e-10,'DiffMaxChange',1e-1,...
                    'Algorithm','Trust-Region', 'Robust','LAR',...
                    'Lower', Lower, 'Upper', Upper);
    g = fittype('a*exp(b*x)+c','indep','x', 'options', fo);
    f0 = fit(t_env', x_env_up', g);
    coefficients = coeffvalues(f0);
    fprintf('aN = %f, bN = %f, cN = %f\n', coefficients(1), coefficients(2), coefficients(3))
end

