function [Y, problem, S] = elliptope_SDP_complex(A, p, Y0)
% Solver for complex semidefinite programs (SDP's) with unit diagonal.
% 
% function [Y, problem, S] = elliptope_SDP_complex(A)
% function [Y, problem, S] = elliptope_SDP_complex(A, p)
% function [Y, problem, S] = elliptope_SDP_complex(A, p, Y0)
%
% A is a Hermitian matrix of size n.
%
% This function uses a local optimization method in Manopt to solve the SDP
%
%   min_X trace(A*X) s.t. diag(X) = 1, X is complex, positive semidefinite.
%
% In practice, the Hermitian matrix X of size n is parameterized as
% X = Y*Y', where Y has size n x p. By default, p is taken large enough
% (that is, sqrt(n)) to ensure that there exists an optimal X whose rank is
% smaller than p. This ensures that the SDP is equivalent to the new
% problem in Y:
%
%   min_Y  trace(Y'*A*Y)  s.t.  diag(Y*Y') = 1, Y complex
%
% The constraints on Y require each row of Y to have unit norm, which is
% why Manopt is appropriate software to solve this problem. An optional
% initial guess can be specified via the input Y0.
%
% See the paper below for theory, specifically, for a proof that, for
% almost all A, second-order critical points of the problem in Y are
% globally optimal. In other words: there are no local traps in Y, despite
% non-convexity.
%
% Outputs:
%
%       Y: is the best point found (an nxp matrix with unit norm rows.)
%          To find X, form Y*Y' (or, more efficiently, study X through Y.)
% 
%       problem: is the Manopt problem structure used to produce Y.
% 
%       S: is a dual optimality certificate (a Hermitian matrix of size n,
%          sparse if A is sparse). The optimality gap (in the cost
%          function) is at most n*min(eig(S)), for both Y and X = Y*Y'.
%          Hence, if min(eig(S)) is close to zero, Y is close to globally
%          optimal. This can be computed via eigs(S, 1, 'SR').
% 
% Paper: https://arxiv.org/abs/1606.04970
%
% @inproceedings{boumal2016bmapproach,
%   author  = {Boumal, N. and Voroninski, V. and Bandeira, A.S.},
%   title   = {The non-convex {B}urer-{M}onteiro approach works on smooth semidefinite programs},
%   booktitle={Neural Information Processing Systems (NIPS 2016)},
%   year    = {2016}
% }
% 
% See also: maxcut elliptope_SDP

% This file is part of Manopt: www.manopt.org.
% Original author: Nicolas Boumal, Oct. 21, 2016
% Contributors:
% Change log:
%
%    Xiaowen Jiang Aug. 20, 2021
%       Added AD to compute the egrad and the ehess

    % If no inputs are provided, since this is an example file, generate
    % a random complex matrix. This is for illustration purposes only.
    if ~exist('A', 'var') || isempty(A)
        n = 100;
        A = randn(n) + 1i*randn(n);
        A = (A+A')/sqrt(2*n);
    end

    n = size(A, 1);
    assert(n >= 2, 'A must be at least 2x2.');
    assert(size(A, 2) == n, 'A must be square.');
    
    % Force A to be Hermitian
    A = (A+A')/2;
    
    % By default, pick a sufficiently large p (number of columns of Y).
    if ~exist('p', 'var') || isempty(p)
        p = floor(sqrt(n)+1);
    end
    
    assert(p >= 1 && p == round(p), 'p must be an integer >= 1.');

    % Pick the manifold of complex n-by-p matrices with unit norm rows.
    manifold = obliquecomplexfactory(n, p, 'rows');
    
    problem.M = manifold;
    
    
    % These three, quick commented lines of code are sufficient to define
    % the cost function and its derivatives. This is good code to write
    % when prototyping. Below, a more advanced use of Manopt is shown,
    % where the redundant computation A*Y is avoided between the gradient
    % and the cost evaluation.
    % % problem.cost  = @(Y) .5*sum(sum(real((A*Y).*conj(Y))));
    % % problem.egrad = @(Y) A*Y;
    % % problem.ehess = @(Y, Ydot) A*Ydot;
    
    % Products with A dominate the cost, hence we store the result.
    % This allows to share the results among cost, grad and hess.
    % This is completely optional.
    function store = prepare(Y, store)
        if ~isfield(store, 'AY')
            AY = A*Y;
            store.AY = AY;
            store.diagAYYt = sum(real(AY .* conj(Y)), 2);
        end
    end
    
    % Define the cost function to be /minimized/.
    problem.cost = @cost;
    function [f, store] = cost(Y, store)
        store = prepare(Y, store);
        f = .5*sum(store.diagAYYt);
    end

    % Define the Riemannian gradient.
    problem.grad = @grad;
    function [G, store] = grad(Y, store)
        store = prepare(Y, store);
        G = store.AY - bsxfun(@times, Y, store.diagAYYt);
    end

    % If you want to, you can specify the Riemannian Hessian as well.
    problem.hess = @hess;
    function [H, store] = hess(Y, Ydot, store)
        store = prepare(Y, store);
        SYdot = A*Ydot - bsxfun(@times, Ydot, store.diagAYYt);
        H = manifold.proj(Y, SYdot);
    end

    % An alternative way to compute the egrad and the ehess is to use 
    % automatic differentiation provided in the deep learning toolbox
    % (slower). AD does not support complex numbers if the Matlab version
    % is R2021a or earlier. The cost function should be defined differently
    % In this case. See complex_example_AD.m and manoptADhelp.m for more
    % information.
    % problem.cost = @cost_AD;
    %    function f = cost_AD(Y)
    %        AY = cprod(A, Y);
    %        diagAYYt = csum(creal(cdottimes(AY, cconj(Y))), 2);
    %        f = .5*csum(diagAYYt);
    %    end
    % Call manoptAD to automatically obtain egrad and ehess:
    % problem = manoptAD(problem);

    % If the version of Matlab installed is R2021b or later, specify the 
    % cost function in the normal way and call manoptAD. 
    % problem.cost = @cost_AD;
    %    function f = cost_AD(Y)
    %        AY = A*Y;
    %        diagAYYt = sum(real(AY .* conj(Y)), 2);
    %        f = .5*sum(diagAYYt);
    %    end
    % problem = manoptAD(problem);


    % If no initial guess is available, tell Manopt to use a random one.
    if ~exist('Y0', 'var') || isempty(Y0)
        Y0 = [];
    end

    % Call your favorite solver.
    opts = struct();
    opts.verbosity = 0;      % Set to 0 for no output, 2 for normal output
    opts.maxinner = 500;     % maximum Hessian calls per iteration
    opts.tolgradnorm = 1e-6; % tolerance on gradient norm
    Y = trustregions(problem, Y0, opts);
    
    % If required, produce an optimality certificate.
    if nargout >= 3
        S = A - spdiags(sum(real((A*Y).*conj(Y)), 2), 0, n, n);
    end

end
