function [X,RESNORM,FVAL,EXITFLAG,OUTPUT,LAMBDA,JACOBIAN] = cnls(caller,FUN,X,A,B,Aeq,Beq,LB,UB,NONLCON,options,defaultopt,optimgetFlag)

%Cnls Constrained nonlinear least square problem solver using fmincon IPM
% Cnls solves problems of the form
%   min     norm(f(x))^2
%   s.t.    A * x <= b,
%           Aeq * x = Beq,
%           LB <= x <= UB,
%           g(x) <= 0,
%           geq(x) = 0.
% the problem is solved by fmincon using IPM for now.

%   Copyright 2022 The MathWorks, Inc.

arguments
    caller
    FUN
    X
    A = []
    B = []
    Aeq = []
    Beq = []
    LB = []
    UB = []
    NONLCON = []
    options = []
    defaultopt = []
    optimgetFlag = 'fast'
end

% Ensure FUN is a function handle or cell array of function handles
if ~functionHandleOrCellArrayOfHandles(FUN)
    error(message('optimlib:lsqfcnchk:InvalidFUN'));
end

% Ensure NONLCON is [] or similar to FUN
if ~isempty(NONLCON) && ~functionHandleOrCellArrayOfHandles(NONLCON)
    error(message('optimlib:optimfcnchk:ConstrNotAFunction', 'NONLCON'));
end

% Populate non-empty defaults for these options
defaultopt.MaxFunEvals = 3000;
defaultopt.MaxIter = 1000;
defaultopt.TolX = 1e-10; % interior-point default

% Create an option for Fmincon solver
% Jacobians of augmented objective and constraint are always provided,
% if user doesn't provide Jacobian for FUN/NONLCON, FD is used automatically

n = numel(X);
sizes.nVar = n;
sizes.nFun = [];
sizes.xShape = size(X);
sizes.lin_eq = size(Aeq,1);
sizes.lin_ineq = size(A,1);

display = optimget(options,'Display',defaultopt,optimgetFlag);
switch display
    case {'off','none'}
        flags.verbosity = 0;
    case {'iter','iter-detailed'}
        flags.verbosity = 2;
    case {'final','final-detailed'}
        flags.verbosity = 1;
    case 'testing'
        flags.verbosity = Inf;
    otherwise
        flags.verbosity = 1;
end

OUTPUT.algorithm = 'interior-point';
[XOUT,LB,UB,msg] = checkbounds(X,LB,UB,sizes.nVar);
if ~isempty(msg)

    EXITFLAG = -2;
    X(:) = XOUT; % checkbounds returns a col vector XOUT

    [RESNORM,LAMBDA,JACOBIAN] = deal([]);
    OUTPUT.firstorderopt = [];
    OUTPUT.iterations = 0;
    OUTPUT.funcCount = 0;
    OUTPUT.constrviolation = [];
    OUTPUT.cgiterations = [];    
    OUTPUT.message = msg;
    OUTPUT.bestfeasible = [];

    if flags.verbosity > 0
        disp(msg)
    end
    return
end

% Account for f-count in this function
FcountCNLS = 0;

% Options for gradients
ConstrHaveJac = strcmpi(optimget(options,'GradConstr',defaultopt,optimgetFlag), 'on');
FunHaveJac = strcmpi(optimget(options,'Jacobian',defaultopt,optimgetFlag), 'on');
checkGrad = strcmpi(optimget(options,'DerivativeCheck',defaultopt,optimgetFlag), 'on');
diagnostics = strcmpi(optimget(options,'Diagnostics',defaultopt,optimgetFlag),'on');

% Get options value for finite difference
options.FinDiffType = optimget(options,'FinDiffType',defaultopt,optimgetFlag);
options.UseParallel = optimget(options,'UseParallel',defaultopt,optimgetFlag);
options.DiffMinChange = optimget(options,'DiffMinChange',defaultopt,optimgetFlag);
options.DiffMaxChange = optimget(options,'DiffMaxChange',defaultopt,optimgetFlag);
options = validateFinDiffRelStep(sizes.nVar,options,defaultopt,optimgetFlag);

[typicalx,ME] = getNumericOrStringFieldValue('TypicalX','ones(numberOfVariables,1)', ...
    ones(sizes.nVar,1),'a numeric value',options,defaultopt);
if ~isempty(ME)
    throw(ME)
end
checkoptionsize('TypicalX', size(typicalx), sizes.nVar);
options.TypicalX = typicalx;

% structs finDiffFlags and sizes needed to call computeFinDiffGradAndJac
finDiffFlags = struct('chkFunEval', true,...
    'fwdFinDiff', strcmpi(options.FinDiffType,'forward'),...
    'isGrad', false,...
    'hasLBs', isfinite(LB),...
    'hasUBs', isfinite(UB),...
    'scaleObjConstr', false,...
    'chkComplexObj', false);

flags.grad = FunHaveJac;
flags.hessflag = false;
flags.constflag = ~isempty(NONLCON);
flags.gradconst = ConstrHaveJac;

if checkGrad
    % FUN and NONLCON is modified later to account for grad.
    % USER_FUN & USER_NONLCON are used only for derivative checks.
    USER_FUN = FUN;
    USER_NONLCON = NONLCON;
end

% Initialize persistent variables. 
xOld = []; % Need this for BFGS updates
xCur = X(:);
Fvec = [];
FvecOld = []; % Need this for BFGS updates/init
JACOBIAN = []; % The output from this function gets the last Jacobian
JacOld = []; % Need this for BFGS updates
BFGSHess = []; % Need this for BFGS updates

if flags.constflag
    % Check if CNLS need to handle fun_then_grad case (no checks)
    if iscell(NONLCON)
        CONGRAD = NONLCON{2};
        NONLCON = NONLCON{1};
        con_then_grad = true;
    else
        con_then_grad = false;
    end
    % fmincon is always called with fun_then_grad syntax
    confun = @(x) nonlconExt(x, false);
    congrad = @(x) nonlconExt(x, true);
    fminconConstr = {confun, congrad};
    
    sizes.mNonlinEq = [];
    sizes.mNonlinIneq = [];
    JacCineqTrans = []; % Need this for FD to work (shape, memory)
    JacCeqTrans = []; % Need this for FD to work (shape, memory)
    JacCineqTransOld = JacCineqTrans; % Need this for BFGS updates
    JacCeqTransOld = JacCeqTrans; % Need this for BFGS updates
    
else
    fminconConstr = [];
end

% get plots/outputs functions and fmincon options
[outputfcn, plotfcns, fminconOptions] = getFminconOptions();

% @todo:
% Cannot handle EXITFLAG -3/5
% Lambda is inaccurate when failed to converge, should we do something?

if iscell(FUN)
    FUNGRAD = FUN{2};
    FUN = FUN{1};
    fun_then_grad = true;
else
    fun_then_grad = false;
end
% Second argument (if true)  will return grad
fun = @(x) objective(x, false);
fungrad = @(x) objective(x, true);

% For parallel finite difference we need to send the function
% handles to the workers now. This avoids sending the function handles in
% every iteration of the solver. The output from 'setOptimFcnHandleOnWorkers' 
% is a onCleanup object that will perform cleanup task on the workers.
UseParallel = optimget(options,'UseParallel',defaultopt,optimgetFlag);
ProblemdefOptions = optimget(options, 'ProblemdefOptions',defaultopt,optimgetFlag);
if ~isempty(ProblemdefOptions) && isfield(ProblemdefOptions, 'FromSolve')
    FromSolve = ProblemdefOptions.FromSolve;
else
    FromSolve = false;
end
if ConstrHaveJac
    worker_nonlcon = {''};
else
    worker_nonlcon = {'fun',[], NONLCON};
end
if FunHaveJac
    worker_fun = {''};
else
    worker_fun = {'fun',[],FUN};
end
cleanupFunctionOnWorkers = setOptimFcnHandleOnWorkers(UseParallel, ... 
    worker_fun,worker_nonlcon,FromSolve);

try
    [X,~,EXITFLAG,OUTPUT,LAMBDA] = fmincon({fun,fungrad},X(:),...
        A,B,Aeq,Beq,LB,UB,fminconConstr,fminconOptions);
catch userFcn_ME
    if strcmp(userFcn_ME.identifier,'optimlib:validateFirstDerivatives:InvalidGrad')
        X  = [];
    else
        optim_ME = MException(userFcn_ME.identifier, ...
            getString(message('optimlib:lsqncommon:FminconError', caller)));
        userFcn_ME = addCause(userFcn_ME,optim_ME);
        rethrow(userFcn_ME)
    end
end

if ~isempty(X)
    X = reshape(X,sizes.xShape);
    FVAL = FUN(X);
    RESNORM = norm(FVAL)^2;
else
    RESNORM = []; 
    FVAL = [];
    EXITFLAG = [];
    LAMBDA = [];
    JACOBIAN = [];
end
% Replace 'fmincon' with caller in the message string
if isstruct(OUTPUT) && isfield(OUTPUT,'message') && ~isempty(OUTPUT.message)
    OUTPUT.message = strrep(OUTPUT.message,'fmincon',caller);
end

% Force a cleanup of the handle object. Sometimes, MATLAB may
% delay the cleanup but we want to be sure it is cleaned up.
% (Don't delete empties to avoid loading graphics libs.)
if ~isempty(cleanupFunctionOnWorkers)
    delete(cleanupFunctionOnWorkers);
end

%  Setup output/plot
    function [outputfcn, plotfcns, fminconOptions] = getFminconOptions()
        % lsq* output and plot functions should get resnorm and residuals so must have a
        % wrapper for fmincon.
        % Handle the output
        outputfcn = optimget(options,'OutputFcn',defaultopt,optimgetFlag);
        if isempty(outputfcn)
            haveoutputfcn = false;
        else
            haveoutputfcn = true;
            % Parse OutputFcn which is needed to support cell array syntax for OutputFcn.
            outputfcn = createCellArrayOfFunctions(outputfcn,'OutputFcn');
        end

        % Handle the output
        plotfcns = optimget(options,'PlotFcns',defaultopt,optimgetFlag);
        if isempty(plotfcns)
            haveplotfcn = false;
        else
            haveplotfcn = true;
            % Parse PlotFcns which is needed to support cell array syntax for PlotFcns.
            plotfcns = createCellArrayOfFunctions(plotfcns,'PlotFcns');
        end

        % other fmincon options
        maxFunEvals = optimget(options,'MaxFunEvals',defaultopt,optimgetFlag) - FcountCNLS;
        
        fminconOptions = optimoptions('Fmincon', 'Display', display, ...
            'ConstraintTolerance', optimget(options,'TolCon',defaultopt,optimgetFlag), ...            
            'MaxFunctionEvaluations', maxFunEvals, ...
            'MaxIterations', optimget(options,'MaxIter',defaultopt,optimgetFlag), ...
            'OptimalityTolerance', optimget(options,'TolFun',defaultopt,optimgetFlag), ...
            'StepTolerance', optimget(options,'TolX',defaultopt,optimgetFlag), ...
            'ScaleProblem', optimget(options,'ScaleProblem',defaultopt,optimgetFlag), ...
            'InitBarrierParam', optimget(options,'InitBarrierParam',defaultopt,optimgetFlag), ...
            'SubproblemAlgorithm', optimget(options,'SubproblemAlgorithm',defaultopt,optimgetFlag), ...
            'FunValCheck', optimget(options,'FunValCheck',defaultopt,optimgetFlag), ...
            'BarrierParamUpdate', optimget(options,'BarrierParamUpdate',defaultopt,optimgetFlag));

        % fixed options
        fminconOptions.Algorithm = 'interior-point';
        fminconOptions.UseParallel = false;
        fminconOptions.CheckGradients = false;
        fminconOptions.SpecifyConstraintGradient = true;
        fminconOptions.SpecifyObjectiveGradient = true;

        fminconOptions.FunctionEvaluationCounter = @getFunctionEvaluationCount;
        fminconOptions.HessianFcn = @hessianfcnSP;
        if haveplotfcn || haveoutputfcn
            fminconOptions.OutputFcn = @lsqOutputPlots;
        end
    end
    
    function out = getFunctionEvaluationCount()
        out = FcountCNLS;
    end
    
    % Objective function (scalar) for fmincon.
    function out = objective(xIn,gradOut)
        if gradOut
            % Update the old value for Fvec and Jac before calculating new values.
            % The update of old values is done when fmincon needs the gradients. The
            % gradients are needed when the next trial is accepted.
            xOld = xCur;
            xCur = xIn;
            JacOld = JACOBIAN;
            FvecOld  = Fvec;
            
            [f,F] = funExt(xIn);
            out = 2*F' * f;
        else
            f = funExt(xIn);
            out = norm(f)^2;
        end        
    end

    % Wrapper around FUN (vector) to manage Jacobian calculations used for BFGS updates
    function [F,Jac] = funExt(xIn)

        Xk = reshape(xIn,sizes.xShape);
        % The first call can be a different than X0 if fmincon shifted the X0. We can
        % detect it.
        if isempty(xOld) 
            % First call
            xOld = xIn;
            xCur = xIn;
        end

        % If there are constraints, then xIn is already counted for evaluation.
        if ~flags.constflag
            FcountCNLS  = FcountCNLS + 1;
        end
        if nargout >= 2
            % Since cnls uses fun_then_grad with fmincon we know that Xk was
            % evaluated for fun before.
            if ~flags.constflag
                FcountCNLS  = FcountCNLS - 1;
            end
            if FunHaveJac
                if fun_then_grad
                    % We need Fvec to compute the gradient for fmincon
                    Fvec = feval(FUN, Xk);
                    JACOBIAN = feval(FUNGRAD, Xk);
                else
                    [Fvec,JACOBIAN] = feval(FUN, Xk);                    
                end
            else
                Fvec = feval(FUN, Xk);
                if isempty(JACOBIAN)
                    % First evaluation, initialize Jac (for FD)
                    sizes.nFun = numel(Fvec);
                    JACOBIAN = zeros(sizes.nFun, sizes.nVar);
                end
                % Finite difference
                [JACOBIAN,~,~,numEvals] = computeFinDiffGradAndJac(xIn,FUN,[],Fvec,[],[], ...
                    JACOBIAN,[],[],LB,UB,[],options,finDiffFlags,sizes);
                % The points for FD at xIn are same as sent for CONSTRAINT JACOBIAN
                % calculation. If both OBJ and CONSTR are using FD then we should not
                % double count the points.                
                if flags.constflag && ~ConstrHaveJac
                    numEvals = 0;
                end                
                FcountCNLS  = FcountCNLS + numEvals;
            end
            Jac = JACOBIAN;
            F = Fvec(:);
        else
            % For one output arg case, do not update the old value.
            F = feval(FUN, Xk);
            if isempty(sizes.nFun)
                sizes.nFun = numel(F);
            end
        end        
        % Diagnostics (if flags.constflag then do it from nonlconExt())
        if diagnostics && ~flags.constflag
            showDiagnostics();
        end
        if checkGrad && ~flags.constflag            
            checkDerivatives()
        end
    end

    % Wrapper around NONLCON to manage Jacobian calculations
    function [cinOrJac,ceqOrJac] = nonlconExt(xIn, gradOut)
        % The first column of the gradient matrix is associated with c(1), the
        % second column is associated with c(2) and so on. This derivative form is the
        % transpose of the form of Jacobians.
        %
        % Note that the objective returns Jacobian (as required by lsq*) but
        % nonlinear constraints returns the transpose of the Jacobian (similar to fmincon)

        Xk = reshape(xIn,sizes.xShape);
        
        % Account for evaluation at xIn
        FcountCNLS  = FcountCNLS + 1;
        if gradOut
            % Store the old values for BFGS update
            JacCineqTransOld = JacCineqTrans;
            JacCeqTransOld = JacCeqTrans;
            % Since cnls uses fun_then_grad we know that Xk was
            % evaluated for fun before.
            FcountCNLS  = FcountCNLS - 1;

            if ConstrHaveJac
                if con_then_grad
                    [JacCineqTrans,JacCeqTrans] = feval(CONGRAD, Xk);                    
                else
                    [~,~,JacCineqTrans,JacCeqTrans] = feval(NONLCON, Xk);                    
                end
            else
                [cin,ceq] = feval(NONLCON, Xk);
                % Finite difference
                [~,JacCineqTrans,JacCeqTrans,numEvals,~] = computeFinDiffGradAndJac(xIn,[],NONLCON,[], ...
                    cin(:),ceq(:),[],JacCineqTrans,JacCeqTrans,LB,UB,[],options,finDiffFlags,sizes);
                % FD evaluation is always counted for CONSTR
                FcountCNLS  = FcountCNLS + numEvals;
            end
            % Output (constr or JacT)
            cinOrJac = JacCineqTrans; ceqOrJac = JacCeqTrans;
        else
            [cinOrJac,ceqOrJac] = feval(NONLCON, Xk);
            % Initialize one time
            if isempty(sizes.mNonlinEq)
                % First evaluation, initialize Jac and sizes (for FD)
                sizes.mNonlinEq = numel(ceqOrJac);
                sizes.mNonlinIneq = numel(cinOrJac);
                JacCineqTrans = zeros(sizes.nVar, sizes.mNonlinIneq); % GTrans
                JacCeqTrans = zeros(sizes.nVar, sizes.mNonlinEq); % GeqTrans
                if diagnostics
                    showDiagnostics();
                end
                if checkGrad
                    checkDerivatives()
                end
            end
        end        
    end

    function Hout = hessianfcnSP(xIn,lambda)
        % LSQ-BFGS update
        if isempty(BFGSHess)
            BFGSHess = eye(sizes.nVar);
            Hout = BFGSHess + 2*(JACOBIAN' * JACOBIAN); % Initial quasi-Newton matrix
            return;
        end
        
        % Hessian of objective        
        delta_x = xIn - xOld;
        %delta_gradLag = 2*(JACOBIAN - JacOld)' * (JacOld * delta_x + FvecOld(:)); Another option to update quasi-Newton matrix
        delta_gradLag = 2*(JACOBIAN'*Fvec(:) - JacOld'*FvecOld(:)) - 2*JACOBIAN'*(JACOBIAN*delta_x);

        if flags.constflag
            % JacCineqTrans and JacCeqTrans are for xIn because we are using
            % fun_then_grad syntax. No need to evaluate again.
            if ~isempty(JacCeqTrans)
                delta_gradLag = delta_gradLag + (JacCeqTrans - JacCeqTransOld) * lambda.eqnonlin;
            end
            if ~isempty(JacCineqTrans)
                delta_gradLag = delta_gradLag + (JacCineqTrans- JacCineqTransOld) * lambda.ineqnonlin;
            end
        end        

        HessS_k = BFGSHess * delta_x;
        s_kTransHessS_k = delta_x' * HessS_k;
        y_kTransS_k = delta_gradLag' * delta_x;

        if s_kTransHessS_k > eps

            if  y_kTransS_k < 0.2*s_kTransHessS_k
                % Note: demonimator below cannot be zero because of condition in if-statement
                theta = 0.8*s_kTransHessS_k / (s_kTransHessS_k - y_kTransS_k);
                delta_gradLag = theta*delta_gradLag + (1.0 - theta)*HessS_k;
                % delta_gradLag updated: also update y_kTransS_k
                y_kTransS_k = delta_x'*delta_gradLag;
            end

            if y_kTransS_k > eps
                BFGSHess = BFGSHess - (HessS_k * HessS_k') / s_kTransHessS_k + ...
                    (delta_gradLag * delta_gradLag') / y_kTransS_k; % BFGS use F_old, G_old, Geq_old, x_old
            end
        end

        % Hessian
        Hout = BFGSHess + 2*(JACOBIAN' * JACOBIAN);        
    end

% Output and Plot fuction wrapper for lsq to provide residual and resnorm because
% fmincon does not have these quanities.
    function stop = lsqOutputPlots(xOutputfcn, optimValues, state)

        stop = false;        
        optimValues.residual = Fvec;
        optimValues.resnorm = norm(Fvec)^2;
        optimValues = rmfield(optimValues,'fval'); % fval is not expected in lsq
        
        % Call output function
        if ~isempty(outputfcn)
            switch state
                case {'iter','init'}
                    stop = callAllOptimOutputFcns(outputfcn,xOutputfcn,optimValues,state) || stop;
                case 'done'
                    callAllOptimOutputFcns(outputfcn,xOutputfcn,optimValues,state);
            end
        end
        % Call plot functions
        if ~isempty(plotfcns)
            switch state
                case {'iter','init'}
                    stop = callAllOptimPlotFcns(plotfcns,xOutputfcn,optimValues,state) || stop;
                case 'done'
                    callAllOptimPlotFcns(plotfcns,xOutputfcn,optimValues,state);
            end
        end

    end
    
% Derivative check
    function checkDerivatives()

        if flags.grad || flags.gradconst            
            options.GradObj = optimget(options,'Jacobian',defaultopt,optimgetFlag);
            funValCheck = strcmp(optimget(options,'FunValCheck',defaultopt,optimgetFlag),'on');
            funfcn_ = lsqfcnchk(USER_FUN,caller,0,funValCheck,flags.grad);
            if flags.constflag % NONLCON is non-empty
                confcn_ = optimfcnchk(USER_NONLCON,caller,0,funValCheck,flags.gradconst,false,true);
            else
                confcn_ = {'','','','',''};
            end

            validateFirstDerivatives(funfcn_,confcn_,X,LB,UB,options,finDiffFlags,sizes);
        end
        checkGrad = false;
    end

% Diagnostics
    function showDiagnostics()

        funValCheck = strcmp(optimget(options,'FunValCheck',defaultopt,optimgetFlag),'on');
        funfcn = lsqfcnchk(FUN,caller,0,funValCheck,flags.grad);        
        if flags.constflag % NONLCON is non-empty
            confcn = optimfcnchk(NONLCON,caller,0,funValCheck,flags.gradconst,false,true);
        else
            confcn = {'','','','',''};
            sizes.mNonlinEq = 0;
            sizes.mNonlinIneq = 0;
        end
        
        diagnose(caller,OUTPUT,flags.grad,flags.hessflag,flags.constflag,flags.gradconst,...
            X(:),sizes.mNonlinEq,sizes.mNonlinIneq,sizes.lin_eq,sizes.lin_ineq,LB,UB,funfcn,confcn);
        diagnostics = false;
    end

end

function TF = functionHandleOrCellArrayOfHandles(value)
TF = isa(value,'function_handle') || ...
    (iscell(value) && length(value) == 2 && ...
    isa(value{1},'function_handle') && ...
    isa(value{2},'function_handle'));
end
