% SPDX-License-Identifier: MIT
%
% Self-contained generator for a linearly constrained elastic-net
% least-squares regression problem.

function [F,params] = test_fn_lin_constr_ML1(N, M, m, seed, dimM, dimN)
%TEST_FN_LIN_CONSTR_ML1 Generate constrained elastic-net regression data.
%
% Problem generated:
%
%   minimize_x  (1/(2*N))*||X*x-y||_2^2 + (m/2)*||x||_2^2
%               + lambda*||x||_1
%   subject to  A*x = b,
%
% where lambda = 1/sqrt(dimN). The design matrix is scaled so that the
% smooth objective has upper curvature M and lower curvature m.
%
% A feasible point is used to define b, while A is constructed so that the
% returned initial point x0 is infeasible. Hence the problem is feasible and
% x0 cannot be an optimal solution.

if nargin < 1 || isempty(N),    N = 200;     end
if nargin < 2 || isempty(M),    M = 10;      end
if nargin < 3 || isempty(m),    m = 1e-2;    end
if nargin < 4 || isempty(seed), seed = 1;    end
if nargin < 5 || isempty(dimM), dimM = 25;   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 dimN < 2
    error('test_fn_lin_constr_ML1:BadDimension', ...
        'dimN must be at least 2.');
end
if m < 0 || M <= m
    error('test_fn_lin_constr_ML1:BadCurvature', ...
        'The curvature inputs must satisfy 0 <= m < M.');
end

% Initialize.
rng(seed);
rho = m;
lambda = 1 / sqrt(dimN);

% Dense rank-deficient design, scaled so that
% lambda_max(X'*X/N) = M-m.
X = dense_rank_deficient_design(N, dimN);
X = X * (sqrt(N * (M - m)) / norm(X, 2));

% Sparse feasible reference point and regression observations.
num_nonzeros = max(1, round(0.10 * dimN));
support = randperm(dimN, num_nonzeros);
x_feasible = zeros(dimN, 1);
x_feasible(support) = randn(num_nonzeros, 1);
x_feasible(support(1)) = 1;
x_feasible = x_feasible / norm(x_feasible);

signal = X * x_feasible;
noise_scale = 0.05 * max(norm(signal) / sqrt(N), 1);
y = signal + noise_scale * randn(N, 1);

% The offset has unit norm, so x0 differs from the feasible reference point.
params.x0 = x_feasible + ones(dimN, 1) / sqrt(dimN);
params.w0 = x_feasible;

% Feasible constraints with a guaranteed infeasible initial point.
[A,b] = consistent_constraints(x_feasible, params.x0, dimM);

% Euclidean topology.
prod_fn = @(a,b) sum(a(:) .* b(:));
norm_fn = @(a) norm(a, '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, 2);

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

% Oracle construction.
H = (X' * X) / N + rho * speye(dimN);
H = (H + H') / 2;
[f_s,grad_f_s] = cached_least_squares_oracles(X, y, rho, norm_fn);
f_n = @(x) lambda * norm(x, 1);
prox_f_n = @(x,varargin) prox_l1(x, lambda, varargin{:});

F.fs = f_s;
F.fn = f_n;
F.grad = grad_f_s;
F.prox = prox_f_n;
F.Residual = @(z) A * z - b;
F.Adjoint = @(u) A' * u;
F.Hessian = @(x) H;

% Curvature information.
params.UpperCurv = M;
params.LowerCurv = m;
params.CondNumber = M / m;

% Local estimates from x0 to w0, as in the existing generators.
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.LipEst = norm_fn(gw0 - gx0) / nx0;

end


function [A,b] = consistent_constraints(x_feasible, x0, dimM)
%CONSISTENT_CONSTRAINTS Ensure feasibility and make x0 infeasible.

d = x0(:) - x_feasible(:);
A = randn(dimM, numel(x_feasible));
A(1,:) = A(1,:) + ((1 - A(1,:) * d) / (d' * d)) * d';
b = A * x_feasible(:);
end


function X = dense_rank_deficient_design(N, dimN)
%DENSE_RANK_DEFICIENT_DESIGN Dense matrix with rank strictly below dimN.

rank_target = min(N, dimN - 1);
G = randn(N, rank_target);
[V,~] = qr(randn(dimN, rank_target), 0);
X = G * V';
end


function [f_s,grad_f_s] = cached_least_squares_oracles(X, y, rho, norm_fn)
%CACHED_LEAST_SQUARES_ORACLES Smooth value and gradient with a one-point cache.

N = size(X, 1);
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)
            residual = X * x - y;
            value = 0.5 * (residual' * residual) / N ...
                + (rho / 2) * norm_fn(x)^2;
            g = (X' * residual) / N + rho * x;

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


function z = prox_l1(x, lambda, varargin)
%PROX_L1 Proximal map of lambda*||.||_1.

if isempty(varargin)
    step = 1;
else
    step = varargin{1};
end
z = sign(x) .* max(abs(x) - step * lambda, 0);
end


function check_positive_integer(x, name)
if ~(isscalar(x) && isnumeric(x) && isfinite(x) ...
        && x == round(x) && x > 0)
    error('test_fn_lin_constr_ML1: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_ML1:BadInput', ...
        '%s must be a finite real scalar.', name);
end
end
