% SPDX-License-Identifier: MIT
% Copyright © 2021 Weiwei "William" Kong

function [F,params] = test_fn_lin_constr_01(N, M, m, seed, dimM, dimN)
% Generator of a test suite of linearly constrained strongly convex QP functions.
%
% Arguments:
%
%   N (int): One of the objective function's hyperparameters.
%
%   dimM (int): One of the objective function's hyperparameters.
%
%   dimN (int): One of the objective function's hyperparameters.
%
%   M (double): The target upper curvature of the objective function.
%
%   m (double): The target lower curvature of the objective function.
%
%   seed (int): The number used to seed MATLAB's random number generator.
%
% Returns:
%
%   A pair consisting of an Oracle and a struct. The oracle is the first-order
%   oracle underlying the optimization problem and the struct contains the
%   relevant hyperparameters of the problem.
%

% Initialize.
rng(seed);
D = sparse(1:dimN, 1:dimN, randi([1, N], 1, dimN), dimN, dimN);
C = rand(dimM, dimN);
Bseed = rand(dimN, dimN);
A = rand(dimM, dimN);
d = rand([dimM, 1]);

% Calibrate the two positive quadratic terms so that the resulting Hessian
% has minimum eigenvalue m and maximum eigenvalue M.
[B, tau, xi] = calibrate_convex_quadratic(Bseed, C, D, M, m);

% Compute the norm of A and the quadratic factors.
Hn = B' * (D' * D) * B;
AtA = A' * A;
Hn = (Hn + Hn') / 2;
AtA = (AtA + AtA') / 2;
norm_A = sqrt(eigs(AtA, 1, 'la')); % same as norm(A,2)

% Set the topology (Euclidean).
prod_fn = @(a,b) sum(dot(a, b));
norm_fn = @(a) norm(a, 'fro');

% Compute the b vector.
b = A * ones(dimN, 1) / dimN;

% 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 params.
params.prod_fn = prod_fn;
params.norm_fn = norm_fn;
x0_base = rand(dimN, 1);
params.x0 = x0_base / sum(x0_base);

% Oracle construction.
H = xi * Hn + tau * (C' * C);
H = (H + H') / 2;
q = -tau * (C' * d);
f_const = 0.5 * tau * norm_fn(d)^2;

[f_s, grad_f_s] = cached_quadratic_oracles(H, q, f_const, prod_fn);
f_n = @(x) 0;
prox_f_n = @(x,step) sm_proj(x, 1);

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

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) H;

% Store the realized curvatures of the actual objective Hessian.
Eigenvalues = eig(H);
params.UpperCurv = max(Eigenvalues);
params.LowerCurv = min(Eigenvalues);
params.M = params.UpperCurv;
params.m = params.LowerCurv;
params.CondNumber = params.UpperCurv / params.LowerCurv;

params.w0 = zeros(dimN,1);
[~, w0_idx] = min(params.x0);
params.w0(w0_idx) = 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 [B, tau, xi] = calibrate_convex_quadratic(Bseed, C, D, M, m)
%CALIBRATE_CONVEX_QUADRATIC Realize the prescribed positive curvatures.
%
% Starting from the random seed matrices, form G and apply the affine
% spectral map H = scale*G + shift*I, whose extreme eigenvalues are m and M.
% The shift is absorbed into the first squared-norm term, so the final
% objective retains the form
%
%   (xi/2)*||D*B*x||^2 + (tau/2)*||C*x-d||^2.

if ~(isscalar(M) && isscalar(m) && isfinite(M) && isfinite(m) ...
        && 0 < m && m < M)
    error('test_fn_lin_constr_01:BadCurvature', ...
        'The curvature inputs must satisfy 0 < m < M.');
end

Hn_seed = Bseed' * (D' * D) * Bseed;
Hc = C' * C;
G = Hn_seed + Hc;
G = (G + G') / 2;

raw_eigenvalues = eig(G);
raw_lower = min(raw_eigenvalues);
raw_upper = max(raw_eigenvalues);
scale = (M - m) / (raw_upper - raw_lower);
shift = m - scale * raw_lower;

% A nonnegative shift is equivalent to requiring that the requested
% condition number not exceed that of the random seed Hessian.
tolerance = 1e-10 * max(1, M);
if shift < -tolerance
    error('test_fn_lin_constr_01:UnattainableCurvaturePair', ...
        ['The requested curvature ratio M/m exceeds that of the random ', ...
         'seed Hessian.']);
end
shift = max(shift, 0);

% Absorb scale*Hn_seed + shift*I into B'*D'*D*B. A random orthogonal
% left factor makes the resulting B dense without changing its Gram matrix.
first_gram = scale * Hn_seed + shift * eye(size(Bseed, 2));
first_gram = (first_gram + first_gram') / 2;
[R, chol_flag] = chol(first_gram);
if chol_flag ~= 0
    error('test_fn_lin_constr_01:FactorizationFailure', ...
        'The calibrated first quadratic factor is not positive definite.');
end
[Q,~] = qr(Bseed, 0);
B = D \ (Q * R);

tau = scale;
xi = 1;
end


function [f_s,grad_f_s] = cached_quadratic_oracles(H,q,f_const,prod_fn)
last_x = [];
last_Hx = [];
last_value = [];

f_s = @valueOracle;
grad_f_s = @gradOracle;

    function value = valueOracle(x)
        [~,value] = evalAt(x);
    end

    function g = gradOracle(x)
        Hx = evalAt(x);
        g = Hx + q;
    end

    function [Hx,value] = evalAt(x)
        if isempty(last_x) || ~isequal(x,last_x)
            last_x = x;
            last_Hx = H*x;
            last_value = 0.5*prod_fn(x,last_Hx) + prod_fn(q,x) + f_const;
        end
        Hx = last_Hx;
        value = last_value;
    end
end
