% SPDX-License-Identifier: MIT
%
% Self-contained generator for linearly constrained Bayesian D-optimal
% experimental design over a compact box.

function [F, params] = test_fn_lin_constr_ML3(N, M, m, seed, dimM, dimN)
%TEST_FN_LIN_CONSTR_ML3 Generate Bayesian D-optimal design instances.
%
%   [F, params] = test_fn_lin_constr_ML3( ...
%       N, M, m, seed, dimM, dimN)
%
% Problem generated:
%
%   minimize_x  psi_s(x) + psi_n(x)
%   subject to  A*x = b,
%
% where
%
%   psi_s(x) = -gamma*logdet(I_N + sum_i x_i*a_i*a_i')
%              + (m/2)*||x||_2^2,
%   psi_n(x) = indicator_[0,2/sqrt(dimN)]^dimN(x).
%
% Here N is the dimension of the Bayesian information matrix and dimN is
% the number of candidate experiments/design weights. The box has Euclidean
% diameter 2 for every dimN. The smooth objective is convex when m=0 and
% m-strongly convex when m>0.
%
% gamma is chosen so that the largest Hessian eigenvalue at x0=0 is M.
% Since the information matrix increases over the nonnegative box, M is a
% global smoothness constant on the box. The condition
%
%   dimN > N*(N+1)/2
%
% makes the log-determinant Hessian rank deficient, so the global lower
% curvature is exactly m.
%
% The constraints consist of one total experimental-budget equality and
% sparse category-budget equalities over disjoint experiment families.
% The uniform reference design is used to define all budget targets.

if nargin < 1 || isempty(N),    N = 10;       end
if nargin < 2 || isempty(M),    M = 20;       end
if nargin < 3 || isempty(m),    m = 0;        end
if nargin < 4 || isempty(seed), seed = 1;     end
if nargin < 5 || isempty(dimM), dimM = 20;    end
if nargin < 6 || isempty(dimN), dimN = 100;   end

check_positive_integer(N,    'N');
check_positive_integer(dimM, 'dimM');
check_positive_integer(dimN, 'dimN');
check_real_scalar(M, 'M');
check_real_scalar(m, 'm');

if m < 0
    error('test_fn_lin_constr_ML3:BadCurvature', ...
        'm must be nonnegative.');
end
if M <= m
    error('test_fn_lin_constr_ML3:BadCurvature', ...
        'M must be strictly larger than m.');
end
if dimN <= N * (N + 1) / 2
    error('test_fn_lin_constr_ML3:BadDimension', ...
        'dimN must exceed N*(N+1)/2.');
end
if dimM > dimN
    error('test_fn_lin_constr_ML3:BadDimension', ...
        'dimM must not exceed dimN for the budget constraints.');
end

% Initialize.
rng(seed);

prior = 1;
box_upper = 2 / sqrt(dimN);

% Candidate experiment vectors. Every column has Euclidean norm 1/2.
candidate_vectors = randn(N, dimN);
column_norms = sqrt(sum(candidate_vectors.^2, 1));
candidate_vectors = 0.5 * bsxfun( ...
    @rdivide, candidate_vectors, max(column_norms, eps));

% At x=0, the log-determinant Hessian before multiplication by gamma is
% (Q0 .* Q0), where Q0 = V'*(prior*I)^{-1}*V.
Q0 = (candidate_vectors' * candidate_vectors) / prior;
base_hessian = Q0 .* Q0;
base_hessian = (base_hessian + base_hessian') / 2;
base_upper_curvature = eigs(base_hessian, 1, 'la');
gamma = (M - m) / base_upper_curvature;

% A feasible interior design and an infeasible initial point.
x_feasible = ones(dimN, 1) / dimN;
params.x0 = zeros(dimN, 1);
params.w0 = x_feasible;

A = budget_constraint_matrix(dimM, dimN);
b = A * x_feasible;
norm_A = norm(full(A), 2);  % MATLAB does not support norm(A,2) for sparse A

% Euclidean topology.
prod_fn = @(u, v) sum(u(:) .* v(:));
norm_fn = @(u) norm(u, 'fro');

% Constraint-map methods.
params.constr_fn = @(z) A * z;
params.grad_constr_fn = @(z) A';
params.set_projector = @(z) b;
params.K_constr = norm_A;

% Basic output parameters.
params.prod_fn = prod_fn;
params.norm_fn = norm_fn;
params.M = M;
params.m = m;

% Oracle construction.
[f_s, grad_f_s] = cached_d_optimal_oracles( ...
    candidate_vectors, prior, gamma, m, norm_fn);
f_n = @(x) indicator_box(x, box_upper);
prox_f_n = @(x, varargin) min(box_upper, max(0, x));

Residual = @(z) A * z - b;
Adjoint = @(u) A' * u;

F.fs = f_s;
F.fn = f_n;
F.grad = grad_f_s;
F.prox = prox_f_n;
F.Residual = Residual;
F.Adjoint = Adjoint;
F.Hessian = @(x) d_optimal_hessian( ...
    candidate_vectors, prior, gamma, m, x);

% Exact global curvature values on the box.
params.UpperCurv = M;
params.LowerCurv = m;
params.CondNumber = params.UpperCurv / params.LowerCurv;

% Simple local estimates from x0 to w0.
dx0 = params.w0 - params.x0;
nx0 = norm_fn(dx0);

gx0 = grad_f_s(params.x0);
gw0 = grad_f_s(params.w0);
params.muEst = 2 * (f_s(params.w0) - f_s(params.x0) ...
    - prod_fn(gx0, dx0)) / nx0^2;
params.muEst = max(params.muEst, 0);
params.LipEst = norm_fn(gw0 - gx0) / nx0;
params.LipEst = max(params.LipEst, eps);

end


function A = budget_constraint_matrix(dimM, dimN)
%BUDGET_CONSTRAINT_MATRIX Total and category experimental budgets.

A = spalloc(dimM, dimN, 2 * dimN - 1);
A(1, :) = 1;

if dimM > 1
    permutation = randperm(dimN);
    category_indices = permutation(1:end-1);  % one ungrouped experiment
    edges = floor((0:(dimM - 1)) * numel(category_indices) / (dimM - 1));

    for j = 2:dimM
        idx = category_indices(edges(j - 1) + 1:edges(j));
        A(j, idx) = 1;
    end
end
end


function [f_s, grad_f_s] = cached_d_optimal_oracles( ...
    candidate_vectors, prior, gamma, m, norm_fn)
%CACHED_D_OPTIMAL_ORACLES Smooth value/gradient with one-point cache.

last_x = [];
last_value = [];
last_grad = [];

f_s = @value_oracle;
grad_f_s = @grad_oracle;

    function value = value_oracle(x)
        [value, ~] = eval_at(x);
    end

    function g = grad_oracle(x)
        [~, g] = eval_at(x);
    end

    function [value, g] = eval_at(x)
        x = x(:);
        if isempty(last_x) || ~isequal(x, last_x)
            [R, inverse_times_vectors] = information_factorization( ...
                candidate_vectors, prior, x);

            log_determinant = 2 * sum(log(diag(R)));
            leverage = sum(candidate_vectors .* inverse_times_vectors, 1)';

            value = -gamma * log_determinant ...
                + (m / 2) * norm_fn(x)^2;
            g = -gamma * leverage + m * x;

            last_x = x;
            last_value = value;
            last_grad = g;
        end

        value = last_value;
        g = last_grad;
    end
end


function H = d_optimal_hessian(candidate_vectors, prior, gamma, m, x)
%D_OPTIMAL_HESSIAN Hessian of the smooth Bayesian D-optimal objective.

dimN = size(candidate_vectors, 2);
[~, inverse_times_vectors] = information_factorization( ...
    candidate_vectors, prior, x);

Q = candidate_vectors' * inverse_times_vectors;
H = gamma * (Q .* Q) + m * speye(dimN);
H = (H + H') / 2;
end


function [R, inverse_times_vectors] = information_factorization( ...
    candidate_vectors, prior, x)
%INFORMATION_FACTORIZATION Cholesky factor and inverse-times-vectors.

N = size(candidate_vectors, 1);
weighted_vectors = bsxfun(@times, candidate_vectors, x(:)');
information_matrix = prior * eye(N) ...
    + weighted_vectors * candidate_vectors';
information_matrix = (information_matrix + information_matrix') / 2;

R = chol(information_matrix);
inverse_times_vectors = R \ (R' \ candidate_vectors);
end


function value = indicator_box(x, upper_bound)
%INDICATOR_BOX Indicator of [0,upper_bound]^n.

tolerance = 1e-12;
if all(x(:) >= -tolerance) ...
        && all(x(:) <= upper_bound + tolerance)
    value = 0;
else
    value = Inf;
end
end


function check_positive_integer(x, name)
if ~(isscalar(x) && isnumeric(x) && isfinite(x) && x == round(x) && x > 0)
    error('test_fn_lin_constr_ML3:BadInput', ...
        '%s must be a positive integer.', name);
end
end


function check_real_scalar(x, name)
if ~(isscalar(x) && isnumeric(x) && isreal(x) && isfinite(x))
    error('test_fn_lin_constr_ML3:BadInput', ...
        '%s must be a finite real scalar.', name);
end
end
