%% Scalar Quanitization 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%EXAMPLES OF USE
%
%[count,z,d,errVec,RRF] = quant(@(z)ones(1,length(z)),0,1,6,opt.method=1);
%[count,z,d,errVec,RRF,DzVec] = quant(@(z)exp((-z.^2)/.1),-1,1,7);
%[count,z,d,errVec] = quant(@(z)6*z.^2.*exp(-2*z.^3),1,2,11);
%[count,z,d,errVec] = quant(@(z)2*z.^2+1,1,2,8);
%
%INPUT
%   p = p(z) density funciton (fhandle)
%   a = Left endpoint
%   b = Right endpoint
%   L = Level for number of partitions, i.e. 2^L
%   opt is an optional data structure with the following elements
%       opt.Guess = 0 Local method in the Coarse grid to fine grid (default) 
%                   1 equidistribution
%                   2 uniform partition
%                   3 skewed partition
%       opt.Pdata = 0 turn data plots off (default)
%                   1 turn data plots on
%       opt.method = 1 for local method
%                    2 for Exact Netwon's method
%                    3 Approximation (default)
%       opt.ErrTol = Control for stopping the iterations (default=10^(-12))
%       opt.dispCount = displays count and error after every iteration
%                       0 off (default)
%                       1 on 
%       opt.Itmax = number of Maximum iterations allowed (default=500)
%
%OUPUT
%   Count = the number of the iterations computed
%   z = Centroid after Approximation using Newton's method
%   d = Computer partitions based off of centroid
%   errVec = |stepsize*DE|_\infty (i.e. independent of stepsize)
%   RRF = Vector of the Residual Reduction Factor
%         Norm(Dz^{k+1})/Norm(Dz^{k})
%   DzVec = Vector of Norms of Residual
%           Only used for opt.method = 2 or opt.method = 3
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Function quant
%   Parent function
function [count,z,d,errVec,RRF,DzVec] = quant(p,a,b,L,opt)
n = 2^L;
%   Set parameters
if nargin == 4
   opt = [];
end
if isfield(opt,'guess') == 0
    guess = 0;
else
    guess = opt.guess;
end
%   Select method
if isfield(opt,'method') == 0
    method = 3;
else
    method = opt.method;
end
%   Set Error Tolerance for stopping interations
if isfield(opt,'ErrTol') == 0
    ErrTol = 10^(-12);
else
    ErrTol = opt.ErrTol;
end
%   Set number of Maximum iterations allowed
if  isfield(opt,'Itmax') == 0
    Itmax = 500;
else
    Itmax = opt.Itmax;
end
if  isfield(opt,'Pdata') == 0
    Pdata = 0;
else
    Pdata = opt.Pdata;
end
if  isfield(opt,'dispCount') == 0
    dispCount = 0;
else
    dispCount = opt.dispCount;
end
%   Set initial guess and partition
[z,d]=Init(a,b,n,p,guess);
%   Preform iterations with selected method
if method == 1
    err=10;
    count = 1;
    RRF=zeros(Itmax,1);
    errVec=zeros(Itmax,1);
    errVec(1)=1;
    while  (count < Itmax) && (err > ErrTol)
        count = count + 1;
        [z,d,err] = Local(z,p,d);
        errVec(count) = err;
        RRF(count) = errVec(count)/errVec(count-1);
    end
    RRF = RRF(1:count);
    errVec = errVec(count);
else
    if method == 2 
        [z,d,count,RRF,DzVec,errVec] = Iter(z,p,d,L,ErrTol,Itmax,method,dispCount);
    else
        method = 3;
	p = @(z) p(z)  + .6
        [z,d,count,RRF,DzVec,errVec] = Iter(z,p,d,L,ErrTol,Itmax,method,dispCount);
    end
end
%   Plot data
if Pdata==1
%% Figure 1 = Partition Vs Difference of Quanitizer, z_{k+1} - z_k' and Density function
    figure(1)
    %subplot(1,2,1)
    plot(z,zeros(1,length(z)),'*b')
    grid off
%    plot(d(2:end-1),(z(2:end)-z(1:end-1)),'--blue','LineWidth',3)
    %grid on
    %title('Partition vs Difference of Quantizer, z_{k+1} - z_k');
    %xlabel('Parition');
    %ylabel('Difference of Quanitizer z_{k+1} - z_k');
    %subplot(1,2,2)
    %plot(a:1/n:b,p(a:1/n:b),'--b','LineWidth',3)
    grid on
%    hold on
%    plot(z(1:32:end),repmat(max(p(d))/2 + min(p(d))/2,1,length(z(1:32:end))),'.black')
    title('Density function vs Partition');
    ylabel('Density function');
    xlabel('Parition');
    hold off
%% Figure 3 = log plot of l_2 norm of Residual vs Iteration
    if method == 3 || method == 2
        figure(3)
        loglog(1:length(DzVec(2:end)),DzVec(2:end),'--b','LineWidth',3)
        grid on
        title('log plot of l_2 norm of Residual vs Iteration');
        xlabel('Iteration');
        ylabel('l_{2} norm of Residual');
    end
%% Figure 4 = Residual Reduction Factor vs Iteration   
    figure(4)
    plot(1:length(RRF(3:end)),RRF(3:end),'--b','LineWidth',3)
    grid on
    title('Residual Reduction Factor vs Iteration');
    xlabel('Iteration');
    ylabel('Residual Reduction Factor');
%% Figure 5 = log plot of calculated error vs Iteration   
    if method == 2 || method == 3
        figure(5)
        loglog(1:count-1,errVec,'--b','LineWidth',3)
                grid on
        title('log plot of ||\nabla DE ||_{\infty} vs Iteration');
        xlabel('Iteration');
        ylabel('||\nabla DE ||_{\infty}');
    end
end
end

%% Function: Init
%   Chooses initial guess for Newton's method
function [z,d] = Init(a,b,n,p,K)
if nargin == 4
    K = 0;
end
if K == 1
    p = @(z)(p(z)).^(1/3);
    z = a:1/n:b;
    d = zeros(1,length(z)+1);
    d(2:end-1) = z(1:end-1)/2+z(2:end)/2;
    d(1) = a;
    d(end) = b;
    bot = quad(p,a,b);
    F = zeros(1,n+1);
    for k = 1:length(z)
        F(k) = quad(p,a,d(k+1))/bot;
    end
    x = linspace(a,b,1000)';
    f = interp1(z,F,x,'linear');
    count = 1;
    zC = 1;
    while count < 1001
        if (abs(f(count) - f(1) - zC*(1/n)) < 10^(-7))
            z(zC) = x(k);
            zC = zC + 1;
        end
        count = count + 1;
    end    
elseif K == 2
    z = a:1/n:b;
    d = zeros(1,length(z)+1);
    d(1) = a;
    d(end) = b;
elseif K == 3
    z = (a+1/n):1/(2*n):(b-1/n);
    d(2:2*n-3) = z(1:2*n-4)/2+z(2:2*n-3)/2;
    d(1) = a;
    d(2*n-2) = b;
else
    m = 2^3;
    err = 10;
    z = a:1/m:b;
    d = zeros(1,length(z)+1);
    d(1) = a;
    d(end) = b;
    count = 0;
    while  (count < 2000) && (err > 10^(-14))
        count = count + 1;
        [z,d,err] = Local(z,p,d);  
    end
    count = 0;
    while n/m > 1        % refine to find grid
        zMid = z(2:end)/2 + z(1:end-1)/2; 
        zTemp = zeros(1,2*length(zMid)+1);
        zTemp(2:2:end) = zMid;
        zTemp(1:2:end) = z;
        z = zTemp;
        n = n/2;
        while  (count < 200) && (err > 10^(-6))
            count = count + 1;
            [z,d,err] = Local(z,p,d);  
        end
    end
    d = zeros(1,length(z)+1);
    d(2:end-1) = z(1:end-1)/2+z(2:end)/2;
end
d(end) = b;
d(1) = a;
end

%% Function: Trap
%   6th order newton-cotes numerical integration on a vector
function total = Trap(f,a,b)
mid = a/2 + b/2;
midL = a/2 + mid/2;
midU = mid/2 + b/2;
x = [a; midL; mid; midU; b];
y = feval(f,x);
sums = 7*y(1,:) + 32*y(2,:) + 12*y(3,:) + 32*y(4,:) +7*y(5,:);
total = ((b-a).*sums)/90;
end

%% Function: Local
%   Computes one iteration of the Local method, updates d, then z
function  [z,d,err] = Local(z,p,d)
T=@(x) x.*p(x);
zTemp=z;
d(2:end-1)=z(1:end-1)/2+z(2:end)/2;
z = Trap(T,d(1:end-1),d(2:end))./Trap(p,d(1:end-1),d(2:end));
err = max(abs(z-zTemp));
end

%% Function: Exact
%   Computes the Gradient and Hessian of the Energy function given by variance.
function [D2E,DE,d] = Exact(z,p,d)
if (size(z,2)==1)
    z=z';
end
n = size(z,2);
d(2:n) = z(1:n-1)/2+z(2:n)/2; %compute decision regions
Ediag = zeros(1,n);
IntTemp = Trap(p,d(1:end-1),d(2:end)); %diagonal for Hessian
Ediag(2:(n-1)) = 2*IntTemp(2:(n-1)) ...
               + z((2:(n-1))).*p(d(3:n)) - z((2:(n-1))).*p(d(2:(n-1))) ...
               - d(3:n).*p(d(3:n)) + d(2:(n-1)).*p(d(2:(n-1)));
Ediag(1) = 2*Trap(p,d(1),d(2)) ...
             +z(1)*p(d(2)) - d(2)*p(d(2));
Ediag(n) = 2*Trap(p,d(n),d(n+1)) ...
             - z(n)*p(d(n)) + d(n)*p(d(n));
EoffU = p(d(2:n)).*(z(1:n-1)-d(2:n)); %upper diagonal for Hessian
EoffL = p(d(2:n)).*(z(2:n)-d(2:n)); %lower diagonal for Hessian
EU = sparse(1:n-1,2:n,EoffU,n,n);
EL = sparse(2:n,1:n-1,EoffL,n,n);
D2E = sparse(1:n,1:n,Ediag,n,n);
D2E = D2E + EU - EL;
DE = 2*z.*Trap(p,d(1:n),d(2:n+1)) - 2*Trap(@(x)x.*p(x),d(1:n),d(2:n+1)); %Gradient
end

%% Function: Approx
%   Computes the Gradient and approximated Hessian of the Energy function given by variance.
function [D2E,DE,d] = Approx(z,p,d)
if (size(z,2)==1)
    z=z';
end
n=size(z,2);
d(2:n) = z(1:n-1)/2+z(2:n)/2;
Beta = (d(2:n+1)-z(1:n)).*p(d(2:n+1)); %upper diagonal for Approx Hessian
Alpha = (z(1:n)-d(1:n)).*p(d(1:n)); %lower diagonal for Approx Hessian
D2EDiag = zeros(1,n);
D2EDiag(2:n-1) = Beta(2:n-1) + Alpha(2:n-1); 
EUOffU = sparse(1:n-1,2:n,Beta(1:n-1),n,n);
ELOffL = sparse(2:n,1:n-1,Alpha(2:n),n,n);
D2EDiag(1) = 2*Alpha(1)+Beta(1); % boundary conditions on diagonal
D2EDiag(n) = Alpha(end)+2*Beta(end);
D2E = sparse(1:n,1:n,D2EDiag,n,n);
D2E = D2E - EUOffU - ELOffL; %Hessian
DE = 2*z.*Trap(p,d(1:n),d(2:n+1)) - 2*Trap(@(x) x.*p(x),d(1:n),d(2:n+1)); %Gradient
end

%% Function: Iter
%   Computes the iterations for newtons method
function [z,d,count,RRF,DzVec,errVec] = Iter(z,p,d,L,ErrTol,maxIt,method,dispCount)
err = 10;
count = 1;
RRF = ones(maxIt,1);
DzTemp = zeros(maxIt,1);
errVec = zeros(maxIt,1);
DzTemp(1) = 1;
if method == 2    
    [D2E,DE,d] = Exact(z,p,d);
else
    [D2E,DE,d] = Approx(z,p,d);
end
N = length(DE);
while (count < maxIt) && (err > ErrTol)
    StepVec = diff(d);
    P = sparse(1:N,1:N,1./StepVec,N,N);    
    if det(P*D2E*P) ~= 0 % && det(P*D2E*P) ~= Inf %rescale Hessian
        Dz = -(P*D2E*P)\(P*DE');
        Dz = P*Dz;
        if dispCount == 1
            disp('P*D2E*P used as Hessian')
        end
    elseif  det(D2E) ~= Inf && det(D2E) ~= 0 
        Dz = -(D2E\DE' + 10^(-16)*spdiags(ones(N,1),0,N,N))\(DE');
        if dispCount == 1
            disp('D2E used as Hessian')
        end
    else
        Dz = -(D2E+40^(-L)*spdiags(ones(N,1),0,N,N))\DE';
        if dispCount == 1
            disp('D2E + 10^L*I used as Hessian')
        end
    end    
    while (z(end) + Dz(end) > d(end)) || (z(1) + Dz(1) < d(1)) %rescale residual
        Dz = Dz/2;
    end
    c = z+Dz';
    while count > 1 && (min(c(2:end)-c(1:end-1)) < 0) %rescale residual
         Dz = Dz/2;
         c = z+Dz';
    end
    while count > 1 && err>errVec(count) %ensure err decreases
        Dz = Dz/2;
        c = c+Dz';
            if method == 2    
                [D2E,DE,d] = Exact(c,p,d);
            else
                [D2E,DE,d] = Approx(c,p,d);
            end
        err = max(abs(DE./StepVec));
        if max(Dz) == 0 && min(Dz) == 0
            break
        end
    end
    if (count > 3) && (norm(Dz)/DzTemp(count-1) > 1) %break if iterations blow up
        break
    end
    %err = max(abs(DE));
    err = max(abs(DE./StepVec));
    count = count + 1;
    errVec(count) = err;
    DzTemp(count) = norm(Dz);
    RRF(count) = DzTemp(count)/DzTemp(count-1);
    z = c;
    if dispCount == 1
       fprintf('The count is: %2.0f, with error: %1.6e \n',count, err)
    end
    if method == 2    
        [D2E,DE,d] = Exact(z,p,d);
    elseif method == 3
        [D2E,DE,d] = Approx(z,p,d);
    end
end
DzVec = DzTemp(1:count);
errVec = errVec(2:count);
RRF = RRF(1:count);
end