%% [ent, ent0, W] = dist_ent (rho, tol, op, dim)
% Computes the one-shot distillable entanglement under a given class of operations.
% 
% Required arguments:
% *  rho: d*d density matrix or d*1 state vector
% *  tol: desired error tolerance in distillation fidelity, 0<=tol<1
%
% 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:
% *  ent: maximum size of maximally entangled state which can be obtained from rho under
%         operations op up to tolerance tol, i.e. log2(ent) is the one-shot distillable entanglement
%            The function makes an attempt to correct for CVX's numerical accuracy
%            (the default accuracy being within sqrt(eps) of the optimal value).
%            One should consult the exact value ent0 if needed.
% *  ent0: the non-floored value of the optimisation problem such that ent = floor(1/ent0)
% *  W: optimal witness in the optimisation problem (see Prop. 4 of the manuscript)
%       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 [ent, ent0, W] = dist_ent(rho, tol, 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
     
    if min(dimrho) ~= 1
        [~,~,v] = svd(frho);
        rho = v(:,1);
    end
    
    [xi,U,V] = SchmidtDecomposition(rho, dim, -1);
    
    cvx_begin quiet

        variable x(d) nonnegative
        
        norm(x) <= 1
        x'*xi >= sqrt(1-tol);
        
        minimize max(x) 

    cvx_end
    
    if nargout > 2
        
        % construct optimal witness
        W = zeros(drho,1);

        for j=1:d
           W = W + x(j) * kron(U(:,j), V(:,j));
        end

        if min(dimrho) ~= 1
            W = W*W';
        end  
    end
    
    ent0 = cvx_optval^2;
        
else % 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==""
            t = norm(PartialTranspose(W,2,dim));
        elseif lower(op)=="pptp"
            variable X(drho,drho) hermitian semidefinite
            W <= X
            t = lambda_max(PartialTranspose(X,2,dim));
        elseif lower(op)=="rains"
            variable X(drho,drho) hermitian semidefinite
            W <= X
            t = norm(PartialTranspose(X,2,dim));
        else
            error('Unsupported type of operations. Allowed choices: ppt, pptp, rains.');
        end
        
        trace(rho*W) >= 1-tol;
        
        minimize t

    cvx_end

    ent0 = cvx_optval;

end

ent = floor(1/ent0);
prec = cvx_precision;
if abs(ent0 - 1/(ent+1)) <= prec(1) % to account for CVX precision
    ent = ent+1;
end

end