% Langevin MCMC via Euler-Maruyama Scheme
% Demonstrates that the Langevin SDE: dX = -U'(X)dt + sqrt(2/beta) dW
% has the Gibbs distribution p(x) ~ exp(-beta * U(x)) as its stationary measure.

function langevin_mcmc()
    %% Parameters
    nSteps   = 20000;   % total number of EM steps
    dt       = 0.005;   % time step
    beta     = 2.0;     % inverse temperature
    x0       = -2.0;    % initial position

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

    %% Precompute trajectory
    X = zeros(1, nSteps+1);
    X(1) = x0;
    noise_coeff = sqrt(2 * dt / beta);
    for k = 1:nSteps
        X(k+1) = X(k) - gradU(X(k)) * dt + noise_coeff * randn;
    end

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

    %% Set up figure with 3 subplots
    fig = figure('Name','Langevin MCMC Demo','Color','w',...
                 'Position',[150 80 900 800]);

    % --- Top panel: potential landscape + moving particle ---
    ax1 = subplot(3,1,1,'Parent',fig);
    hold(ax1,'on'); grid(ax1,'on');
    plot(ax1, xGrid, U(xGrid), 'b-', 'LineWidth', 2);
    particleDot = plot(ax1, x0, U(x0), 'ro', 'MarkerSize', 12, ...
                       'MarkerFaceColor','r', 'LineWidth', 1.5);
    xlim(ax1, [-2.5 2.5]);
    ylim(ax1, [-0.2 2.5]);
    xlabel(ax1, 'x'); ylabel(ax1, 'U(x)');
    title(ax1, 'Double-Well Potential  U(x) = (x^2 - 1)^2', 'FontSize', 13);
    legend(ax1, {'U(x)', 'Particle'}, 'Location', 'northeast');

    % --- Middle panel: trajectory over time ---
    ax2 = subplot(3,1,2,'Parent',fig);
    hold(ax2,'on'); grid(ax2,'on');
    trajLine = animatedline(ax2, 'Color', [0 0.45 0.74], 'LineWidth', 0.5);
    xlim(ax2, [0 nSteps*dt]);
    ylim(ax2, [-2.5 2.5]);
    xlabel(ax2, 'Time');  ylabel(ax2, 'X(t)');
    title(ax2, 'Sample Path (Euler-Maruyama)', 'FontSize', 13);

    % --- Bottom panel: histogram vs Gibbs ---
    ax3 = subplot(3,1,3,'Parent',fig);
    hold(ax3,'on'); grid(ax3,'on');
    gibbsLine = plot(ax3, xGrid, gibbs_pdf, 'r-', 'LineWidth', 2.5);
    histBars  = histogram(ax3, X(1), 'BinEdges', linspace(-2.5,2.5,60), ...
                          'Normalization', 'pdf', ...
                          'FaceColor', [0.3 0.75 0.93], 'EdgeColor', 'w', ...
                          'FaceAlpha', 0.7);
    xlim(ax3, [-2.5 2.5]);
    ylim(ax3, [0 max(gibbs_pdf)*1.3]);
    xlabel(ax3, 'x'); ylabel(ax3, 'Density');
    title(ax3, 'Histogram of Samples vs Gibbs Distribution', 'FontSize', 13);
    legend(ax3, {'Gibbs  p(x) \propto e^{-\beta U(x)}', 'Samples'}, ...
           'Location', 'northeast');

    stepText = annotation(fig, 'textbox', [0.62 0.01 0.35 0.04], ...
                          'String', 'Step: 0', 'FontSize', 11, ...
                          'EdgeColor', 'none', 'HorizontalAlignment', 'right');

    %% Animate
    % Show every Nth frame for smooth animation
    frameSkip = 20;
    burnIn = round(nSteps * 0.1);  % mark burn-in visually

    for k = 1:frameSkip:nSteps+1
        % Update particle on potential
        set(particleDot, 'XData', X(k), 'YData', U(X(k)));

        % Update trajectory (add points in batch)
        times = ((k-frameSkip):min(k,nSteps+1)-1) * dt;
        if k <= frameSkip
            times = (0:k-1) * dt;
            addpoints(trajLine, times, X(1:k));
        else
            idx = max(1,k-frameSkip+1):k;
            addpoints(trajLine, (idx-1)*dt, X(idx));
        end

        % Update histogram with samples so far (use post-burn-in for cleaner view)
        startIdx = max(1, min(k, burnIn));
        set(histBars, 'Data', X(startIdx:k));

        % Rescale histogram y-axis
        ylim(ax3, [0 max(gibbs_pdf)*1.3]);

        % Step counter
        set(stepText, 'String', sprintf('Step: %d / %d', k-1, nSteps));

        drawnow limitrate;
    end

    % Final full histogram with all post-burn-in samples
    set(histBars, 'Data', X(burnIn:end));
    ylim(ax3, [0 max(gibbs_pdf)*1.3]);
    set(stepText, 'String', sprintf('Step: %d / %d  (done)', nSteps, nSteps));

    % Change trajectory burn-in portion color
    % Add a vertical line indicating burn-in boundary
    plot(ax2, [burnIn*dt, burnIn*dt], [-2.5 2.5], 'r--', 'LineWidth', 1);
    text(ax2, burnIn*dt + 0.5*dt, 2.2, 'burn-in', 'Color', 'r', 'FontSize', 9);

    % Print summary
    fprintf('\n=== Langevin MCMC Summary ===\n');
    fprintf('Potential: U(x) = (x^2 - 1)^2  (double well)\n');
    fprintf('beta = %.1f,  dt = %.4f,  steps = %d\n', beta, dt, nSteps);
    fprintf('Post burn-in samples: %d\n', nSteps - burnIn + 1);
    fprintf('Sample mean:  %.4f  (theory: 0 by symmetry)\n', mean(X(burnIn:end)));
    fprintf('Sample std:   %.4f\n', std(X(burnIn:end)));
    fprintf('============================\n');

    hold(ax1,'off'); hold(ax2,'off'); hold(ax3,'off');
end
