% Number of Bernoulli draws until the first "1" is geometric
% Textbook: Part I, "Continuous Random Variable", exercise (2022, Classwork)
%
% Draw X ~ Ber(p) repeatedly and count the draws up to and including the
% first "1". The count K ~ Geo(p), p_K[k] = p(1-p)^(k-1), k = 1,2,...,
% with E[K] = 1/p.

clearvars; close all
rng(0)

p = 0.2;        % success probability
N = 1e5;        % number of experiments

K = zeros(1, N);
for n = 1:N
    k = 1;
    while rand >= p     % "0" with probability 1-p
        k = k + 1;
    end
    K(n) = k;
end

k = 1:30;
p_emp = sum(K == k', 2)'/N;
p_theory = p*(1-p).^(k-1);

figure
bar(k, p_emp)
hold on
stem(k, p_theory, 'r', 'LineWidth', 1.5)
hold off
xlabel('k')
ylabel('p_K[k]')
legend('simulation', 'theory, Geo(p)')
grid on

fprintf('E[K]:   theory = %.3f, simulation = %.3f\n', 1/p, mean(K))
fprintf('Var[K]: theory = %.3f, simulation = %.3f\n', (1-p)/p^2, var(K))