function prob = adjoint_add(prob)
% ADJOINT_ADD   Zero functions corresponding to the generalized adjoint problem
%
% This function appends the generalized adjoint zero function 
%               (u,lambda,eta)->(lambda^T, eta^T)*C(u)
% and the corresponding monitor functions 
%               (u,lambda,eta)->eta
% with continuation parameters obtained by prepending "d." to the label of
% the corresponding continuation parameters in the original extended
% continuation problem.

efunc   = prob.efunc;
adjoint = prob.adjoint;
data    = coco_func_data;
data    = init_data(efunc, adjoint, data);
prob    = coco_add_func(prob, 'coco_adjoint', @adeqs, data, 'zero', ...
  'uidx', efunc.x_idx, 'x0', adjoint.l0, 't0', adjoint.tl0, ...
  'fdim', adjoint.a_dim(2), 'f+df', 'remesh', @remesh);
if ~isempty(adjoint.active)
  prob    = coco_add_pars(prob, 'coco_adjoint_active_pars', ...
    data.l_idx(adjoint.l_idx(adjoint.active)), ...
    adjoint.lnames(adjoint.active), 'active');
end
if ~isempty(adjoint.inactive)
  prob    = coco_add_pars(prob, 'coco_adjoint_inactive_pars', ...
    data.l_idx(adjoint.l_idx(adjoint.inactive)), ...
    adjoint.lnames(adjoint.inactive), 'inactive');
end
prob    = coco_add_slot(prob, 'coco_adjoint', @coco_save_data, data, ...
  'save_full');

end

function data = init_data(efunc, adjoint, data)

dim = efunc.x_dim;
data.x_idx = 1:dim;
data.l_idx = dim + (1:numel(adjoint.l0));

data.xtr     = zeros(numel(adjoint.l0),1);
data.xtrstay = data.l_idx(adjoint.l_idx) - dim;
data.xtr(data.xtrstay) = data.xtrstay;

ncols    = adjoint.a_dim(2);
data.shp = [ncols, dim];
adjoint.dax_idx = 1:(ncols*dim);
for i = 1:numel(adjoint.identifyers)
  func   = adjoint.funcs(i);
  ax_idx = func.ax_idx;
  x_idx  = func.x_idx;
  
  dax_idx = kron(ones(1,numel(x_idx)), ax_idx);
  dx_idx  = kron(x_idx, ones(1,numel(ax_idx)));
  dax_idx = dax_idx+ncols*(dx_idx-1);
  
  adjoint.funcs(i).dax_idx = dax_idx;
end

data.adjoint = adjoint;

end

function [data, f, J] = adeqs(prob, data, u)

pr = data.pr;

x = u(pr.x_idx);
l = u(pr.l_idx);

[pr, C] = adfunc(prob, pr, x);
f = l'*C;
f = f(:);

if nargout>=3
  [pr, dC] = adfunc_DU(prob, pr, x);
  J = [reshape(l'*dC, pr.shp),  C'];
end

data.pr = pr;

end

function [data, C] = adfunc(prob, data, x)

% create cache for variables that grow
persistent urows ucols uvals

if nargin==0
  % clear cache
  urows = [];
  ucols = [];
  uvals = [];
  return;
end

%% extract u from xp=[u;p]

adjoint = data.adjoint;
nrows   = adjoint.a_dim(1);
ncols   = adjoint.a_dim(2);
idx_MX  = 0;

%% compute adjoint coefficient matrix

for i = 1:numel(adjoint.identifyers)
  func   = adjoint.funcs(i);
  af_idx = func.af_idx;
  ax_idx = func.ax_idx;
  x_idx  = func.x_idx;
  fdata  = func.data;
  
  % evaluate adjoint of function
  [fdata, A] = func.F(prob, fdata, x(x_idx));
  
  % merge adjoints
  [r, c, v]  = find(A);
  idx        = idx_MX+(1:numel(r));
  idx_MX     = idx_MX+numel(r);
  urows(idx) = af_idx(r)';
  ucols(idx) = ax_idx(c)';
  uvals(idx) = v;
  
  adjoint.funcs(i).data = fdata;
end

data.adjoint = adjoint;
%% apply permutations

idx = 1:idx_MX;
C = sparse(adjoint.af_idx(urows(idx)), adjoint.ax_idx(ucols(idx)), uvals(idx), nrows, ncols);

end

function [data, dC] = adfunc_DU(prob, data, x)

% create cache for variables that grow
persistent urows ucols uvals

if nargin==0
  % clear cache
  urows = [];
  ucols = [];
  uvals = [];
  return;
end

%% extract u from xp=[u;p]

adjoint = data.adjoint;
nrows   = adjoint.a_dim(1);
ncols   = adjoint.a_dim(2);
idx_MX  = 0;

%% compute Jacobian of adjoint matrix

for i = 1:numel(adjoint.identifyers)
  func    = adjoint.funcs(i);
  af_idx  = func.af_idx;
  dax_idx = func.dax_idx;
  x_idx   = func.x_idx;
  fdata   = func.data;
  
  % evaluate jacobian of adjoint of function
  if ~isempty(func.DFDX)
    [fdata, dA] = func.DFDX(prob, fdata, x(x_idx));
  else
    [fdata, dA] = coco_ezDFDX('f(o,d,x)', prob, fdata, @func.F, x(x_idx));
  end
  
  % reshape jacobian
  dA = reshape(dA, [numel(af_idx), numel(dA)/numel(af_idx)]);
    
  % merge adjoints
  [r, c, v]  = find(dA);
  idx        = idx_MX+(1:numel(r));
  idx_MX     = idx_MX+numel(r);
  urows(idx) = af_idx(r)';
  ucols(idx) = dax_idx(c)';
  uvals(idx) = v;
  
  adjoint.funcs(i).data = fdata;
end

data.adjoint = adjoint;

%% apply permutations

idx = 1:idx_MX;
dC = sparse(adjoint.af_idx(urows(idx)), adjoint.dax_idx(ucols(idx)), uvals(idx), nrows, ncols*numel(data.x_idx));

end

function [prob, stat, xtr] = remesh(prob, data, chart, ub, Vb) %#ok<INUSD>
% Remesh function for adjoint problem.

efunc   = prob.efunc;
adjoint = prob.adjoint;
xtr     = data.xtr;
xtrstay = data.xtrstay;
data    = init_data(efunc, adjoint, data);
xtr(xtrstay) = data.xtr(data.xtrstay);
prob    = coco_change_func(prob, data, 'uidx', efunc.x_idx, ...
  'u0', adjoint.l0, 'fdim', adjoint.a_dim(2), 'vecs', adjoint.Vl0);
stat = 'success';

end
