%% Problem Sheet 8
% Stochastic Simulation

close all; clear all; clc;

%% Exercise 3

% a)

m = 20;
N = 10^5;

% energy function and probability distribution
f = @(x) - exp(-((x-3).^2)/2) - exp(-((x-9).^2)/3) - exp(-((x-16).^2)/4);
T = 0.1;

% the corresponding Boltzmann distribution
pi = exp(-f(1:m)/T);
pi = pi/sum(pi);

% initialize and choose random initial state
X = zeros(N+1,1);
X(1) = ceil(rand()*m);
for i = 1:N
    % generate a proposal
    if rand() <= 0.5
        Y = min(X(i)+1, m);
    else
        Y = max(X(i)-1, 1);
    end
    alpha = pi(Y)/pi(X(i));
    if rand() <= alpha
        X(i+1) = Y;
    else
        X(i+1) = X(i);
    end
end

% theoretical distribution
figure(1);
bar(1:m, pi, 1);
title('Theoretical distribution');

% empirical distribution
figure(2);
max_count = hist(X,1:m);
bar(1:m, max_count/sum(max_count), 1);
title('Empirical distribution using a)');


% b)

% now: more than one temperature
T = [0.1 0.3 0.6];
% probability of switching between chains
beta = 0.1;

% random initial states
X = zeros(N+1, length(T));
X(1,:) = ceil(rand(1,length(T))*m);
for i = 1:N
    if rand() <= beta
        % choose index for swapping
        t = ceil(rand()*(length(T)-1));
        % calculate acceptance probability
        alpha = exp(- f(X(i,t+1)) / T(t) ...
                    - f(X(i,t)) / T(t+1) ...
                    + f(X(i,t)) / T(t) ...
                    + f(X(i,t+1)) / T(t+1));
        Y = X(i,:);
        % if accepted: swap
        if rand() <= alpha
            Y(t) = X(i,t+1);
            Y(t+1) = X(i,t);
        end
        X(i+1,:) = Y;
    else
        % take independent steps
        U = rand(length(T),1);
        Y = zeros(length(T),1);
        for j = 1:length(T)
            if U(j) <= 0.5
                Y(j) = min(X(i,j)+1, m);
            else
                Y(j) = max(X(i,j)-1, 1);
            end
            alpha = exp(( f(X(i,j)) - f(Y(j)) ) / T(j));
            if rand() <= alpha
                X(i+1,j) = Y(j);
            else
                X(i+1,j) = X(i,j);
            end
        end
    end
end

% empirical distribution
figure(3);
h = hist(X(:,1),1:m);
bar(1:m, h/sum(h), 1);
title('Empirical distribution using b)');


%% Exercise 4 

clear all; close all; 

% sample size and parameters
J = 3;
M = 5;
m = 10;
T = 1; % (increase the temperature for a better mixed distribution)
N = 10^3; %10^4;

% we represent every state (originally an m x m matrix) by a vector
% of length m*m. The index (i,j) in the matrix corresponds to (i-1)*m + j
% in the vector.
index = @(i,j) (i-1)*m + j;

% adjacency matrix (for m x m lattice with periodic boundary conditions)
edges = zeros(m*m, m*m);
for i = 1:m
    for j = 1:m
        above = [mod(i-2,m)+1, j];
        right = [i, mod(j,m)+1];
        edges(index(i,j), index(above(1),above(2))) = 1;
        edges(index(above(1),above(2)), index(i,j)) = 1;
        edges(index(i,j), index(right(1),right(2))) = 1;
        edges(index(right(1),right(2)), index(i,j)) = 1;
    end
end

% empty vectors for sums and most frequent values
sums = zeros(N,1);
mostfrequent = zeros(N,1);
frequency = zeros(N,1); % optional

X = ceil(rand(m*m,1)*M); % initial state: independent uniform RVs
for steps = 1:N
    Y = X;
    % simulate U_{i,j}
    U = zeros(m*m, m*m);
    for s = 1:size(U,1)
        for t = 1:size(U,2)
            if edges(s,t) == 1
                U(s,t) = rand()*exp(J*(X(s)==X(t))/T);
            end
        end
    end
    % determine where to place edges for clusters
    B = (U >= 1);
    % collect clusters
    clusters = zeros(m*m, 1);
    for s = 1:length(clusters)
        if clusters(s) == 0
            clusters(s) = max(clusters)+1;
            todo = s;
            while ~isempty(todo)
                neighbors = find(edges(todo(1),:)==1);
                for k=1:length(neighbors)
                    if (clusters(neighbors(k)) == 0) && (B(todo(1), neighbors(k)) == 1)
                        clusters(neighbors(k)) = clusters(todo(1));
                        todo = [todo, neighbors(k)];
                    end
                end
                todo = todo(2:end);
            end
        end
    end
    % draw a new value for each cluster
    for i = 1:max(clusters)
        X(clusters==i) = ceil(rand()*M);
    end
    % remember sum of values ...
    sums(steps) = sum(sum(X));
    % ... and most frequent value
    [max_count, max_index] = max(hist(X(:), 1:M));
    mostfrequent(steps) = max_index;
    frequency(steps) = max_count / (m*m);
end

% rearrange values into matrix for output
X_mat = zeros(m,m);
for i=1:m
    for j=1:m
        X_mat(i,j) = X(index(i,j));
    end
end

X_mat
hm = HeatMap(X_mat, 'Colormap', 'autumn');

figure(1);
hist(mostfrequent, 1:M);
title('Histogram of most frequent values');

fprintf('The estimated expected value: %f\n', mean(sums));

% optional
figure(2);
plot(frequency);
xlabel('iteration');
title('Frequency of the most frequent value');