% SPDX-License-Identifier: MIT
%
% Self-contained generator for linearly constrained Huberized quantum
% state tomography over the density-matrix spectrahedron.

function [F,params] = test_fn_lin_quant_SDP4(N, M, m, seed, dimM, dimN)
%TEST_FN_LIN_QUANT_SDP4 Generate a constrained quantum tomography problem.
%
% Problem generated:
%
%   minimize_x  (gamma/N)*sum_i h_delta(B_i*x-y_i)
%               + (m/2)*||x||_2^2 + indicator_D(x)
%   subject to  A*x = b,
%
% where D is the set of d-by-d complex density matrices represented by an
% isometric Hermitian vectorization, dimN=d^2, and h_delta is the symmetric
% Huber loss with delta=0.5. The measurement matrix B is scaled through gamma
% so that the smooth objective has upper curvature M and lower curvature m.
%
% A density matrix x_feasible is used to define b, while A is constructed so
% that the maximally mixed initial state x0 is infeasible. Hence the problem
% is feasible and x0 cannot be an optimal solution.

if nargin < 1 || isempty(N),    N = 50;      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 = 10;   end
if nargin < 6 || isempty(dimN), dimN = 64;   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 || M <= m
    error('test_fn_lin_quant_SDP4:BadCurvature', ...
        'The curvature inputs must satisfy 0 <= m < M.');
end

d = round(sqrt(dimN));
if d^2 ~= dimN || d < 2
    error('test_fn_lin_quant_SDP4:BadDimension', ...
        'dimN must be the square of an integer at least two.');
end
if N >= dimN - 1
    error('test_fn_lin_quant_SDP4:TooManyMeasurements', ...
        'N must be smaller than dimN-1.');
end

% Initialize.
rng(seed);
delta = 0.5;

% Generate a full-rank synthetic quantum state by mixing a random pure
% state with the maximally mixed state.
g = randn(d,1) + 1i * randn(d,1);
g = g / norm(g);
rho_feasible = 0.35 * (g * g') + 0.65 * eye(d) / d;
rho_feasible = (rho_feasible + rho_feasible') / 2;
x_feasible = hermitian_svec(rho_feasible);

% The initial point is the maximally mixed state.
rho0 = eye(d) / d;
params.x0 = hermitian_svec(rho0);
params.w0 = x_feasible;

% Generate dense traceless Hermitian observables with unit Frobenius norm.
B = zeros(N, dimN);
for i = 1:N
    G = randn(d) + 1i * randn(d);
    observable = (G + G') / 2;
    observable = observable - real(trace(observable)) * eye(d) / d;
    observable = observable / norm(observable, 'fro');
    B(i,:) = hermitian_svec(observable)';
end

% Scale the Huberized data-fitting term to realize the curvature pair (M,m).
sigma_max = norm(B, 2);
gamma = N * delta * (M - m) / sigma_max^2;

% Noisy expectation measurements. Since the observables have unit
% Frobenius norm and the true state is a moderate perturbation of the
% maximally mixed state, every residual at x0 lies in the quadratic
% Huber region, so the upper curvature M is attained at x0.
noise = 0.02 * (2 * rand(N,1) - 1);
y = B * x_feasible + noise;

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

% Euclidean topology induced by the isometric Hermitian vectorization.
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_quantum_huber_oracles( ...
    B, y, gamma, m, delta, norm_fn);
f_n = @(x) indicator_density_matrix(x, d);
prox_f_n = @(x,varargin) project_density_matrix(x, d);

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) quantum_huber_hessian(B, y, gamma, m, 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,b] = consistent_constraints(x_feasible, x0, dimM)
%CONSISTENT_CONSTRAINTS Ensure feasibility and make x0 infeasible.

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


function [f_s,grad_f_s] = cached_quantum_huber_oracles( ...
    B, y, gamma, m, delta, norm_fn)
%CACHED_QUANTUM_HUBER_ORACLES Smooth value and gradient with a one-point cache.

N = size(B, 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 = B * x - y;
            [loss,slope] = symmetric_huber(residual, delta);
            value = gamma * mean(loss) + (m / 2) * norm_fn(x)^2;
            g = gamma * (B' * slope) / N + m * x;

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


function [loss,slope] = symmetric_huber(t, delta)
%SYMMETRIC_HUBER Value and derivative of the symmetric Huber loss.

absolute_t = abs(t);
quadratic = (absolute_t <= delta);
linear = ~quadratic;

loss = zeros(size(t));
slope = zeros(size(t));
loss(quadratic) = t(quadratic).^2 / (2 * delta);
slope(quadratic) = t(quadratic) / delta;
loss(linear) = absolute_t(linear) - delta / 2;
slope(linear) = sign(t(linear));
end


function H = quantum_huber_hessian(B, y, gamma, m, delta, x)
%QUANTUM_HUBER_HESSIAN Piecewise Hessian of the smooth objective.

N = size(B, 1);
dimN = size(B, 2);
residual = B * x(:) - y;
active = double(abs(residual) < delta);
H = gamma * (B' * bsxfun(@times, B, active)) / (N * delta) ...
    + m * speye(dimN);
H = (H + H') / 2;
end


function value = indicator_density_matrix(x, d)
%INDICATOR_DENSITY_MATRIX Indicator of {rho >= 0 : trace(rho)=1}.

rho = hermitian_smat(x, d);
eigenvalues = eig(rho);
tolerance = 1e-10;
if min(real(eigenvalues)) >= -tolerance ...
        && abs(real(trace(rho)) - 1) <= tolerance
    value = 0;
else
    value = Inf;
end
end


function z = project_density_matrix(x, d)
%PROJECT_DENSITY_MATRIX Euclidean projection onto the density matrices.

rho = hermitian_smat(x, d);
rho = (rho + rho') / 2;
[U,S] = eig(rho);
eigenvalues = real(diag(S));
projected_eigenvalues = project_simplex(eigenvalues);
rho_projected = U * diag(projected_eigenvalues) * U';
rho_projected = (rho_projected + rho_projected') / 2;
z = hermitian_svec(rho_projected);
end


function z = project_simplex(x)
%PROJECT_SIMPLEX Euclidean projection onto {z >= 0 : sum(z)=1}.

x = x(:);
sorted_x = sort(x, 'descend');
cumulative = cumsum(sorted_x) - 1;
indices = (1:numel(x))';
rho = find(sorted_x - cumulative ./ indices > 0, 1, 'last');
theta = cumulative(rho) / rho;
z = max(x - theta, 0);
end


function x = hermitian_svec(H)
%HERMITIAN_SVEC Isometric real vectorization of a Hermitian matrix.

H = (H + H') / 2;
d = size(H, 1);
x = zeros(d^2, 1);
x(1:d) = real(diag(H));
root_two = sqrt(2);
index = d;

for i = 1:d-1
    for j = i+1:d
        index = index + 1;
        x(index) = root_two * real(H(i,j));
    end
end
for i = 1:d-1
    for j = i+1:d
        index = index + 1;
        x(index) = root_two * imag(H(i,j));
    end
end
end


function H = hermitian_smat(x, d)
%HERMITIAN_SMAT Inverse of hermitian_svec.

x = x(:);
H = complex(zeros(d, d));
H(1:d+1:end) = x(1:d);
root_two = sqrt(2);
num_upper = d * (d - 1) / 2;
real_part = x(d+1:d+num_upper);
imag_part = x(d+num_upper+1:end);
index = 0;

for i = 1:d-1
    for j = i+1:d
        index = index + 1;
        H(i,j) = (real_part(index) + 1i * imag_part(index)) / root_two;
        H(j,i) = conj(H(i,j));
    end
end
H = (H + H') / 2;
end


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