% EE278B tree reconstruction homework problem
% Bernd Bandemer/Abbas El Gamal, October 2011
% based on material by Min-Wook Jeong/Balaji Prabakhar, Fall 2010

function treeReconstruction
    % Problem parameters
    nLayers = 12;  % depth of the tree (number of layers)
    nTrials = 50000; % number of tree realizations
    vEpsilon = linspace(0,0.5,15);    % values for the noise probability in each edge of the tree
    
    % Create variables to hold data
    nEpsilon = length(vEpsilon);
    errorRateMAP = zeros(1,nEpsilon);  % this will hold the error rates
    errorRateMaj = zeros(1,nEpsilon); 
    
    % Run the experiment
    tic
    for iEpsilon = 1:nEpsilon
        disp(sprintf('Step %d/%d...',iEpsilon,nEpsilon));
        epsilon = vEpsilon(iEpsilon);
        for iTrial = 1:nTrials
            % Generate a realization of the tree (number of leafs: 2^(nLayers-1) )
            root = binornd(1, 1/2);  % generate a root with Bern(1/2)
            noise = generateData(nLayers, epsilon);  % generate the noise 
            leafs = mod(root+noise, 2);  % the leafs nodes are root+noise
            
            % Generate the likelihood values at the leafs
            leafLikelihoods = zeros(2, 2^(nLayers-1) );  
                % first row: L_{i,0}, second row: L_{i,1}
            leafLikelihoods(1,:) = 1-leafs;   % as in lecture slide 4-7
            leafLikelihoods(2,:) = leafs;
            
            % MAP detection of the root node based on the leafs
            rootLikelihoods = MAP(leafLikelihoods, epsilon);
            rootHatMAP = rootLikelihoods(2) > rootLikelihoods(1);  
                % the root is detected to be one iff L_{1,1}>L_{1,0}
                
            % Detection of root node using majority vote
            rootHatMaj = MajorityVote( leafs );
            
            % Count error statistics
            if rootHatMAP~=root  % if wrongly detected...
                errorRateMAP(iEpsilon) = errorRateMAP(iEpsilon) + 1;  % ... count it as an error
            end
            if rootHatMaj~=root  % if wrongly detected...
                errorRateMaj(iEpsilon) = errorRateMaj(iEpsilon) + 1;  % ... count it as an error
            end
        end
    end
    timeTaken = toc;
    errorRateMAP = errorRateMAP / nTrials;  % convert error count to error probability
    errorRateMaj = errorRateMaj / nTrials;  
    
    % Make a plot
    figure; hold on;
    plot(vEpsilon, errorRateMAP, 'bx-');
    plot(vEpsilon, errorRateMaj, 'r+-');
    legend('MAP rule','Majority vote');
    xlabel('epsilon');
    ylabel('error probability');
    title(sprintf('%d layers, %d trials, %.1f seconds', nLayers, nTrials, timeTaken));
    ylim([0 0.7]);
    print('-dpdf',sprintf('treeReconstruction_%dlayers_%dtrials',nLayers,nTrials));
end

function leafs = generateData(nLayers, epsilon)
    if nLayers == 1  % first layer -- no noise involved
        leafs = 0;  
    else   % k-th layer: generate (k-1)th layer and add independent Bern(epsilon) noise
        layerAbove = generateData(nLayers-1, epsilon);
        leafs = kron(layerAbove, [1 1]);   % make two replicas of every node of the layer above
        leafs = leafs + binornd(1,epsilon, 1,2^(nLayers-1));  % add independent noise
        leafs = mod( leafs, 2 );  % modulo 2 addition, not regular addition
    end
end

function rootLikelihoods = MAP( leafLikelihoods, eps )
% Compute the likelihoods at the root node
% Inputs: 
%   leafLikelihoods [2 x N], where
%       N is the number of leaves
%       leafLikelihood(1,i) is L_{i,0} of the lecture notes (page 4-5)
%       leafLikelihood(2,i) is L_{i,1} of the lecture notes
%   epsilon
%       the parameter of the Bern(epsilon) noise across each edge in the tree
% Outputs: 
%   rootLikelihoods [2 x 1], where
%       rootLikelihoods(1) is L_{1,0}
%       rootLikelihoods(2) is L_{1,1}

%%% FILL in your implementation here.
    epsbar  = 1-eps;
    nleaves = length(leafLikelihoods(1,:));
    while (nleaves ~= 1)
        nleaves = nleaves/2;     % The layer above has half the nodes
        temp = zeros(2,nleaves); % Temporary scratch space
        for i = 1:nleaves
            temp(1,i) = ([epsbar eps]*leafLikelihoods(:,2*i-1))*...
                        ([epsbar eps]*leafLikelihoods(:,2*i));
            temp(2,i) = ([eps epsbar]*leafLikelihoods(:,2*i-1))*...
                        ([eps epsbar]*leafLikelihoods(:,2*i));
        end
        temp = temp/max(max(temp));
        leafLikelihoods = temp;
    end
    rootLikelihoods = leafLikelihoods;
end

function root = MajorityVote( leafs )
% Compute the likelihoods at the root node
% Inputs: 
%   leafs [1 x N], where
%       N is the number of leaves
%       leafs(i) is the value (0 or 1) of the i-th leaf random variable
% Outputs: 
%   root [1 x 1], your guess of the root random variable (0 or 1)

%%% FILL in your implementation here.
    N = length(leafs);
    root = sum(leafs)>(N/2);
end
