% Control3A_3.m
% Solution to problem 3 of control 3. May 14, 2021.
% JMS, May.2021

% Set data
clear
N = 60;     % number of balls
nop = 50;   % number of extraction operations
nr = 1e4;   % number of experiment realizations

% Perform one experiment of nop substitutions
b = zeros(N,nop+1);    % value of each ball before each substitution
b(:,1) = 1;            % initial value of each ball
for iop = 1:nop                     % loop on extractions
  ib = randi(N);                    % index of the extracted ball
  b(:,iop+1) = b(:,iop);            % ball values are not changed
  b(ib,iop+1) = mod(b(ib,iop),3)+1; % ... except that of ball ib
     % notice: mod(j,3) returns 1,2,0 for j=1,2,3
end

% Plot number of balls of each value, as a function of time (extractions)
figure(11)
t = 0:nop;
plot(t,sum(b==1),t,sum(b==2),t,sum(b==3),'LineWidth',1.5)
legend('1','2','3')
xlabel('Number of substitutions')
ylabel('Number of balls')
set(gca,'FontSize',15)
grid on

% Repeat the experiment nr times
for ir = 1:nr                 % loop on experiments
  b = ones(N,1);              % initial ball values
  for iop = 1:nop             % loop on extractions
    ib = randi(N);            % index of the extracted ball
    b(ib) = mod(b(ib),3)+1;   % change value of ball ib
  end
  nf(:,ir) = [sum(b==1);sum(b==2);sum(b==3)]; % final number of balls of each value
end

% Plot probabilities for final number of balls of each value
for h = 0:N
  P(:,h+1) = sum(nf==h,2)/nr; % estimated prob. that nf=h, for each ball value
  % An equivalent way to do the same (estimating also the uncertainties)
%   for iv = 1:3                           % loop on ball values
%     x = (nf(iv,:)==h);                   % random variable (true or false)
%     P(iv,h+1) = meanValue(x);            % prob. is fraction of true x
%     DP(iv,h+1) = standDev(x)/sqrt(nr-1); % uncertainty of mean values
%   end
end
figure(12)
h = 0:N;
plot(h,P,'LineWidth',1.5)
legend('1','2','3')
% hold on, plot(h,P-DP,h,P+DP), hold off
xlabel('Final number of balls')
ylabel('Probability estimate')
set(gca,'FontSize',15)
grid on

% Print some probability estimates
x = (nf(2,:)>nf(1,:));         % true or false for each experiment
P1 = meanValue(x);             % prob = fraction of times that it was true
DP1 = standDev(x)/sqrt(nr-1);  % meanVal uncertainty (not requested in the exam)
x = (20<=nf(1,:) & nf(1,:)<=25);
P2 = meanValue(x);
DP2 = standDev(x)/sqrt(nr-1);
x = (nf(2,:)==15 | nf(3,:)==15);
P3 = meanValue(x);
DP3 = standDev(x)/sqrt(nr-1);
fprintf('         P( n(2)>n(1) ) = %5.3f +/- %5.3f\n',P1,DP1)
fprintf('      P( 20<=n(1)<=25 ) = %5.3f +/- %5.3f\n',P2,DP2)
fprintf('P( n(2)==15 | n(3)=15 ) = %5.3f +/- %5.3f\n',P3,DP3)
