%% [fid, W] = dist_fid (rho, m, op, dim)
% Computes the fidelity of entanglement distillation under a given class of operations.
% 
% Required arguments:
% *  rho: d*d density matrix or d*1 state vector
% *  m: desired size of output maximally entangled state
%
% Optional arguments:
% *  op: class of operations (default: ppt)
%        allowed choices: ppt (PPT operations), pptp ((PPT+)-preserving), rains (Rains-preserving)
%        disregarded when input is pure (all classes including LOCC give the same result)
% *  dim: vector [dA,dB] of dimensions of subsystems (default: both systems of equal dimension)
%
% Output:
% *  fid: fidelity of distillation
% *  W: optimal witness such that fid = trace(rho*W)
%       if the input is provided as a vector, W is also returned as one
%
% Requires: CVX (http://cvxr.com/cvx), QETLAB (http://qetlab.com)
%
% Based on:
%   Bartosz Regula, Kun Fang, Xin Wang, and Mile Gu, arXiv:1906.01648

function [fid, W] = dist_fid(rho, m, op, dim)

dimrho = size(rho);
drho = max(dimrho);

if nargin < 4 % assume dimensions of subsystems are equal
    dim = [sqrt(drho),sqrt(drho)];
    d = sqrt(drho);
else
    d = min(dim);
end

frho = full(rho);

if rank(frho) == 1 % pure state case
    
    % obtain Schmidt decomposition
    
    if min(dimrho) ~= 1
        [~,~,v] = svd(frho);
        rho = v(:,1);
    end
    
    [xi,U,V] = SchmidtDecomposition(rho, dim, -1);
    
    % compute m-distillation norm of Schmidt vector
    
    xi = sort(xi,'descend');
    k = 1;
    minsum = sum(xi(m:d).^2);

    for j=2:m
        jsum = sum(xi(m-j+1:d).^2)/j;
        if jsum < minsum
            minsum = jsum;
            k = j;
        end
    end
    
    fid = sum(xi(1:m-k))+k*sqrt(minsum);
    fid = fid^2/m;
    
    if nargout > 1
        
        % construct optimal witness
        
        wxi = xi./sqrt(minsum);
        [~,wp] = maxk(xi,m-k);
        wxi(wp) = 1;
        W = zeros(drho,1);
        for j=1:d
           W = W + wxi(j) * kron(U(:,j), V(:,j));
        end
        W = W/sqrt(m);

        if min(dimrho) ~= 1
            W = W*W';
        end
        
    end
        
else % general mixed state case

    I = eye(drho);
    rho = (rho+rho')/2; % avoids numerical problems

    cvx_begin sdp quiet

        variable W(drho,drho) hermitian semidefinite

        W <= I

        if nargin<3 || lower(op)=="ppt" || op==""
            -I/m <= PartialTranspose(W,2,dim) <= I/m
        elseif lower(op)=="pptp"
            variable X(drho,drho) hermitian semidefinite
            W <= X
            PartialTranspose(X,2,dim) <= I/m
        elseif lower(op)=="rains"
            variable X(drho,drho) hermitian semidefinite
            W <= X
            -I/m <= PartialTranspose(X,2,dim) <= I/m
        else
            error('Unsupported type of operations. Allowed choices: ppt, pptp, rains.');
        end

        maximize trace(rho*W)

    cvx_end

    fid = cvx_optval;

end