% function [Rxx, bsum , bsum_lin] = macmax(Esum, h, Lxu, N , cb)
%
% Simultaneous water-filling Esum MAC max rate sum (linear & nonlinear GDFE)
% The input is space-time domain h, and the user can specify a temporal
% block symbol size N (essentially an FFT size). 
%
% This program uses the CVX package
%
% the inputs are: 
% Esum The sum-user energy/SAMPLE scalar in time-domain. 
%      This will be increased by cb*N by this program.
%      Each user energy should be scaled by N/(N+nu)if there is cyclic prefix 
%      This energy is the trace of the corresponding user Rxx (u)
%      The sum energy is computed as the sum of the Eu components
%      internally. 
% h    The TIME-DOMAIN Ly x sum(Lx(u)) x N channel for all users
% Lxu  The number of antennas for each user 1 x U
% N    The FFT size (equally spaced over (0,1/T) at 1/(NT).
% cb   cb = 1 for complex, cb=2 for real baseband
%
% the outputs are:
% Rxx A block-diagonal psd matrix with the input autocorrelation for each
%     user on each tone. Rxx has size (sum(Lx(u)) x sum(Lx(u)) x N .
%     sum trace(Rxx) over tones and spatial dimensions equal the Eu 
%     FREQUENCY DOMAIN
% bsum the maximum rate sum, when cb=2, this is effectively over lower half
% of tones, or equivalently the 1/2*log2 form of data rates are summed
% bsum  bsum_lin - the maximum sum rate with a linear receiver
% 
%     b is an internal convergence (vector, rms) value, but not sum rate

function [Rxx, bsum, bsum_lin] = macmax(Esum, h, Lxu, N , cb)

H = fft(h, N, 3); % scaled so that N factor no longer needed in rate calculation
[Ly, Lx, ~] = size(H);
Esum=cb*N*Esum;
if numel(Lxu) > 1
    U = numel(Lxu);
    if sum(Lxu)~=Lx
        error('mismatch between sum of Lxu and Lx');
    end
elseif (Lx/Lxu)~= floor(Lx/Lxu)
    error('invalid Lxu');
else
    U = Lx/Lxu;
    Lxu = Lxu*ones(1,U);
end

idx_end = cumsum(Lxu);
idx_start = [1, idx_end(1:end-1)+1];
idx_exp_end = N*idx_end;
idx_exp_start = [1, idx_exp_end(1:end-1)+1];

Hcell=mat2cell(H, Ly, Lxu, ones(1,N));
Hexpand = zeros(N*Ly,N*Lx);
for u = 1:U
    Hexpand(:,idx_exp_start(u):idx_exp_end(u)) = blkdiag(Hcell{1,u,:});
end

Lxu_max = max(Lxu);
cvx_begin quiet
cvx_solver mosek
    variable Rxxun(Lxu_max,Lxu_max,N,U) hermitian semidefinite
    Rxxun_cell = reshape(num2cell(Rxxun, [1,2,3]), 1, U);
    for u = 1:U
        Rxxun_cell{u} = Rxxun_cell{u}(1:Lxu(u), 1:Lxu(u),:);
        Rxxun_cell{u} = reshape(num2cell(Rxxun_cell{u},[1,2]), 1, N);
    end
    Rxxun_cell = [Rxxun_cell{:}];
    Rxx = blkdiag(Rxxun_cell{:});
    maximize (1/cb*log2(exp(1))*log_det(eye(N*Ly) + Hexpand*Rxx*Hexpand'))
    subject to
        trace(Rxx) <= Esum;
cvx_end

bsum=cvx_optval;

Rxxun_cell = reshape(num2cell(Rxxun, [1,2,3]), U, 1);
for u = 1:U
    Rxxun_cell{u} = Rxxun_cell{u}(1:Lxu(u), 1:Lxu(u),:);
    Rxxun_cell{u} = reshape(num2cell(Rxxun_cell{u},[1,2]), 1, N);
end
Rxxun_cell = vertcat(Rxxun_cell{:}); % U*N cell array

Rxx = zeros(Lx,Lx,N);
for n=1:N
    Rxx(:,:,n) = blkdiag(Rxxun_cell{:,n});
end

bs=zeros(1,U);
bsum_lin=0;
for u=1:U
    indices=idx_start(u):idx_end(u);
    for n=1:N
     bs(u)=bs(u)+(1/cb)*(log2(det(eye(Ly)+H(:,:,n)*Rxx(:, ...
       :,n)*H(:,:,n)')) - log2(det(eye(Ly)+H(:,:,n)*Rxx(:, ...
       :,n)*H(:,:,n)'- H(:,indices,n)*Rxx(indices, ...
       indices,n)*H(:,indices,n)')));
    end
    bsum_lin=bsum_lin+real(bs(u));
end

end
