% L6Bdirichlet   Demo of Dirichlet distribution -- with animation
%
% This time we are working with a six-sided (possibly loaded) die.
% The parameter vector contains SIX probability parameters,
% and its distribution is a six-dimensional Dirichlet distribution.
%
% Visualizations are lower-dimensional projections because we
% don't have a 5-dimensional display.
%
% Input is a sequence of die rolls (integers from 1 to 6).
% Alternatively, if no input (or empty input) is given,
% then the user is prompted for the rolls, one at a time.

function [postpara,X]=L6Bdirichlet(X, plot_every)

askuser=false;
if nargin<1, X=[]; end
if isempty(X), askuser=true; end
if nargin<2, plot_every=1; end

if askuser
    nmax = 1e6;
else
    nmax = length(X);
end

% Prior parameters for Dirichlet
priorpara = ones(1,6);
postpara  = priorpara;

% Fine grid of values for a single parameter, for plotting
gridsize = 1000;
u = linspace(0,1,gridsize)';

clf

if askuser
    fprintf('Input roll results (between 1 and 6) one at a time, or 0 to terminate.\n');
end

for n=0:nmax
    if n>0
        % One more observation => increase one of the parameters.
        postpara(X(n)) = postpara(X(n))+1;
    end
    
    if mod(n,plot_every)==0
        
        % Marginal densities of each of the six parameters
        f = zeros(gridsize, 6);
        for i=1:6
            % Marginal distribution of i'th parameter is
            % beta distribution that lumps together the
            % other parameters as the "other" choice.
            a = postpara(i);
            b = sum(postpara) - a;
            f(:,i) = betapdf(u, a, b);
        end
    
        plot(u,f, 'linewidth', 2);
        legend('P1','P2','P3','P4','P5','P6');
        title(sprintf('after %d observations', n));
        xlabel('parameter value');
        ylabel('density');
        grid on
    
        if not(askuser) && n<nmax
            pause
        end
    end
    
    if askuser
        x = input('next roll: ');
        while x<0 || x>6 || x~=round(x)
            x = input('invalid input\nnext roll: ');
        end
        if x==0, break; end
        X = [X x];
    end
end