function [X,r] = mlambda(W,a,p)
%
% [X,r] = mlambda(W,a,p) produces p optimal solutions to the integer 
%         least squares problem min_{x}(x-a)'*W^{-1}*(x-a) and the 
%         corresponding values of the objective function 
%
% Inputs:
%    W - n by n symmetric positive definite matrix. In GNSS, it is  
%        the covariane matrix of the real least squares estimator  
%        of the integer ambiguity vector.
%    a - n-dimensional real vector. In GNSS, it is the real least 
%        squares estimator of the ambiguity vector.
%    p - number of required optimal solutions with default value 1.
%
% Outputs:
%    X - n by p matrix. Its j-th column is the j-th optimal solution,
%        i.e., its objective function is the j-th smallest.
%    r - p-dimensional real vector for the p values of the objective 
%        function at the p optimal solutions, i.e.,
%        r(j)=(X(:,j)-a)'*W^{-1}*(X(:,j)-a)
%

% ------------------------------------------------------------------------
% Main references:
% [1] M. Al Borno, X.-W. Chang, and X. Xie. On "Decorrelation" in Solving 
%     Integer Least-Squares Problems for Ambiguity Determination, 
%     Survey Review, 46:37-49, 2014. 
% [2] X.-W. Chang, X. Yang, and T. Zhou, MLAMBDA: A Modified LAMBDA Method 
%     for Integer Least-squares Estimation, Journal of Geodesy, 79 (2005), 
%     pp. 552-565. 
%
% Authors: Xiao-Wen Chang, www.cs.mcgill.ca/~chang
%          Xiaohu Xie
% Copyright (c) 2011-2022. Scientific Computing Lab, McGill University.
% September 2011. Last revision: December 2022
% ------------------------------------------------------------------------

    % Check input arguments
    if nargin < 2 % input error
        error('Not enough input arguments!')
    end
    
    if nargin < 3
        p = 1;
    end
    
    if p <= 0 % input error
        error('Third input argument must be an integer bigger than 0!')
    end
    
    [m,n] = size(W);
    if m ~= n || m ~= size(a,1) || size(a,2) ~= 1  % input error
        error('Input arguments have a matrix dimension error!')
    end
    
    % Transform the ILS problem by reduction
    [L,d,Z,a] = reduction(W,a);
    
    % Search to find the p optimal solutions of the reduced ILS problem
    [Optis,r] = search(L,d,a,p);
    
    % Find the p optimal solutions of the original ILS problem by transformation
    X = Z'*Optis;
end

%======================================================================
%====================== Subfunctions of MLAMBDA =======================
%======================================================================

function [L,d,Z,a] = reduction(W,a)
%
% [L,d,Z,a] = reduction(W,a) computes the factorization W=Z'*L'*D*L*Z
%             and a:=Z^{-T}a.
%
% Inputs:
%    W - n by n symmetric positive definite matrix
%    a - n-dimensional real vector.
%
% Outputs:
%    L - n by n unit lower triangular matrix.
%    d - vector formed by the diagonal of the n by n diagonal matrix D.
%    Z - n by n unimodular matrix.
%    a - updated a.

% Subfunction: ldlp

% ------------------------------------------------------------------------
% Main reference:
%   M. Al Borno, X.-W. Chang, and X. Xie. On "Decorrelation" in Solving 
%   Integer Least-Squares Problems for Ambiguity Determination, 
%   Survey Review, 46:37-49, 2014. 
%
% Authors: Xiao-Wen Chang, www.cs.mcgill.ca/~chang
%          Xiaohu Xie
% Copyright (c) 2011-2016. Scientific Computing Lab, McGill University.
% September 2011. Last revision: April 2016
% ------------------------------------------------------------------------
    n = size(W,1);
    
    % Compute the L'DL factorization of W with minimum symmetric pivoting
    [L,d,Z,a] = ldlp(W,a);
    
    k = n - 1;
    
    while k > 0
        
        kp1 = k + 1;
        
        tmp = L(kp1,k);
        mu = round(tmp);
        if mu ~= 0
            tmp = tmp - mu*L(kp1,kp1);
        end
        
        delta = d(k) + tmp^2*d(kp1);
        
        if delta < d(kp1)
            if mu ~=0  % Perform size reductions
                % Size reduction on L(k+1,k)
                L(kp1:n,k) = L(kp1:n,k) - mu*L(kp1:n,kp1);
                Z(kp1,:) = Z(kp1,:) + mu*Z(k,:);
                a(k) = a(k) - mu*a(kp1);
                
                % Size reductions on L(k+2:n,k) for numerical stability
                for i = k+2:n
                    mu = round(L(i,k));
                    if mu ~= 0 
                        L(i:n,k) = L(i:n,k) - mu*L(i:n,i);
                        Z(i,:) = Z(i,:) + mu*Z(k,:);
                        a(k) = a(k) - mu*a(i);
                    end
                end
            end
    
            % Perform a symmetric permutation 
            eta = d(k)/delta;
            lambda = d(kp1)*L(kp1,k)/delta;
            d(k) = eta*d(kp1);
            d(kp1) = delta;
            L(k:kp1,1:k-1) = [-L(kp1,k),1; eta,lambda]*L(k:kp1,1:k-1);
            L(kp1,k) = lambda;
            L(k+2:n,[k kp1]) = L(k+2:n,[kp1 k]);
            Z([k kp1],:) = Z([kp1 k],:);
            a([k kp1]) = a([kp1 k]);
    
            if k < n-1
                k = k + 1;
            end
        else
            k = k - 1;
        end
    end
end

%======================================================================
%===================== Subfunctions of Reduction ======================
%======================================================================

function [L,d,P,a] = ldlp(W,a)
%
% [L,d,P,a] = ldlp(W,a) computes the L'DL factorization of W with minimum
%             symmetric piovting: W = P'L'DLP and computes a:=Pa.
%
% Inputs:
%    W - n by n symmetric positive definite matrix. 
%    a - n-dimensional vector.
%
% Outputs:
%    L - n by n unit lower triangular matrix.
%    d - vector formed by the diagonal of the n by n diagonal matrix D.
%    P - n by n permutation matrix.
%    a - updated a.

% ------------------------------------------------------------------------
% Main references:
% [1] X.-W. Chang, X. Yang, and T. Zhou, MLAMBDA: A Modified LAMBDA Method 
%     for Integer Least-squares Estimation, Journal of Geodesy, 79 (2005), 
%     pp. 552-565. 
%
% Authors: Xiao-Wen Chang, www.cs.mcgill.ca/~chang
%          Xiaohu Xie 
% Copyright (c) 2006-2016. Scientific Computing Lab, McGill University.
% September 2011. Last revision: April 2016
% ------------------------------------------------------------------------
    n = size(W,1);
    perm = 1:n;  % record the permutation information
    
    for k = n:-1:1
        [mindiag, q] = min(diag(W(1:k,1:k)));
	    if mindiag <= 0
		    error ('Input matrix W is not positive definite!');
        end
        perm([k q]) = perm([q k]);
        W([k q],:) = W([q k],:);
        W(:,[k q]) = W(:, [q k]);
        W(k,1:k-1) = W(k,1:k-1)/W(k,k);
        W(1:k-1,1:k-1) = W(1:k-1,1:k-1) - W(k,1:k-1)'*W(k,k)*W(k,1:k-1);   
        a([k q]) = a([q k]);
    end
    
    d = diag(W)';
    L = tril(W,-1) + eye(n);
    P = zeros(n);
    for j = 1:n
        P(j,perm(j)) = 1;
    end

end

%======================================================================
%====================== Subfunctions of MLAMBDA =======================
%======================================================================

function [Optis, fun] = search(L,d,a,p)
%  
% [Optis,fun] = search(L,d,a,p) produces the p optimal solutions to the ILS
%             problem  min_{z}f(z)=(z-a)'*(L'DL)^{-1}*(z-a) by search.
%
% Inputs:
%    L - n by n unit lower triangular matrix
%    d - n-dimensional real positive vector, d = diag(D)
%    a - n-dimensional real vector
%    p - number of required optimal solutions with default value 1
%
% Outputs:
%    Optis - n by p matrix whose j-th column is the j-th optimal solution
%    fun -  p-dimensional real vector for the p-values of the objective 
%           function at the p optimal solutions, referred to as residuals 
%           or distances, i.e., fun(j)=f(Optis(:,j))

% ------------------------------------------------------------------------
% Main reference:
%   X.-W. Chang, X. Yang, and T. Zhou, MLAMBDA: A Modified LAMBDA Method 
%   for Integer Least-squares Estimation, Journal of Geodesy, 79 (2005), 
%   pp. 552-565. 
%
% Authors: Xiao-Wen Chang, www.cs.mcgill.ca/~chang
%          Taha Ghassemi, Xiaohu Xie
% Copyright (c) 2011-2016. Scientific Computing Lab, McGill University.
% September 2011. Last revision: April 2016
% ------------------------------------------------------------------------
    n = length(d);
    
    % Initialization 
    z = zeros(n,1);        % current candidate solution
    zb = zeros(n,1);       % used to compute z
    step = zeros(n,1);     % step(k) is used to enumerate z(k) at level k
    dist = zeros(n+1,1);   % partial distance of a candidate solution
    S = zeros(n,n);        % used for computing zb
    path = n*ones(n,1);    % path(k) used for updating S(k,k:path(k)-1) 
    Optis = zeros(n,p);    % to store the p candidate solutions  
    fun = zeros(1,p);      % residuals at the p candidate solutions
    imax = p;
    maxDist = inf;         % initial ellipsoid bound 
    count = 0;             % initial number of candidate solutions
    ulevel = 1;
    
    % Determine which level to move to after z(1) is chosen at level 1.
    if p == 1            
        k0 = 2;
    else
        k0 = 1;
    end
    
    % Initialization at level n
    zb(n) = a(n);
    z(n) = round(zb(n));
    y(n) = zb(n) - z(n);
    step(n) = sign(y(n));
    k = n;
    
    endSearch = false;
    
    
    while ~endSearch
        path(ulevel:k-1) = k;
        for j = ulevel-1 : -1 : 1
            if path(j) < k 
                path(j) = k;
            else
                break;
            end
        end
        
        newDist = dist(k) + y(k) * y(k) / d(k);
        
        while newDist < maxDist
            if k ~= 1      % move to level k-1
                k = k - 1;
                dist(k) = newDist;
                for j = path(k) : -1 : k + 1
                    S(j-1,k) = S(j,k) - y(j) * L(j,k);
                end
                zb(k) = a(k) + S(k,k);
                z(k) = round(zb(k));
                y(k) = zb(k) - z(k);
                step(k) = sign(y(k));
            else      % store the found point and try next valid integer
                if count < p - 1     % store the first p-1 initial points
                    count  =  count + 1;
                    Optis(:, count) = z;
                    fun(count) = newDist;  % store f(z)
                else
                    Optis(:, imax) = z;
                    fun(imax) = newDist;
                    [maxDist, imax] = max(fun);
                end
                k = k0;
                z(k) = z(k) + step(k);   % next valid integer      
                y(k) = zb(k) - z(k);
                step(k) = -step(k) - sign(step(k));
            end
            newDist = dist(k) + y(k) * y(k) / d(k);
        end
        ulevel = k;
        
        while newDist >= maxDist   % exit or move to level k+1
            if k == n             
                endSearch = true;
                break;
            end
            k = k + 1;      % move to level k+1
            z(k) = z(k) + step(k);    % next integer
            y(k) = zb(k) - z(k);
            step(k) = -step(k) - sign(step(k));
            newDist = dist(k) + y(k) * y(k) / d(k);
        end
    end
    
    % Sort the solutions by their corresponding residuals
    [fun,p] = sort(fun);
    Optis = Optis(:,p);
end
