% Script to implement the Lloyd-Max Algorithm
%
% Functions you must define:
%
% p(x)  density of X
% F(x)  cdf of X
% G(x)  int_{-\infty}^x t p(t) dt
%
% Additional functions that use the above functions:
%
% cbar(c)
%
% If c is a vector of coefficients, with c1 < ... < cN,
% then the best partition is
%
%       (-inf,v1] (v1,v2] ... (vN-2,vN-1] (vN-1,inf)
%
% where v1 = (c1+c2)/2, ... , vN-1 = (cN-1+cN)/2.
% The function cbar takes c as input and returns
% v = [ v1 ... vN-1 ].  The function cbar first sorts the vector c.
%
% getnum(v)
%
% This function takes v and returns n = [ n1 ... nN ],
% where nj = int_{V_j} x p(x) d x.
%
% getden(v)
%
% This function takes v and returns d = [ d1 ... dN ],
% where dj = int_{V_j} p(x) d x.
%
%==============================================================
%
% INITIALIZE
%
xmin = -3;  % range to plot density
xmax =  6;
x = linspace(xmin,xmax,200);
maxloop = 200;
fprintf('Max iterations = %g\n',maxloop)
%
% PLOT DENSITY FUNCTION
%
figure(1)
subplot(2,1,1)
plot(x,p(x))
title('Density Function')
%
% GET DESIRED NUMBER OF LEVELS  N  AND PLOT INITIAL UNIFORM
% PARTITION ON TOP OF DENSITY
%
N = input('How many levels (N) do you want? ');
v = linspace(xmin,xmax,N+1);
v = v(2:N); % v = [ v_1 ... v_{N-1} ]'
fprintf('Press any key to continue...\n')
pause
%
% DRAW PARTITION LINES
%
set(gca,'XTick',v,'XTickLabel',[],'XGrid','on','GridLineStyle','-')
title('Density and Initial Partition') % REPLACE TITLE
%
% INITIALIZE FOR FIRST PASS OF LLOYD-MAX ALG AND PRINT MSE.
%
num = getnum(v);
den = getden(v);
c = num./den;
oldTERM= c*num'; % num is a row vector
fprintf('Initial MSE = E[X^2] - %g\n',oldTERM)
subplot(2,1,2)
%
% BEGIN LOOP
%
% flag = 1;    % Use with WHILE to step through algorithm
% while flag
for i = 1:maxloop
   v = cbar(c);
   num = getnum(v);
   den = getden(v);
   c = num./den;
%   fprintf('MSE = E[X^2] - %g\n',c*num')  % num is a row vector
%   plot(x,p(x))
%   title('Density and Current Partition')
%   set(gca,'XTick',v,'XTickLabel',[],'XGrid','on','GridLineStyle','-')
%   flag = input('Continue? Enter 0 to stop: ');
end
%
% PRINT FINAL PARTITION AND COEFFICIENTS
%
%  v
%  c
%
% COMPUTE FINAL MSE
%
newTERM = c*num';
fprintf('Final   MSE = E[X^2] - %g\n',c*num')  % num is a row vector
% fprintf('Reduction in MSE = %g\n',newTERM-oldTERM)
%
% PLOT FINAL PARTITION ON TOP OF DENSITY
%
plot(x,p(x))
title('Density and Final Partition')
set(gca,'XTick',v,'XTickLabel',[],'XGrid','on','GridLineStyle','-')
%
% PLOT LLOYD-MAX QUANTIZER IN DIFFERENT FIGURE
%
% figure(2)
% [xx,yy] = stepfcn(c,v,xmin,xmax);
% plot(xx,yy,x,x)
% title('Lloyd-Max Quantizer')
