function [prob,C,X,Y,lam] = prob_dio(rho, m, tol)
% prob_dio(rho,m,tol) computes the maximal probability of distillation of an m-dimensional maximally coherent state under DIO with tolerance tol
% rho: dxd density matrix or normalised state vector
% m: target state dimension
% tol: error tolerance 0<=tol<=1

d=size(rho,2);
if d==1
    rho=rho*rho';
    d=size(rho,1);
end
rho=(rho+rho')/2;

cvx_begin sdp quiet

    variable C(d,d) hermitian semidefinite
    dual variable X
    dual variable Y
    dual variable lam

    maximize ( m*trace(diag(diag(C))*rho) )

    Y : C - m*diag(diag(C)) <= 0
    X : m*diag(diag(C)) <= eye(d)

    lam : trace(C*rho) >= (1-tol)*trace(m*diag(diag(C))*rho)

cvx_end

prob = cvx_optval;

end