%% ========= Params =========
sigma_e = 0.05;
steps   = 100;                 % inner EM steps per unit time
dt      = 1 / steps;
N_tau   = 1;
T       = 5000;                % recorded outer steps
alpha   = 0.7;                 % single alpha
trials  = 100;                 % number of simulations

%% ========= Storage =========
num_abund = nan(trials,1);   % numerical abundance corr per trial
num_grow  = nan(trials,1);   % numerical growth corr per trial
th_abund  = nan(trials,1);   % analytical abundance corr per trial
th_grow   = nan(trials,1);   % analytical growth corr per trial
kept      = false(trials,1);

%% ========= Trials =========
for jj = 1:trials
    jj
    [gamma, lambda] = GammaLambda_Angle(alpha);
    [CT, x0, CRT]   = compute_correlation_coefficient(gamma, lambda);
    th_abund(jj) = CT;
    th_grow(jj)  = CRT;

    if ~all(x0 > 0), continue; end

    % init state
    x = x0(1:2);
    R = x0(3:5);

    x_hist  = zeros(2, T);
    x_hist1 = zeros(2, T);

    for k = 1:T
        x_old = x;
        for kk = 1:steps
            x = x + dt * ((-1 + gamma * R') .* x')';
            drift  = R - R.^2 - R .* (x * lambda');
            noise  = randn(1,3) .* R * sqrt(dt) * sigma_e;
            ItoFix = R * (sigma_e^2) / 2;
            R = R + dt * (drift + ItoFix) + noise;
        end
        x_hist(:, k)  = x';
        x_hist1(:, k) = log(x ./ x_old)';
    end

    tail = floor(T/10)+1 : T;
    num_abund(jj) = corr(x_hist(1,tail)',  x_hist(2,tail)');
    num_grow(jj)  = corr(x_hist1(1,tail)', x_hist1(2,tail)');

    kept(jj) = true;
end

%% ========= Keep only successful trials =========
num_abund = num_abund(kept);
num_grow  = num_grow(kept);
th_abund  = th_abund(kept);
th_grow   = th_grow(kept);
trial_ids = 1:sum(kept);

fprintf('Successful trials: %d / %d\n', numel(num_abund), trials);

%% ========= Plots =========
figure('Name','Correlation comparison','Color','w');
tiledlayout(2,1);

% --- Abundance correlations ---
nexttile; hold on; grid on;
plot(trial_ids, num_abund, 'x', 'MarkerSize', 6, 'LineWidth', 1.2);
plot(trial_ids, th_abund,  'o', 'MarkerSize', 7, 'LineWidth', 1.2);
xlabel('Trial'); ylabel('CR(n_1, n_2)');
title('Abundance correlations');
legend('Numerical (X)','Analytical (O)','Location','best');

% --- Growth correlations ---
nexttile; hold on; grid on;
plot(trial_ids, num_grow, 'x', 'MarkerSize', 6, 'LineWidth', 1.2);
plot(trial_ids, th_grow,  'o', 'MarkerSize', 7, 'LineWidth', 1.2);
xlabel('Trial'); ylabel('CR(r_1, r_2)');
title('Growth correlations');
legend('Numerical (X)','Analytical (O)','Location','best');