% SPDX-License-Identifier: MIT
%
% Self-contained generator for a linearly constrained group-sparse
% Huberized-hinge SVM problem.

function [F,params] = test_fn_lin_constr_ML2(N, M, m, seed, dimM, dimN)
%TEST_FN_LIN_CONSTR_ML2 Generate a constrained Huberized-hinge SVM.
%
% Problem generated:
%
%   minimize_x  mean_i h_delta(1-y_i*X_i*x) + (m/2)*||x||_2^2
%               + lambda*sum_g ||x_g||_2
%   subject to  A*x = b,
%
% where delta = 2, lambda = 1/sqrt(dimN), and the variables are divided
% into contiguous groups of at most ten coordinates. The design matrix is
% scaled so that the smooth objective has upper curvature M and lower
% curvature m.
%
% The equality constraints compare average linear scores across matched
% synthetic cohorts, giving a data-derived linear fairness operator. The
% reference classifier defines the target score gaps.

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_ML2:BadDimension', ...
        'dimN must be at least 2.');
end
if m < 0 || M <= m
    error('test_fn_lin_constr_ML2:BadCurvature', ...
        'The curvature inputs must satisfy 0 <= m < M.');
end
if 2 * dimM > N
    error('test_fn_lin_constr_ML2:BadDimension', ...
        'N must be at least 2*dimM for the cohort constraints.');
end

% Initialize.
rng(seed);
rho = m;
lambda = 1 / sqrt(dimN);
delta = 2;
group_size = min(10, dimN);
groups = contiguous_groups(dimN, group_size);

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

% Group-sparse feasible reference point and binary labels.
num_groups = numel(groups);
num_active_groups = max(1, ceil(0.20 * num_groups));
active_groups = randperm(num_groups, num_active_groups);
x_feasible = zeros(dimN, 1);
for k = 1:num_active_groups
    idx = groups{active_groups(k)};
    x_feasible(idx) = randn(numel(idx), 1);
end
x_feasible(groups{active_groups(1)}(1)) = 1;
x_feasible = x_feasible / norm(x_feasible);

scores = X * x_feasible;
noise_scale = 0.10 * max(norm(scores) / sqrt(N), 1);
latent = scores + noise_scale * randn(N, 1);
y = ones(N, 1);
y(latent < 0) = -1;

% x0=0 places every sample in the quadratic part of h_delta at x0.
params.x0 = zeros(dimN, 1);
params.w0 = x_feasible;

% Matched-cohort fairness constraints on average linear scores.
A = matched_cohort_fairness_constraints(X, x_feasible, dimM);
b = A * x_feasible;

% 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.
[f_s,grad_f_s] = cached_huberized_svm_oracles( ...
    X, y, rho, delta, norm_fn);
f_n = @(x) group_lasso_value(x, lambda, groups);
prox_f_n = @(x,varargin) prox_group_lasso( ...
    x, lambda, groups, 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) huberized_svm_hessian(X, y, rho, delta, x);

% 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 = matched_cohort_fairness_constraints(X, x_reference, dimM)
%MATCHED_COHORT_FAIRNESS_CONSTRAINTS Average-score contrast rows.

N = size(X, 1);
dimN = size(X, 2);
group_size = floor(N / (2 * dimM));
scores = X * x_reference(:);
[~, order] = sort(scores, 'ascend');

low_group = order(1:group_size);
high_group = order(end-group_size+1:end);
first_row = mean(X(high_group, :), 1) - mean(X(low_group, :), 1);
first_gap = first_row * x_reference(:);
if first_gap <= eps
    error('test_fn_lin_constr_ML2:DegenerateFairnessRow', ...
        'Could not construct a nonzero reference score gap.');
end

A = zeros(dimM, dimN);
A(1, :) = first_row / first_gap;

if dimM > 1
    used = false(N, 1);
    used(low_group) = true;
    used(high_group) = true;
    pool = find(~used);
    pool = pool(randperm(numel(pool)));
    cursor = 1;

    for j = 2:dimM
        group_one = pool(cursor:cursor + group_size - 1);
        group_two = pool(cursor + group_size:cursor + 2 * group_size - 1);
        row = mean(X(group_one, :), 1) - mean(X(group_two, :), 1);
        A(j, :) = row / max(norm(row, 2), eps);
        cursor = cursor + 2 * group_size;
    end
end
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 groups = contiguous_groups(dimN, group_size)
%CONTIGUOUS_GROUPS Partition 1:dimN into contiguous groups.

starts = 1:group_size:dimN;
groups = cell(numel(starts), 1);
for k = 1:numel(starts)
    groups{k} = starts(k):min(starts(k) + group_size - 1, dimN);
end
end


function [f_s,grad_f_s] = cached_huberized_svm_oracles( ...
    X, y, rho, delta, norm_fn)
%CACHED_HUBERIZED_SVM_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)
            shortfall = 1 - y .* (X * x);
            [loss,slope] = huberized_hinge(shortfall, delta);
            value = mean(loss) + (rho / 2) * norm_fn(x)^2;
            g = -(X' * (y .* slope)) / N + rho * x;

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


function [loss,slope] = huberized_hinge(t, delta)
%HUBERIZED_HINGE Value and derivative of the Huberized hinge.

loss = zeros(size(t));
slope = zeros(size(t));
quadratic = (t > 0) & (t < delta);
linear = (t >= delta);
loss(quadratic) = t(quadratic).^2 / (2 * delta);
slope(quadratic) = t(quadratic) / delta;
loss(linear) = t(linear) - delta / 2;
slope(linear) = 1;
end


function H = huberized_svm_hessian(X, y, rho, delta, x)
%HUBERIZED_SVM_HESSIAN Piecewise Hessian of the smooth objective.

N = size(X, 1);
dimN = size(X, 2);
shortfall = 1 - y .* (X * x(:));
active = double((shortfall > 0) & (shortfall < delta));
H = X' * bsxfun(@times, X, active) / (N * delta) ...
    + rho * speye(dimN);
H = (H + H') / 2;
end


function value = group_lasso_value(x, lambda, groups)
%GROUP_LASSO_VALUE Evaluate lambda*sum_g ||x_g||_2.

x = x(:);
value = 0;
for k = 1:numel(groups)
    value = value + norm(x(groups{k}), 2);
end
value = lambda * value;
end


function z = prox_group_lasso(x, lambda, groups, varargin)
%PROX_GROUP_LASSO Block soft-thresholding.

if isempty(varargin)
    step = 1;
else
    step = varargin{1};
end

x = x(:);
z = zeros(size(x));
threshold = step * lambda;
for k = 1:numel(groups)
    idx = groups{k};
    block = x(idx);
    block_norm = norm(block, 2);
    if block_norm > threshold
        z(idx) = (1 - threshold / block_norm) * block;
    end
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_ML2: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_ML2:BadInput', ...
        '%s must be a finite real scalar.', name);
end
end
