function sum_rate = Digital_MIMO_precoding(K,M,H_k,Psi_k,v,P_T,sigma_2,lambda)

Q = 3;

sum_rate_all = [];

for qq = 1:Q

    [sum_rate,SINR_k] = Capacity_MIMO(K,M,H_k,v,sigma_2);

    roh_k = SINR_k;   %更新gamma

    Psi_k = Psi_k_update(K,M,H_k,roh_k,v,sigma_2);    %更新eps

    v = v_optimize(K,M,H_k,roh_k,Psi_k,P_T,lambda);


end

end

function [sum_rate,SINR_k] = Capacity_MIMO(K,M,H_k,v,sigma_2)
v_k = zeros(3*M,K);
for k = 1:K
    v_k(:,k) = v(3*M*(k-1)+1:1:3*M*k,1);
end

SINR_k = zeros(K,1);
sum_rate = 0;
for k = 1:K

    temp1 = squeeze(H_k(:,:,k)*v_k(:,k));
    temp2 = zeros(3,3);
    for jj = 1:K
        temp2 = temp2 + (H_k(:,:,k)*v_k(:,jj))*((H_k(:,:,k))*v_k(:,jj))';
    end
    temp2 = temp2 - (H_k(:,:,k)*v_k(:,k))*(H_k(:,:,k)*v_k(:,k))';
    temp2 = temp2 + eye(3)*sigma_2;
%     temp2 = temp2^(-1);

    SINR_k(k) = temp1' * (temp2 \ temp1);
    sum_rate = sum_rate + log2(1+SINR_k(k));
end

sum_rate = abs(sum_rate);
end
function Psi_k = Psi_k_update(K,M,H_k,roh_k,v,sigma_2)
v_k = zeros(3*M,K);
for k = 1:K
    v_k(:,k) = v(3*M*(k-1)+1:1:3*M*k,1);
end

Psi_k = zeros(3,K);

for k = 1:K
    temp1 = H_k(:,:,k)*v_k(:,k);
    temp2 = zeros(3,3);

    for jj = 1:K
        temp2 = temp2 + (H_k(:,:,k)*v_k(:,jj))*((H_k(:,:,k))*v_k(:,jj))';
    end
    temp2 = temp2 + eye(3)*sigma_2;

    Psi_k(:,k) = sqrt(1+roh_k(k))*temp2\temp1;
end

end
function v = v_optimize(K,M,H_k,roh_k,Psi_k,P_T,lambda)
beta = zeros(3*M*K,1);
A = zeros(3*M,3*M);
for k = 1:K
    beta((k-1)*3*M+1:1:k*3*M,1) = sqrt(1+roh_k(k)) * (squeeze(H_k(:,:,k)))' * Psi_k(:,k);
    A = A + (squeeze(H_k(:,:,k)))' * Psi_k(:,k)*((squeeze(H_k(:,:,k)))' * Psi_k(:,k))';
end
A  = kron(A,eye(K));
A = (A+A')/2;
% v = A^(-1)*beta;
% v = v/norm(v)*sqrt(P_T/(lambda/2)^2);

% cvx_begin quiet
% cvx_precision low
% variable v(3*M*K,1) complex;
% minimize((v')*A*v-2*real((beta')*v))
% subject to
% norm(v) <= sqrt(P_T/(lambda/2)^2);
% cvx_end

left_zeta = 0;
right_zeta = 10;

v = (A)\beta;

PP =  sqrt(P_T/(lambda^2/4/pi));

if norm(v) <= PP
    return;

else
    v = (A+right_zeta*eye(3*M*K))\beta;

    while(norm(v)>PP)
        right_zeta = 10* right_zeta;
        v = (A+right_zeta*eye(3*M*K))\beta;
        if norm(v)<=PP
            left_zeta = right_zeta/10-1;
        end
    end

    while(abs(norm(v)-PP)>0.001*PP)
        center_zeta = (left_zeta + right_zeta)/2;
        v = (A + center_zeta*eye(3*M*K))\beta;
        if norm(v)>PP
            left_zeta = center_zeta;
        else
            right_zeta = center_zeta;
        end
    end

    v = (A + center_zeta*eye(3*M*K))\beta;    
end


end

