function varargout = FitObjective(m, con, obj, opts)
%FITOBJECTIVE Optimizes the parameters of a model to minimize a set of
%   objective functions.
%
%   [m, con, G, D] = FitObjective(m, con, obj, opts)
%
%   Inputs:
%       m    - Kronecker Bio model
%       con  - An vector of experimental conditions under which the model
%              will be simulated
%       obj  - A matrix length of con by any number of objective functions
%              that will be evaluated
%       opts - Scalar structure of options
%           UseParams - A vector of linear indexes or a vector of logicals
%                       length of nP indicating which rate parameters are
%                       active for fitting. Default = 1:nP
%           UseICs - If UseModelICs is true then UseICs can be a vector of
%                    linear indexes or a vector of logicals length of nX.
%                    If UseModelICs is false then UseICs can be a matrix of
%                    logicals size nX by nCon. It can also be a vector of
%                    length nX, and every experiment will be considered to
%                    have the same active IC parameters. It can also be a
%                    vector of linear indexes into the nX vector and
%                    assumed the same for all conditions.
%                    Default = []
%           UseModelICs - A scalar logical indicating that the model IC
%                         parameters should be used for simulation and as
%                         fitting parameters. Default = true
%           LowerBound - A vector for the lower bounds on the fitting
%                        parameters. It can be length nP+nCon*nX, nP+nX,
%                        nV, or just nP if nVX = 0. The bounds will be
%                        interpreted in that order if the length matches
%                        multiple orders. Default = zeros(nP*nX,1)
%           UpperBound - A vector for the upper bounds on the fitting
%                        parameters. It must be the same length and type as
%                        the lower bound. Default = zeros(nP*nX,1)
%           ObjWeights - A matrix of size(obj) indicating a post evaluation
%                        weight on each objective function in terms of how
%                        much it will contribute to the final value.
%                        Default = ones(size(obj))
%           UseAdjoint - A scalar logical indicating whether the
%                        sensitivites should be calculated via the adjoint
%                        method or the forward method. Default = true
%           TolOptim - A scalar indicating the objective tolerance. The
%                      optimization stops when it is predicted that the
%                      objective function cannot be improved more than this
%                      in the next iteration. Default = 1e-5
%           RelTol - Numerical scalar for the relative tolerance of the
%                    integration. Default = 1e-6
%           AbsTol - Cell vector of vectors, numerical scalar, or numerical
%                    vector for the absolute tolerance of the integration.
%                    Default = 1e-9 
%           AdaptAbsTol - A scalar logical determining if the absolute
%                         tolerance should adaptively optimize in order to
%                         most efficiently compute the objective and
%                         gradient. Default = true
%           Normalized - A scalar logical determining if the optimization
%                        should be done in log parameters space.
%                        Default = true
%           Restart - A scalar integer determining how many times the
%                     optimzation should restart once optimization has
%                     stopped. Default = 0
%           RestartJump - A function handle @(iter, G) determining the
%                         schedule for the noise that will be added to the
%                         parameters before each restart. The parameters
%                         for the next iteration will be normally
%                         distributed in log space with a mean equal to the
%                         previous iteration and a standard deviation equal
%                         to the value returned by this function. The value
%                         returned should usually be a scalar, but it can
%                         also be a vector with length equal to the number
%                         of active parameters. It can also be numeric, and
%                         the noise will be treated as this constant value.
%                         Default = @(iter,G)(0.001), a constant noise
%                         equal to 0.1%.
%           TerminalGoal - Optimization is halted when this objective
%                          function value is reached. Default = -inf
%           MaxStepSize - Scalar fraction indicator of the maximum relative 
%                         step size that any parameter can take in a single
%                         interation. Default = 1
%           Algorithm - fmincon string option indicating which algorithm to
%                       use. Default = 'active-set'
%           MaxIter - fmincon scalar option. Maximum number of iterations
%                     allowed before optimization will be terminated.
%                     Default = 1000
%           MaxFunEvals - fmincon scalar option. Maximum number of
%                         objective function evaluations allowed before
%                         optimization will be terminated. Default = 5000
%           Verbose - Scalar depth of verbosity. Higher numbers lead to
%                     more detailed print out.
%
%   Outputs
%       m - The model with all the optimum rate parameters applied, as
%           well as the IC parameters, if UseModelICs = true.
%       con - The experimental conditions which will have the optimum IC
%             parameters applied, if UseModelICs = false.
%       G - The optimum objective function value
%       D - The gradient at the optimum parameter set
%
%   FitObjective uses the derivatives in the Kronecker model and in the
%   objective functions to build a function that can not only evaluate the
%   objective functions at particular parameter sets, but also evaluate the
%   gradient at those parameter sets, thereby pointing in the direction of
%   a more optimum parameter set. This function is built around Matlab's
%   fmincon, which is a gradient descent minimizer. It varies the
%   parameters attempting to find the parameter set that will minimize the
%   objective function while keeping with the bounds.

% (c) 2010 David R Hagen & Bruce Tidor
% This work is released under the MIT license.

%TODO: eliminate the need for the jSol vector being passed to objective
%functions if possible

% Clean up inputs
if nargin < 4
    opts = [];
end

assert(isscalar(m), 'KroneckerBio:FitObjective:MoreThanOneModel', 'The model structure must be scalar')

%% Options
% Default options
defaultOpts.UseParams               = 1:m.nP;
defaultOpts.UseICs                  = [];
defaultOpts.UseModelICs             = true;
defaultOpts.LowerBound              = zeros(m.nP+m.nX, 1);
defaultOpts.UpperBound              = zeros(m.nP+m.nX, 1) + inf;
defaultOpts.ObjWeights              = ones(size(obj));
defaultOpts.UseAdjoint              = true;

defaultOpts.TolOptim                = 1e-5;
defaultOpts.AbsTol                  = NaN; % in fixAbsTol
defaultOpts.RelTol                  = NaN; % 1e-6
defaultOpts.AdaptAbsTol             = true;
defaultOpts.Normalized              = true;
defaultOpts.Restart                 = 0;
defaultOpts.RestartJump             = 0.001;
defaultOpts.TerminalGoal            = -inf;

defaultOpts.MaxStepSize             = 1;
defaultOpts.Algorithm               = 'active-set';
defaultOpts.MaxIter                 = 1000;
defaultOpts.MaxFunEvals             = 5000;

defaultOpts.Verbose                 = 1;

% Combine user options with default options
opts = mergestruct(defaultOpts, opts);

% Constants
nX = m.nX;
nP = m.nP;
nCon = numel(con);
nObj = size(obj,2);

% Ensure UseRates is column vector of logical indexes
[opts.UseParams, nVP] = fixUseParams(opts.UseParams, nP);

% Ensure UseICs is a matrix of logical indexes
[opts.UseICs, nVX] = fixUseICs(opts.UseICs, opts.UseModelICs, nX, nCon);

nV = nVX + nVP;

% Store complete parameter sets
pAll = m.p;

if opts.UseModelICs
    icAll = m.ic;
else
    icAll = zeros(nX, nCon);
    for i = 1:nCon
        icAll(:,i) = con(i).ic;
    end
end

% Construct starting variable parameter set
v0 = [pAll(opts.UseParams); vec(icAll(opts.UseICs))];

% Ensure Restart is a positive integer
if ~(opts.Restart >= 0)
    opts.Restart = 0;
    warning('KroneckerBio:FitObjective', 'opts.Restart was not nonegative. It has been set to 0.')
end

if ~(opts.Restart == floor(opts.Restart))
    opts.Restart = floor(opts.Restart);
    warning('KroneckerBio:FitObjective', 'opts.Restart was not a whole number. It has been floored.')
end

% Ensure RestartJump is a function handle
if isnumeric(opts.RestartJump)
    opts.RestartJump = @(iter,G)(opts.RestartJump);
end

%% Standardize con and obj
con = pastestruct(Uzero(m), con);
obj = pastestruct(Gzero(m), obj);
for iCon = 1:nCon
    con(iCon) = pastestruct(Uzero(m), con(iCon).update(m, con(iCon).ic));
    for iObj = 1:nObj
        obj(iCon,iObj) = pastestruct(Gzero(m), obj(iCon,iObj).update(m));
    end
end

%% Integration type: simple, continuous, complex, or both
[opts.continuous, opts.complex, opts.tGet] = fixIntegrationType(obj);

%% Tolerances
if isempty(opts.RelTol) || isnan(opts.RelTol)
    opts.RelTol = 1e-6;
end

% Fix AbsTol to be a cell array of vectors appropriate to the problem
opts.AbsTol = fixAbsTol(opts.AbsTol, 2, opts.continuous, nX, nCon, opts.UseAdjoint, opts.UseParams, opts.UseICs, opts.UseModelICs);

%% Options structure for integration
intOpts = opts;
intOpts.Verbose = max(intOpts.Verbose-1, 0);

%% Minimization options
minOpts.Algorithm               = opts.Algorithm;
minOpts.TolFun                  = opts.TolOptim;
minOpts.TolX                    = 0;
minOpts.OutputFcn               = @isTerminalGoal;
minOpts.GradObj                 = 'on';
minOpts.Hessian                 = 'off'; % unused
minOpts.HessMult                = []; % unused
minOpts.MaxFunEvals             = opts.MaxFunEvals;
minOpts.MaxIter                 = opts.MaxIter;
minOpts.RelLineSrchBnd          = opts.MaxStepSize;
minOpts.RelLineSrchBndDuration  = inf;
minOpts.TolCon                  = 1e-6;

if opts.Verbose
    minOpts.Display = 'iter';
else
    minOpts.Display = 'off';
end

% Optimization bounds
if opts.UseModelICs
    % Bounds can be (nP+nX), nV, or nP long
    assert(numel(opts.LowerBound) == numel(opts.UpperBound))
    l = numel(opts.LowerBound);
    
    if l == nP+nX
        optimLower = opts.LowerBound([opts.UseParams; opts.UseICs]);
        optimUpper = opts.UpperBound([opts.UseParams; opts.UseICs]);
    elseif l == nV
        optimLower = opts.LowerBound;
        optimUpper = opts.UpperBound;
    elseif nVX == 0 && l == nP
        optimLower = opts.LowerBound(opts.UseParams);
        optimUpper = opts.UpperBound(opts.UseParams);
    elseif isscalar(l)
        optimLower = zeros(nV,1) + opts.LowerBound;
        optimUpper = zeros(nV,1) + opts.UpperBound;
    else
        error('KroneckerBio:FitObjective:BoundSize', ...
            'LowerBound and UpperBound must be vectors the length of m.nP+m.nX, number of varible parameters, m.nP if there are no variable ICs, or scalar')
    end
else
    % Bounds can be nP+(nX*nCon), nP+nX, nV, or nP long
    assert(numel(opts.LowerBound) == numel(opts.UpperBound))
    l = numel(opts.LowerBound);

    if l == nP+(nX*nCon)
        optimLower = opts.LowerBound([opts.UseParams; vec(opts.UseICs)]);
        optimUpper = opts.UpperBound([opts.UseParams; vec(opts.UseICs)]);
    elseif l == nP+nX
        optimLower = opts.LowerBound(nP+1:nP+nX); % Extract nX vector
        optimLower = vec(kron(optimLower,ones(1,nCon))); % Expand it
        optimLower = optimLower(vec(opts.UseICs)); % Extract active IC parameters
        optimLower = [opts.LowerBound(opts.UseParams); optimLower]; % Concatenate active rate parameters
        optimUpper = opts.UpperBound(nP+1:nP+nX); % Extract nX vector
        optimUpper = vec(kron(optimUpper,ones(1,nCon))); % Expand it
        optimUpper = optimUpper(vec(opts.UseICs)); % Extract active IC parameters
        optimUpper = [opts.UpperBound(opts.UseParams); optimUpper]; % Concatenate active rate parameters
    elseif l == nV
        optimLower = opts.LowerBound;
        optimUpper = opts.UpperBound;
    elseif nVX == 0 && l == nP
        optimLower = opts.LowerBound(opts.UseParams);
        optimUpper = opts.UpperBound(opts.UseParams);
    elseif isscalar(l)
        optimLower = zeros(nV,1) + opts.LowerBound;
        optimUpper = zeros(nV,1) + opts.UpperBound;
    else
        error('KroneckerBio:FitObjective:BoundSize', ...
            'LowerBound and UpperBound must be vectors the length of m.nP+(m.nX*length(con)), m.nP+m.nX, number of varible parameters, m.nP if there are no variable ICs, or scalar')
    end
end

%% Normalize parameters
if opts.Normalized
    % Normalize starting parameters and bounds
    v0 = log(v0);
    optimLower = log(optimLower);
    optimUpper = log(optimUpper);
    
    % Change relative line search bound to an absolute scale in log space
    % Because fmincon lacks an absolute option, this hack circumvents that
    minOpts.TypicalX = zeros(nV,1) + log(1 + opts.MaxStepSize)*log(realmax);
    minOpts.RelLineSrchBnd = 1 / log(realmax);
end

%% Apply bounds to starting parameters before optimizing
% fmincon will choose a wierd value if a starting parameter is outside the bounds
v0(v0 < optimLower) = optimLower(v0 < optimLower);
v0(v0 > optimUpper) = optimUpper(v0 > optimUpper);

%% Run optimization
% Initialize loop
vhat = v0;
Gbest = inf;
vbest = v0;

for iRestart = 1:opts.Restart+1
    % Init abort parameters
    aborted = false;
    vAbort = vhat;
    
    if opts.Verbose; fprintf('Beginning gradient descent...\n'); end
    [vhat, G, unused, unused, unused, D]  = fmincon(@objective, vhat, [], [], [], [], optimLower, optimUpper, [], minOpts);
    
    % Check abortion status
    % Abortion values are not returned by fmincon and must be retrieved
    if aborted
        vhat = vAbort;
        [G, D] = objective(vhat);
    end
    
    % Re-apply stiff bounds
    vhat(vhat < optimLower) = optimLower(vhat < optimLower);
    vhat(vhat > optimUpper) = optimUpper(vhat > optimUpper);
    
    % See if these parameters are better than previous iterations
    if G < Gbest
        Gbest = G;
        vbest = vhat;
    end

    % Terminate if goal has been met
    if G <= opts.TerminalGoal
        break
    end
    
    % Jump parameters before restarting
    % Retain the deterministic nature of fitting by fixing the stream
    defaultStream = RandStream.getDefaultStream;
    prodvhat = prod(vhat); % model fingerprint
    RandStream.setDefaultStream(RandStream.create('mt19937ar', 'seed', mod(prodvhat/eps(prodvhat)*iRestart,2^32)));
    if opts.Normalized
        vhat = vbest + randn(nV,1) .* vec(opts.RestartJump(iRestart,G));
    else
        vhat = exp(log(vbest) + randn(nV,1) .* vec(opts.RestartJump(iRestart,G)));
    end
    RandStream.setDefaultStream(defaultStream);
    
    % Prevent jumps from leaving bounds
    while any(vhat < optimLower) || any(vhat > optimUpper)
        vhat(vhat < optimLower) = 2*(optimLower(vhat < optimLower)) - vhat(vhat < optimLower);
        vhat(vhat > optimUpper) = 2*(optimUpper(vhat > optimUpper)) - vhat(vhat > optimUpper);
    end
end

% Unnormalize
if opts.Normalized
    vbest = exp(vbest);
    D = D ./ vbest;
end

% Update parameter sets
pAll(opts.UseParams) = vbest(1:nVP);
icAll(opts.UseICs) = vbest(nVP+1:nV);

if opts.UseModelICs
    % Update the model
    m = m.update(pAll, icAll);
else
    % Update the model
    m = m.update(pAll, m.ic);
    % Update the conditions
    temp = con;
    con = constructExperiment(size(con));
    for i = 1:nCon
        con(i) = temp(i).update(m, icAll(:,i));
    end
end

varargout{1} = m;
varargout{2} = con;

if nargout >= 3
    varargout{3} = G;
    if nargout >= 4
        varargout{4} = D;
    end
end

% End of function
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%% Objective function for fmincon %%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [G, D] = objective(v)
        % Reset answers
        G = 0;
        D = zeros(nV,1);
        
        % Unnormalize
        if opts.Normalized
            v = exp(v);
        else
            % If fmincon chooses negative values, force them to be zero
            v(v < 0) = 0;
        end
        
        % Update parameter sets
        pAll(opts.UseParams) = v(1:nVP);
        icAll(opts.UseICs) = v(nVP+1:nV);
        
        newCon = Uzero(nCon);
        if opts.UseModelICs
            m = m.update(pAll, icAll);
            for iCon = 1:nCon
                newCon(iCon) = pastestruct(Uzero(m), con(iCon).update(m));
            end
        else
            m = m.update(pAll, m.ic);
            for iCon = 1:nCon
                newCon(iCon) = pastestruct(Uzero(m), con(iCon).update(m, icAll(:,iCon)));
            end
        end
        con = newCon;
        
        newObj = Gzero([nCon, nObj]);
        for iCon = 1:nCon
            for iObj = 1:nObj
                newObj(iCon,iObj) = pastestruct(Gzero(m), obj(iCon,iObj).update(m));
            end
        end
        obj = newObj;
        
        % Integrate system to get objective function value
        if nargout == 1
            G = computeObj(m, con, obj, intOpts);
        end
        
        % Integrate sensitivities or use adjoint to get objective gradient
        if nargout == 2
            if opts.UseAdjoint
                [G, D] = integrateObjSensAdj(m, con, obj, intOpts);
            else
                if opts.AdaptAbsTol
                    [G, D, intOpts.AbsTol] = computeObjSens(m, con, obj, intOpts);
                else
                    [G, D] = computeObjSens(m, con, obj, intOpts);
                end
            end
            
            % Normalize gradient
            if opts.Normalized
                D = vec(D) .* v;
            end
        end
    end

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%% Halt optimization on terminal goal %%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function stop = isTerminalGoal(x, optimValues, state)
        if optimValues.fval  < opts.TerminalGoal
            aborted = true;
            vAbort = x;
            stop = true;
        else
            stop = false;
        end
    end

end
