% Comparison: Unadjusted Langevin (ULA) vs Metropolis-Adjusted Langevin (MALA)
% Both use the same potential, beta, step size, and random seed so the
% comparison is fair. Shows that ULA develops bias at larger dt while MALA
% remains exact thanks to the Metropolis-Hastings correction.

function compare_ula_mala()
    %% Parameters
    nSteps   = 20000;
    dt       = 0.05;    % deliberately large to expose ULA bias
    beta     = 2.0;
    x0       = -2.0;

    % Potential: double-well  U(x) = (x^2 - 1)^2
    U      = @(x) (x.^2 - 1).^2;
    gradU  = @(x) 4*x.*(x.^2 - 1);

    %% Precompute both trajectories with the SAME noise sequence
    rng(42);  % fix seed for reproducibility
    noiseSeq = randn(1, nSteps);
    unifSeq  = rand(1, nSteps);  % for MH accept/reject

    sigma2 = 2 * dt / beta;
    sigma  = sqrt(sigma2);
    propMean = @(x) x - gradU(x) * dt;
    logq     = @(x, y) -0.5 * (y - propMean(x)).^2 / sigma2;

    % --- ULA ---
    Xula = zeros(1, nSteps+1);
    Xula(1) = x0;
    for k = 1:nSteps
        Xula(k+1) = propMean(Xula(k)) + sigma * noiseSeq(k);
    end

    % --- MALA ---
    Xmala    = zeros(1, nSteps+1);
    accepted = zeros(1, nSteps);
    Xmala(1) = x0;
    for k = 1:nSteps
        xk = Xmala(k);
        Y  = propMean(xk) + sigma * noiseSeq(k);
        log_alpha = (-beta*U(Y)  + logq(Y, xk)) ...
                  - (-beta*U(xk) + logq(xk, Y));
        if log(unifSeq(k)) < log_alpha
            Xmala(k+1) = Y;
            accepted(k) = 1;
        else
            Xmala(k+1) = xk;
            accepted(k) = 0;
        end
    end

    %% Theoretical Gibbs distribution
    xGrid = linspace(-3, 3, 500);
    gibbs_unnorm = exp(-beta * U(xGrid));
    Zn = trapz(xGrid, gibbs_unnorm);
    gibbs_pdf = gibbs_unnorm / Zn;

    burnIn = round(nSteps * 0.1);
    binEdges = linspace(-3, 3, 60);

    %% =====================================================================
    %  Figure 1: Side-by-side animated comparison
    %  =====================================================================
    fig1 = figure('Name','ULA vs MALA Comparison','Color','w',...
                  'Position',[80 80 1200 850]);

    % ---- Row 1: Potential with both particles ----
    ax_pot = subplot(3,2,[1 2],'Parent',fig1);
    hold(ax_pot,'on'); grid(ax_pot,'on');
    plot(ax_pot, xGrid, U(xGrid), 'k-', 'LineWidth', 2);
    dotULA  = plot(ax_pot, x0, U(x0), 'o', 'MarkerSize', 11, ...
                   'MarkerFaceColor', [0.85 0.33 0.1], 'MarkerEdgeColor','w');
    dotMALA = plot(ax_pot, x0, U(x0), 'o', 'MarkerSize', 11, ...
                   'MarkerFaceColor', [0 0.45 0.74], 'MarkerEdgeColor','w');
    xlim(ax_pot,[-3 3]); ylim(ax_pot,[-0.2 3]);
    xlabel(ax_pot,'x'); ylabel(ax_pot,'U(x)');
    title(ax_pot,'Double-Well Potential  U(x) = (x^2-1)^2','FontSize',13);
    legend(ax_pot,{'U(x)','ULA particle','MALA particle'},'Location','northeast');

    % ---- Row 2 left: ULA trajectory ----
    ax_traj_ula = subplot(3,2,3,'Parent',fig1);
    hold(ax_traj_ula,'on'); grid(ax_traj_ula,'on');
    trajULA = animatedline(ax_traj_ula,'Color',[0.85 0.33 0.1],'LineWidth',0.4);
    xlim(ax_traj_ula,[0 nSteps]); ylim(ax_traj_ula,[-3 3]);
    xlabel(ax_traj_ula,'Step'); ylabel(ax_traj_ula,'X_k');
    title(ax_traj_ula,'ULA Trajectory','FontSize',12);

    % ---- Row 2 right: MALA trajectory ----
    ax_traj_mala = subplot(3,2,4,'Parent',fig1);
    hold(ax_traj_mala,'on'); grid(ax_traj_mala,'on');
    trajMALA = animatedline(ax_traj_mala,'Color',[0 0.45 0.74],'LineWidth',0.4);
    xlim(ax_traj_mala,[0 nSteps]); ylim(ax_traj_mala,[-3 3]);
    xlabel(ax_traj_mala,'Step'); ylabel(ax_traj_mala,'X_k');
    title(ax_traj_mala,'MALA Trajectory','FontSize',12);

    % ---- Row 3 left: ULA histogram ----
    ax_hist_ula = subplot(3,2,5,'Parent',fig1);
    hold(ax_hist_ula,'on'); grid(ax_hist_ula,'on');
    plot(ax_hist_ula, xGrid, gibbs_pdf, 'r-', 'LineWidth', 2.5);
    hULA = histogram(ax_hist_ula, Xula(1), 'BinEdges', binEdges, ...
                     'Normalization','pdf', ...
                     'FaceColor',[0.85 0.33 0.1],'EdgeColor','w','FaceAlpha',0.6);
    xlim(ax_hist_ula,[-3 3]); ylim(ax_hist_ula,[0 max(gibbs_pdf)*1.4]);
    xlabel(ax_hist_ula,'x'); ylabel(ax_hist_ula,'Density');
    title(ax_hist_ula,'ULA Histogram vs Gibbs','FontSize',12);
    legend(ax_hist_ula,{'Gibbs','ULA'},'Location','northeast');

    % ---- Row 3 right: MALA histogram ----
    ax_hist_mala = subplot(3,2,6,'Parent',fig1);
    hold(ax_hist_mala,'on'); grid(ax_hist_mala,'on');
    plot(ax_hist_mala, xGrid, gibbs_pdf, 'r-', 'LineWidth', 2.5);
    hMALA = histogram(ax_hist_mala, Xmala(1), 'BinEdges', binEdges, ...
                      'Normalization','pdf', ...
                      'FaceColor',[0 0.45 0.74],'EdgeColor','w','FaceAlpha',0.6);
    xlim(ax_hist_mala,[-3 3]); ylim(ax_hist_mala,[0 max(gibbs_pdf)*1.4]);
    xlabel(ax_hist_mala,'x'); ylabel(ax_hist_mala,'Density');
    title(ax_hist_mala,'MALA Histogram vs Gibbs','FontSize',12);
    legend(ax_hist_mala,{'Gibbs','MALA'},'Location','northeast');

    infoText = annotation(fig1,'textbox',[0.30 0.005 0.45 0.03],...
        'String','Step: 0','FontSize',11,'EdgeColor','none',...
        'HorizontalAlignment','center');

    %% Animate figure 1
    frameSkip = 30;
    for k = 1:frameSkip:nSteps+1
        % Particles on potential
        set(dotULA,  'XData', Xula(k),  'YData', U(Xula(k)));
        set(dotMALA, 'XData', Xmala(k), 'YData', U(Xmala(k)));

        % Trajectories
        if k <= frameSkip
            idx = 1:k;
        else
            idx = max(1,k-frameSkip+1):k;
        end
        addpoints(trajULA,  idx-1, Xula(idx));
        addpoints(trajMALA, idx-1, Xmala(idx));

        % Histograms (use post-burn-in when available)
        si = max(1, min(k, burnIn));
        set(hULA,  'Data', Xula(si:k));
        set(hMALA, 'Data', Xmala(si:k));
        ylim(ax_hist_ula,  [0 max(gibbs_pdf)*1.4]);
        ylim(ax_hist_mala, [0 max(gibbs_pdf)*1.4]);

        % Info
        if k > 1
            ar = 100*mean(accepted(1:min(k-1,nSteps)));
        else
            ar = 0;
        end
        set(infoText, 'String', sprintf('Step %d / %d   |   MALA accept: %.1f%%   |   dt = %.3f', ...
            k-1, nSteps, ar, dt));

        drawnow limitrate;
    end

    % Final histograms with post-burn-in data
    set(hULA,  'Data', Xula(burnIn:end));
    set(hMALA, 'Data', Xmala(burnIn:end));
    ylim(ax_hist_ula,  [0 max(gibbs_pdf)*1.4]);
    ylim(ax_hist_mala, [0 max(gibbs_pdf)*1.4]);

    % Burn-in markers
    plot(ax_traj_ula,  [burnIn burnIn], [-3 3], 'r--', 'LineWidth',1);
    plot(ax_traj_mala, [burnIn burnIn], [-3 3], 'r--', 'LineWidth',1);
    set(infoText, 'String', sprintf('Done  |  MALA accept: %.1f%%  |  dt = %.3f', ...
        100*mean(accepted), dt));

    %% =====================================================================
    %  Figure 2: Quantitative diagnostics
    %  =====================================================================
    fig2 = figure('Name','ULA vs MALA Diagnostics','Color','w',...
                  'Position',[100 60 1000 700]);

    % ---- Panel 1: KL divergence proxy (bin-wise comparison) over time ----
    ax_kl = subplot(2,2,1,'Parent',fig2);
    hold(ax_kl,'on'); grid(ax_kl,'on');
    checkpoints = round(linspace(burnIn, nSteps+1, 80));
    kl_ula  = zeros(size(checkpoints));
    kl_mala = zeros(size(checkpoints));
    binCenters = (binEdges(1:end-1) + binEdges(2:end)) / 2;
    gibbs_at_bins = exp(-beta * U(binCenters));
    gibbs_at_bins = gibbs_at_bins / sum(gibbs_at_bins);  % discrete pmf

    for j = 1:length(checkpoints)
        n = checkpoints(j);
        % ULA histogram counts
        cULA = histcounts(Xula(burnIn:n), binEdges, 'Normalization','probability');
        cULA(cULA == 0) = 1e-12;
        kl_ula(j) = sum(gibbs_at_bins .* log(gibbs_at_bins ./ cULA));
        % MALA histogram counts
        cMALA = histcounts(Xmala(burnIn:n), binEdges, 'Normalization','probability');
        cMALA(cMALA == 0) = 1e-12;
        kl_mala(j) = sum(gibbs_at_bins .* log(gibbs_at_bins ./ cMALA));
    end
    semilogy(ax_kl, checkpoints, kl_ula, '-o', 'Color',[0.85 0.33 0.1],...
             'LineWidth',1.5,'MarkerSize',4,'MarkerFaceColor',[0.85 0.33 0.1]);
    semilogy(ax_kl, checkpoints, kl_mala,'-s', 'Color',[0 0.45 0.74],...
             'LineWidth',1.5,'MarkerSize',4,'MarkerFaceColor',[0 0.45 0.74]);
    set(ax_kl,'YScale','log');
    xlabel(ax_kl,'Step'); ylabel(ax_kl,'KL( Gibbs || samples )');
    title(ax_kl,'KL Divergence from Gibbs (log scale, lower is better)','FontSize',12);
    legend(ax_kl,{'ULA','MALA'},'Location','northeast');

    % ---- Panel 2: Running mean (should converge to 0 by symmetry) ----
    ax_mean = subplot(2,2,2,'Parent',fig2);
    hold(ax_mean,'on'); grid(ax_mean,'on');
    idx_range = burnIn:50:nSteps+1;
    runMeanULA  = cumsum(Xula(burnIn:end))  ./ (1:(nSteps-burnIn+2));
    runMeanMALA = cumsum(Xmala(burnIn:end)) ./ (1:(nSteps-burnIn+2));
    plot(ax_mean, burnIn:nSteps+1, runMeanULA,  '-','Color',[0.85 0.33 0.1],'LineWidth',1.5);
    plot(ax_mean, burnIn:nSteps+1, runMeanMALA, '-','Color',[0 0.45 0.74],'LineWidth',1.5);
    plot(ax_mean, [burnIn nSteps+1], [0 0], 'k--', 'LineWidth',1);
    xlabel(ax_mean,'Step'); ylabel(ax_mean,'Running Mean');
    title(ax_mean,'Running Mean (theory = 0)','FontSize',12);
    legend(ax_mean,{'ULA','MALA','Theory'},'Location','northeast');

    % ---- Panel 3: Autocorrelation ----
    ax_acf = subplot(2,2,3,'Parent',fig2);
    hold(ax_acf,'on'); grid(ax_acf,'on');
    maxLag = 200;
    acfULA  = autocorr_manual(Xula(burnIn:end),  maxLag);
    acfMALA = autocorr_manual(Xmala(burnIn:end), maxLag);
    plot(ax_acf, 0:maxLag, acfULA,  '-','Color',[0.85 0.33 0.1],'LineWidth',2);
    plot(ax_acf, 0:maxLag, acfMALA, '-','Color',[0 0.45 0.74],'LineWidth',2);
    xlabel(ax_acf,'Lag'); ylabel(ax_acf,'ACF');
    title(ax_acf,'Autocorrelation Function','FontSize',12);
    legend(ax_acf,{'ULA','MALA'},'Location','northeast');

    % ---- Panel 4: Final overlaid histograms ----
    ax_ovl = subplot(2,2,4,'Parent',fig2);
    hold(ax_ovl,'on'); grid(ax_ovl,'on');
    histogram(ax_ovl, Xula(burnIn:end), 'BinEdges', binEdges, ...
              'Normalization','pdf','FaceColor',[0.85 0.33 0.1],...
              'EdgeColor','none','FaceAlpha',0.45);
    histogram(ax_ovl, Xmala(burnIn:end), 'BinEdges', binEdges, ...
              'Normalization','pdf','FaceColor',[0 0.45 0.74],...
              'EdgeColor','none','FaceAlpha',0.45);
    plot(ax_ovl, xGrid, gibbs_pdf, 'r-', 'LineWidth', 2.5);
    xlim(ax_ovl,[-3 3]); ylim(ax_ovl,[0 max(gibbs_pdf)*1.4]);
    xlabel(ax_ovl,'x'); ylabel(ax_ovl,'Density');
    title(ax_ovl,'Final Histograms vs Gibbs','FontSize',12);
    legend(ax_ovl,{'ULA','MALA','Gibbs'},'Location','northeast');

    % Effective sample size (via ACF)
    essULA  = ess_from_acf(acfULA,  nSteps - burnIn + 1);
    essMALA = ess_from_acf(acfMALA, nSteps - burnIn + 1);

    %% Print summary
    fprintf('\n=========================================\n');
    fprintf('   ULA vs MALA Comparison Summary\n');
    fprintf('=========================================\n');
    fprintf('Potential: U(x) = (x^2 - 1)^2\n');
    fprintf('beta = %.1f,  dt = %.4f,  steps = %d\n', beta, dt, nSteps);
    fprintf('-----------------------------------------\n');
    fprintf('                     ULA        MALA\n');
    fprintf('-----------------------------------------\n');
    fprintf('Accept rate:         100%%       %.1f%%\n', 100*mean(accepted));
    fprintf('Sample mean:     %8.4f    %8.4f\n', mean(Xula(burnIn:end)), mean(Xmala(burnIn:end)));
    fprintf('Sample std:      %8.4f    %8.4f\n', std(Xula(burnIn:end)), std(Xmala(burnIn:end)));
    fprintf('Final KL div:    %8.4f    %8.4f\n', kl_ula(end), kl_mala(end));
    fprintf('ESS (from ACF):  %8.0f    %8.0f\n', essULA, essMALA);
    fprintf('=========================================\n');
    fprintf('(Theory: mean=0 by symmetry)\n\n');
end

%% Helper: manual autocorrelation (avoids Econometrics Toolbox dependency)
function acf = autocorr_manual(x, maxLag)
    x = x(:) - mean(x);
    n = length(x);
    c0 = sum(x.^2);
    acf = zeros(1, maxLag+1);
    for lag = 0:maxLag
        acf(lag+1) = sum(x(1:n-lag) .* x(1+lag:n)) / c0;
    end
end

%% Helper: effective sample size from ACF (initial positive sequence estimator)
function ess = ess_from_acf(acf, n)
    % Sum ACF until it first goes negative (Geyer's initial positive criterion)
    tau = 1;
    for k = 2:length(acf)
        if acf(k) < 0
            break;
        end
        tau = tau + 2 * acf(k);
    end
    ess = n / tau;
end
