function [TARGET APPROX] = p22_writeup() 

function pmf = makeRandPmf()
  %% computes a random pmf for our cube, normalizing
  %% so that it integrates to 1
  pmf = rand([5 5 5]);
  pmf = pmf / sum(sum(sum(pmf)));
end;

function marg = marginalX(jointpmf)
  %% returns the marginal probability density of the PDF
  %% as a function of X (i.e., marginalizes over Y and Z)
   marg = sum(sum(jointpmf, 2),3);
end

function marg = marginalY(jointpmf) 
  %% returns the marginal probability density of the PDF
  %% as a function of Y (i.e., marginalizes over X and Z)
   marg = shiftdim(sum(sum(jointpmf,1),3));
end

function marg = marginalZ(jointpmf)  
  %% returns the marginal probability density of the PDF
  %% as a function of Z (i.e., marginalizes over X and Y)
   marg = shiftdim(sum(sum(jointpmf,1),2),2);
end;

function nextapprox = ipfIter(truemargX, truemargY, truemargZ, approxpmf, iteration),
   %% given our true observed marginals (truemargX, truemargY,
   %% truemargZ) and our current approximation for the pmf, 
   %% computes the next approximation.  iteration is our iteration index

   nextapprox = [];
   amX = marginalX(approxpmf);
   amY = marginalY(approxpmf);
   amZ = marginalZ(approxpmf);
   for i=1:size(approxpmf,1),
     for j=1:size(approxpmf,2),
       for k=1:size(approxpmf,3),
	 
           if mod(iteration,3) == 0,
               nextapprox(i,j,k) = approxpmf(i,j,k)*(truemargX(i)/amX(i));
           elseif mod(iteration,3) == 1,
               nextapprox(i,j,k) = approxpmf(i,j,k)*(truemargY(j)/amY(j));
           elseif mod(iteration,3) == 2,
               nextapprox(i,j,k) = approxpmf(i,j,k)*(truemargZ(k)/amZ(k));
           end;

       end;
     end;
   end;
end

function approximation = doipf(truepmf, approximation, iteration)
   approximation = ipfIter(marginalX(truepmf), marginalY(truepmf), marginalZ(truepmf), approximation, iteration);
end;

function divergence = computeDivergence(P,Q),
   % computes the KL-divergence between 3D distributions P and Q
   divergence = sum(sum(sum(  P.*log(P./Q) )));
end;


%% first come up with a random 5x5x5 distribution 
TARGET = makeRandPmf
MARGINAL_X = marginalX(TARGET)
MARGINAL_Y = marginalY(TARGET)
MARGINAL_Z = marginalZ(TARGET)

APPROX = makeRandPmf;
margXs = [marginalX(TARGET)];
margYs = [marginalY(TARGET)];
margZs = [marginalX(TARGET)];
KLdiverge = [];

NUM_ITERATIONS = 20;

for i=1:NUM_ITERATIONS,
  APPROX = doipf(TARGET,APPROX,i);
  margXs = [margXs marginalX(APPROX)];
  margYs = [margYs marginalY(APPROX)];
  margZs = [margZs marginalZ(APPROX)];
  KLdiverge = [KLdiverge computeDivergence(TARGET, APPROX)];
end;

% plot marginals per iteration
figure(1);
clf;
subplot(3,1,1);
plot([1:5],margXs);
title('marginal P(a)');
subplot(3,1,2);
plot([1:5],margYs);
title('marginal P(b)');
subplot(3,1,3);
plot([1:5],margZs);
title('marginal P(c)');

figure(2);
clf;
plot([1:NUM_ITERATIONS],KLdiverge);
title('KL divergence between actual and approximated distributions as a function of IPF iteration');

% plot the actual distributions using slices
figure(3);
subplot(1,2,1);
[x y z] = meshgrid([1:5],[1:5],[1:5]);
slice(x,y,z,TARGET, [0:0.1:5], 0, 0);
title('Actual distribution');
subplot(1,2,2);
[x y z] = meshgrid([1:5],[1:5],[1:5]);
slice(x,y,z,APPROX, [0:0.1:5], 0, 0);
title('Approximated distribution');


%MARGINAL_X_DIFF = MARGINAL_X - marginalX(APPROX)
%MARGINAL_Y_DIFF = MARGINAL_Y - marginalY(APPROX)
%MARGINAL_Z_DIFF = MARGINAL_Z - marginalZ(APPROX)


end


