% Metropolis-Adjusted Langevin Algorithm (MALA)
% Extends the unadjusted Langevin (ULA) by adding a Metropolis-Hastings
% accept/reject step, making the chain exactly target the Gibbs distribution
% p(x) ~ exp(-beta * U(x)) even for finite step sizes.
%
% Proposal:  Y = X_k - gradU(X_k)*dt + sqrt(2*dt/beta) * Z,  Z ~ N(0,1)
% Accept Y with probability:  alpha = min(1, pi(Y)*q(Y->X) / (pi(X)*q(X->Y)))
% where pi(x) ~ exp(-beta*U(x)) and q(x->y) is the Langevin proposal density.

function mala_mcmc()
    %% Parameters
    nSteps   = 20000;   % total number of MALA steps
    dt       = 0.05;    % step size (can be larger than ULA thanks to MH correction)
    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 MALA trajectory
    X = zeros(1, nSteps+1);
    accepted = zeros(1, nSteps);  % 1 if accepted, 0 if rejected
    X(1) = x0;
    sigma2 = 2 * dt / beta;       % proposal variance
    sigma  = sqrt(sigma2);

    % Langevin proposal mean: mu(x) = x - gradU(x)*dt
    propMean = @(x) x - gradU(x) * dt;

    % Log proposal density: q(x -> y) = N(y; mu(x), sigma^2)
    logq = @(x, y) -0.5 * (y - propMean(x)).^2 / sigma2;

    for k = 1:nSteps
        xk = X(k);

        % Propose
        Y = propMean(xk) + sigma * randn;

        % Log acceptance ratio:
        % log alpha = [ -beta*U(Y) + logq(Y, xk) ] - [ -beta*U(xk) + logq(xk, Y) ]
        log_alpha = (-beta * U(Y)  + logq(Y, xk)) ...
                  - (-beta * U(xk) + logq(xk, Y));

        % Accept or reject
        if log(rand) < log_alpha
            X(k+1) = Y;
            accepted(k) = 1;
        else
            X(k+1) = xk;
            accepted(k) = 0;
        end
    end

    acceptRate = mean(accepted);

    %% 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 4 subplots
    fig = figure('Name','MALA Demo','Color','w',...
                 'Position',[150 50 950 900]);

    % --- Panel 1: potential landscape + moving particle ---
    ax1 = subplot(4,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');

    % --- Panel 2: trajectory over time with accept/reject coloring ---
    ax2 = subplot(4,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]);
    ylim(ax2, [-2.5 2.5]);
    xlabel(ax2, 'Step');  ylabel(ax2, 'X_k');
    title(ax2, 'MALA Sample Path (blue = accepted, red = rejected)', 'FontSize', 13);

    % --- Panel 3: running acceptance rate ---
    ax4 = subplot(4,1,3,'Parent',fig);
    hold(ax4,'on'); grid(ax4,'on');
    acceptLine = animatedline(ax4, 'Color', [0.47 0.67 0.19], 'LineWidth', 1.5);
    plot(ax4, [0 nSteps], [0.574 0.574], 'k--', 'LineWidth', 1);  % optimal rate reference
    xlim(ax4, [0 nSteps]);
    ylim(ax4, [0 1]);
    xlabel(ax4, 'Step'); ylabel(ax4, 'Rate');
    title(ax4, 'Running Acceptance Rate (dashed = 0.574 optimal)', 'FontSize', 13);

    % --- Panel 4: histogram vs Gibbs ---
    ax3 = subplot(4,1,4,'Parent',fig);
    hold(ax3,'on'); grid(ax3,'on');
    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)}', 'MALA Samples'}, ...
           'Location', 'northeast');

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

    %% Animate
    frameSkip = 20;
    burnIn = round(nSteps * 0.1);

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

        % Update trajectory
        if k <= frameSkip
            idx = 1:k;
        else
            idx = max(1,k-frameSkip+1):k;
        end
        addpoints(trajLine, idx-1, X(idx));

        % Mark rejected steps with red dots (sparse, every frameSkip)
        if k > 1
            rejIdx = idx(idx > 1);
            rejIdx = rejIdx(accepted(rejIdx-1) == 0);
            if ~isempty(rejIdx)
                plot(ax2, rejIdx-1, X(rejIdx), '.', 'Color', [0.85 0.33 0.1], 'MarkerSize', 4);
            end
        end

        % Update running acceptance rate
        if k > 1
            runRate = cumsum(accepted(1:min(k-1,nSteps))) ./ (1:min(k-1,nSteps));
            clearpoints(acceptLine);
            addpoints(acceptLine, 1:min(k-1,nSteps), runRate);
        end

        % Update histogram
        startIdx = max(1, min(k, burnIn));
        set(histBars, 'Data', X(startIdx:k));
        ylim(ax3, [0 max(gibbs_pdf)*1.3]);

        % Step counter
        set(stepText, 'String', sprintf('Step: %d / %d  |  Accept: %.1f%%', ...
            k-1, nSteps, 100*mean(accepted(1:max(1,k-1)))));

        drawnow limitrate;
    end

    % Final update
    set(histBars, 'Data', X(burnIn:end));
    ylim(ax3, [0 max(gibbs_pdf)*1.3]);
    set(stepText, 'String', sprintf('Step: %d / %d  |  Accept: %.1f%%  (done)', ...
        nSteps, nSteps, 100*acceptRate));

    % Burn-in marker on trajectory
    plot(ax2, [burnIn burnIn], [-2.5 2.5], 'r--', 'LineWidth', 1);
    text(ax2, burnIn + nSteps*0.01, 2.2, 'burn-in', 'Color', 'r', 'FontSize', 9);

    % Print summary
    fprintf('\n=== MALA 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('Acceptance rate: %.2f%%\n', 100*acceptRate);
    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'); hold(ax4,'off');
end
