function X=diag(X,k)
%DIAG (overloaded)

% Author Johan Löfberg 
% $Id: diag.m,v 1.7 2010-01-13 14:18:15 joloef Exp $  

n = X.dim(1);
m = X.dim(2);

if nargin == 1
    k = 0;
end

if min([n m])==1
    Y = X;
	diagX = diag(X.basis(:,1),k);
	Y.basis = diagX(:);
	for i = 1:length(X.lmi_variables)
		diagX = diag(X.basis(:,i+1),k);
		Y.basis(:,i+1) = diagX(:);
    end
	Y.lmi_variables = X.lmi_variables;
	Y.dim(1) = size(diagX,1);
	Y.dim(2) = size(diagX,2);
    X=Y;
else
    index = diag(reshape(1:n*m,n,m),k);
    X.basis = X.basis(index,:);
	X.dim(1) = length(index);
	X.dim(2) = 1;
    X = clean(X);
end
% Reset info about conic terms
if isa(X,'sdpvar')
    X = flush(X);
    X.conicinfo = [0 0];
end