function [A,chi2,CM,AT,CHI2,DCHI2,DA] = levmarq(FUN,X0,Y0,A0,mfit,da,tolfun,maxiter,maxsamenum)
% LEVMARQ   An implementation of the Levenberg-Marquardt method.
%   LEVMARQ(FUN,X0,Y0,A0,mfit,da,tolfun,maxiter,maxsamenum) implements the
%   Levenberg-Marquardt method for non-linear least-squares function
%   fitting, as transcribed from "Numerical Recipes in C", p. 542-547.
% 
% INPUTS:   FUN         A handle to the non-linear function
%           X0          The measured X points (N x 1)
%           Y0          The measured Y points (N x P)
%           A0          Initial coefficients for the non-linear function
%           mfit        The number of coefficients in A0 to adjust
%           da          The step size to try when adjusting A0
%           tolfun      The threshold for the stopping criteria
%           maxiter     The maximum number of iterations to try
%           maxsamenum  The maximum number of same chi^2 values to get 
%                         before we give up
%
% OUTPUTS:  A           The optimized parameters
%           chi2        The value of the objective function
% 
% Written by:       Jason Cardema
% Last modified:    06 September 2006


%% STEP 1: Initialization
N = length(X0);     % Data length
P = size(Y0,2);     % Number of output variables
M = length(A0);     % Parameter length
A = A0(:);          % Set the parameters initially to A0
lambda = 1e-3;      % Pick a modest value for lambda as the initial guess
chi2 = chi2_fun(FUN,A,X0,Y0);   % Initial value of the chi^2 function
DECREASED_FLAG = 0;

CHI2 = zeros(1,maxiter);    % Save all the values of CHI2
AT = zeros(mfit,maxiter);   % Save all the values of A_temp
LAMBDA = zeros(1,maxiter);  % Save all the values of lambda
DCHI2 = zeros(mfit,maxiter);    % Save all the values of DCHI2DA
DA = zeros(mfit,maxiter);       % Save all the values of delta_a
CHI2(1) = chi2;
AT(:,1) = A;
LAMBDA(1) = lambda;

% Create the alpha and beta matrices
alpha = zeros(mfit,mfit);   % Initialize the alpha matrix
beta = zeros(mfit,1);       % Initialize the beta matrix

% Keep looping until the stopping criteria is met
iternum = 0;
same_count = 0;
while (chi2 > tolfun) && (iternum < maxiter) && (same_count < maxsamenum)
    
    iternum = iternum + 1;  % Keep track of the iteration number
    
    % Find the partial derivatives of chi2 wrt the different a's
    A_last = zeros(mfit,1);
    A_next = zeros(mfit,1);
    dchi2da = zeros(mfit,1);
    for jj = 1:mfit
        A_last = A;
        A_last(jj) = A(jj) - da;
        chi2_last(jj) = chi2_fun(FUN,A_last,X0,Y0);
        A_next = A;
        A_next(jj) = A(jj) + da;
        chi2_next(jj) = chi2_fun(FUN,A_next,X0,Y0);
        dchi2da(jj) = (chi2_next(jj) - chi2_last(jj)) / (2*da);
        
    end
    
    % Update the values of alpha and beta
    alpha = dchi2da*dchi2da';
    beta = -dchi2da;
    
    %% STEP 2: Solve for delta_a
    % Alter the linearized fitting matrix by augmenting diagonal elements
    alpha_d = alpha + lambda*eye(mfit);
    if rank(alpha_d) == mfit
        delta_a = alpha_d \ beta;
    else
        fprintf('The alpha_d matrix is rank-deficient.  Adjusting...\n');
        if alpha_d(1,1) > alpha_d(2,2)
            delta_a(1,1) = beta(1)/alpha(1,1);
            delta_a(2,1) = 0;
        else
            delta_a(1,1) = 0;
            delta_a(2,1) = beta(2)/alpha(2,2);
        end
    end

    
    %% STEP 3: Check the value of chi^2 to see if we've improved
    chi2_old = chi2;                    % Save the old value of chi^2
    A_temp = A + delta_a;               % Create the trial matrix A_temp
    chi2 = chi2_fun(FUN,A_temp,X0,Y0);  % Calculate new chi^2 value
    
    CHI2(iternum+1) = chi2;
    AT(:,iternum+1) = A_temp;
    LAMBDA(iternum+1) = lambda;
    DCHI2(:,iternum+1) = dchi2da;
    DA(:,iternum+1) = delta_a;

    % Check if we've reached the same chi^2 value too many times in a row
    if abs(chi2 - chi2_old) < 1e-12
        same_count = same_count + 1;
    else
        same_count = 0;         % Reset the value of the same_count
    end

    % Compare the old and new chi^2 values
    if (chi2 <= chi2_old) && (DECREASED_FLAG == 0)
        % If the new chi^2 is less, we're moving in the right direction
        lambda = lambda*0.10;   % Decrease lambda (use Newton method more)
        A = A_temp;             % Update the value of A
        DECREASED_FLAG = 1;
    elseif chi2 > chi2_old
        % If we're not improving, increase lambda
        lambda = lambda*10; % Increase lambda (use steepest descent more)
        chi2 = chi2_old;    % Reset the value of chi^2
        DECREASED_FLAG = 0;
    else
        A = A_temp;
    end
    
end % Repeat the while loop until the value of chi^2 is less than tolfun

% Compute the covariance matrix
CM = alpha;

if (iternum >= maxiter)
    fprintf('STOPPED: Maximum number of iterations reached!\n');
elseif (same_count >= maxsamenum)
    fprintf('STOPPED: Got the same chi^2 value %d times in a row!\n',same_count);
end

CHI2 = CHI2(1:iternum+1);
AT = AT(:,1:iternum+1);
LAMBDA = LAMBDA(1:iternum+1);
DCHI2 = DCHI2(:,1:iternum+1);
DA = DA(:,1:iternum+1);

% figure
% plot(CHI2,'o-')
% title(sprintf('Chi^2 values, step size of %f',da))
% xlabel('Iteration number')
% ylabel('Chi^2')
% 
% figure
% plot(AT','o-')
% title(sprintf('A values, step size of %f',da))
% xlabel('Iteration number')
% ylabel('A')
% 
% figure
% plot(DA','o-')
% title(sprintf('DA values, step size of %f',da))
% xlabel('Iteration number')
% ylabel('DA')
% 
% figure
% plot(log10(LAMBDA),'o-')
% title(sprintf('log_{10}(lambda), step size of %f',da))
% xlabel('Iteration number')
% ylabel('lambda')



function F = chi2_fun(FUN,A,X,Y)
F = sum(sum((FUN(A,X) - Y).^2));