% SPDX-License-Identifier: MIT
% Copyright © 2021 Weiwei "William" Kong
%
% Self-contained one-file version for generating a linearly constrained,
% strongly convex, L1-regularized logistic regression test problem.

function [F, params] = test_fn_lin_constr_logreg(N, M, m, seed, dimM, dimN)
%TEST_FN_LIN_CONSTR_LOGREG_L1_01 Generate a constrained logistic-regression test instance.
%
%   [F, params] = test_fn_lin_constr_logreg_l1_01(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) = mean_i log(1 + exp(-y_i * X_i*x)) + (rho/2)*||x||_2^2,
%   psi_n(x) = lambda*||x||_1,
%   rho      = m,
%   lambda   = 1/sqrt(dimN).
%
% The data matrix X is scaled so that the logistic-loss Hessian contribution
% has curvature bounded by M-m, making the smooth part strongly convex with
% lower curvature at least m and upper curvature at most M.
%
% Inputs:
%   N     - number of logistic-regression samples
%   M     - target upper curvature of smooth objective
%   m     - target lower curvature / ridge parameter; must be positive
%   seed  - random seed
%   dimM  - number of linear constraints
%   dimN  - number of optimization variables
%
% Outputs:
%   F      - oracle struct with fields:
%              fs, fn, grad, prox, Residual, Adjoint, Hessian
%            plus aliases finit and gradinit for compatibility with older code.
%   params - struct containing generated data, constraints, topology, and
%            curvature diagnostics.
%
% Notes:
%   - This file is self-contained. No external helper functions are needed.
%   - F.prox accepts both F.prox(z) and F.prox(z, step). If step is omitted,
%     it defaults to 1.

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 m < 0
    error('test_fn_lin_constr_logreg_l1_01:BadCurvature', ...
        'm must be nonnegative so that the smooth objective is strongly convex.');
end
if M <= m
    error('test_fn_lin_constr_logreg_l1_01:BadCurvature', ...
        'M must be strictly larger than m.');
end

% Initialize.
rng(seed);

X = randn(N, dimN);
% XtX = X' * X;
maxEigXtX = (max(svd(X)))^2;
if maxEigXtX <= 0
    error('test_fn_lin_constr_logreg_l1_01:BadDataMatrix', ...
        'Could not scale X because X''*X has nonpositive largest eigenvalue.');
end

% Since logistic''(t) <= 1/4, scaling X this way gives
% lambda_max(X'*D*X/N) <= M-m for every diagonal 0 <= D_ii <= 1/4.
X = X / sqrt(maxEigXtX) * sqrt(4 * N * (M - m));

y = 2 * randi([0, 1], N, 1) - 1;
A = randn(dimM, dimN);
rho = m;
lambda = 1 / sqrt(dimN);


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

% Initial point and feasible right-hand side.
params.x0 = rand(dimN, 1);

vec=rand(dimN,1)*0.4-2*rand(dimN,1);
b = A * vec;

% Constraint map methods.
params.constr_fn = @(z) A * z;
params.grad_constr_fn = @(z) A';
params.set_projector = @(z) b; %#ok<NASGU> compatible with existing drivers
params.K_constr = 0;

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


% Oracle construction.
[f_s, grad_f_s] = cached_logistic_oracles(X, y, rho, norm_fn);
f_n = @(x) lambda * norm(x, 1);
prox_f_n = @(x, varargin) prox_l1(x, lambda, varargin{:});

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) logistic_hessian(X, y, rho, x);


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


% Simple local estimates from x0 to w0
params.w0 = randn(dimN,1);

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 [f_s, grad_f_s] = cached_logistic_oracles(X, y, rho, norm_fn)
%CACHED_LOGISTIC_ORACLES Smooth logistic value and gradient with 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)
        if isempty(last_x) || ~isequal(x, last_x)
            x = x(:);
            margin = y .* (X * x);
            loss = mean(softplus_stable(-margin));
            value = loss + (rho / 2) * norm_fn(x)^2;

            % d/dx mean softplus(-y_i X_i x)
            % = -(1/N) X' * (y .* sigmoid(-y .* X*x)).
            weights = y .* sigmoid_stable(-margin);
            g = -(X' * weights) / N + rho * x;

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

        value = last_value;
        g = last_grad;
    end
end

function H = logistic_hessian(X, y, rho, x)
%LOGISTIC_HESSIAN Hessian of the smooth logistic-plus-ridge term.

N = size(X, 1);
dimN = size(X, 2);
x = x(:);

margin = y .* (X * x);
s = sigmoid_stable(margin);
weights = s .* (1 - s);

H = X' * bsxfun(@times, X, weights) / N + rho * speye(dimN);
H = (H + H') / 2;
end

function z = prox_l1(x, lambda, varargin)
%PROX_L1 Proximal map of lambda*||.||_1.
%
%   z = prox_l1(x, lambda)       uses step = 1.
%   z = prox_l1(x, lambda, step) uses z = prox_{step*lambda||.||_1}(x).

if nargin < 3 || isempty(varargin)
    step = 1;
else
    step = varargin{1};
end

check_real_scalar(step, 'step');
if step < 0
    error('prox_l1:BadStep', 'step must be nonnegative.');
end

z = soft_threshold(x, step * lambda);
end

function z = soft_threshold(x, tau)
%SOFT_THRESHOLD Componentwise soft-thresholding.
z = sign(x) .* max(abs(x) - tau, 0);
end

function value = softplus_stable(z)
%SOFTPLUS_STABLE Numerically stable log(1 + exp(z)).
value = max(z, 0) + log1p(exp(-abs(z)));
end

function s = sigmoid_stable(z)
%SIGMOID_STABLE Numerically stable sigmoid, 1/(1+exp(-z)).
s = zeros(size(z));
idx = z >= 0;
s(idx) = 1 ./ (1 + exp(-z(idx)));

ez = exp(z(~idx));
s(~idx) = ez ./ (1 + ez);
end


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