function air = air_garcia_implementation(x,r,idx_tx,Px,M_training)

MAX_MEMORY = 200e6; % maximum allowed size for a matrix

if nargin == 3
    Px = [];
    M_training = [];
end

if nargin == 4
    M_training = [];
end


% if input is complex, separate into real and imaginary parts
if any(imag(x(:))~=0) || any(imag(r(:))~=0)
    x = [real(x); imag(x)];
    r = [real(r); imag(r)];
end

D = size(x, 1); % D = 2 if complex x
N = size(x, 2); % number of constellation points
M = size(r, 2); % number of samples

% set default training set size
if isempty(M_training)
    M_training = ceil(0.3*M);
end

M_testing = M - M_training;

% Training: estimate parameters of the conditionally Gaussian model
% sort according to transmit index
[idx_tx_training, idx_sort] = sort(idx_tx(1:M_training));
r_training = r(:, idx_sort);
i_bounds = zeros(1, N+1);

% compute conditional means and covariance matrices
C_n = zeros(D, D, N);
det_n = zeros(1, N);

for n=1:N
    % find how many times x(:, n) was transmitted and update i_bounds
    N_current_x = find(idx_tx_training((i_bounds(n)+1):end)==n, 1, 'last');
    if isempty(N_current_x), N_current_x=0; end
    i_bounds(n+1) = i_bounds(n) + N_current_x;

    if N_current_x > 0
        % Compute mu_n=E[Y|X=x_n] according to Eq. (14) and store it in
        % x(:, n) to save space
        x(:, n) = sum(r_training(:, (i_bounds(n)+1):i_bounds(n+1)), 2)/(i_bounds(n+1)-i_bounds(n));

        % compute C_n=cov[Y|X=x_n] according to Eq. (15)
        r_meanfree = r_training(:, (i_bounds(n)+1):i_bounds(n+1)) - x(:, n);
        C_n(:, :, n) = (r_meanfree*r_meanfree')/(i_bounds(n+1)-i_bounds(n));
        % store also the determinant of C(:, :, n)
        det_n(n) = det(C_n(:, :, n));

        % if the determinant is 0, or if the matrix is badly conditioned,
        % regularize by adding a small identity matrix. Note that we do
        % need the check for 0 determinant, in case a cloud has exactly 0
        % variance according to the training set
        if det_n(n)==0 || cond(C_n(:, :, n))>1e16
            C_n(:, :, n) = C_n(:, :, n) + 5 * eps * eye(D);
            det_n(n) = (5*eps)^D;
        end

    end
end

% uniform input pmf Px if not provided
if isempty(Px)
    Px = repmat(1/N, [1, N]);
end


% extract testing set and sort it according to transmit index
[idx_tx_testing, idx_sort] = sort(idx_tx((M_training+1):M));
r_testing = r(:, M_training+idx_sort);

% computation of h(Y|X)
h_Y_X = 0;
i_bounds_testing = zeros(1, N+1);
% loop over constellation points to compute h(Y|X)
for n = 1:N
    % find how many times x(:, n) was transmitted and update
    % i_bounds_testing
    N_current_x = find(idx_tx_testing((i_bounds_testing(n)+1):end)==n, 1, 'last');
    if isempty(N_current_x), N_current_x=0; end
    i_bounds_testing(n+1) = i_bounds_testing(n) + N_current_x;
    % add the corresponding contribution to the mutual information (two
    % first lines of Eq. (17)). This, together with
    % D/2*log2(2*pi) after the end of the loop, gives h(Y|X)
    h_Y_X = h_Y_X + N_current_x * log2(det_n(n))/2+...
        sum(sum(conj(r_testing(:, (i_bounds_testing(n)+1):i_bounds_testing(n+1))-x(:, n)).*(C_n(:, :, n)\(r_testing(:, (i_bounds_testing(n)+1):i_bounds_testing(n+1))-x(:, n)))))/2/log(2);

    
end
h_Y_X = D/2*log2(2*pi) + h_Y_X/M_testing;

% When computing log(py), we might run out of memory. If necessary, we
% doe the computation in blocks
logpy = zeros(1, M_testing);
BLOCK_SIZE = floor(MAX_MEMORY/N);
N_blocks = ceil(M_testing/BLOCK_SIZE);

% loop over blocks of symbols. This loop can be replaced by parfor to allow
% parallel computation
for i_block = 1:N_blocks

    logpy_cur = zeros(1, M_testing);

    % beginning of block
    i_start = (i_block-1) * BLOCK_SIZE + 1;
    % end of block
    i_end = min(M_testing, i_block*BLOCK_SIZE);
    % block size
    current_block_size = i_end-i_start+1;

    % compute exponents of third line of (17)
    exponents = zeros(N, current_block_size);
    for n = 1:N
        exponents(n, :) = -log(det_n(n))/2-real(sum(conj(r_testing(:, i_start:i_end)-x(:, n)).*(C_n(:, :, n)\(r_testing(:, i_start:i_end)-x(:, n))), 1))/2;
        %sum über 2 einträge von r
    end

    % compute third line of Eq. (17). Use a custom function
    % that computes log(sum(exp(x))) avoiding overflow errors
    logpy_cur(i_start:i_end) = math_logsumexp(log(Px(:))+exponents, 1);
    logpy = logpy + logpy_cur;
end

% output entropy h(Y)
h_Y = D/2*log2(2*pi) - mean(logpy)/log(2);%log basis change

% compute mutual information
air = h_Y - h_Y_X;


end


function [y] = math_logsumexp(x, dim)
%[y] = math_logsumexp(x, dim)
%  Computes log(sum(exp(x), dim)), avoiding overflow errors when one of the
%  x is large.

if nargin<2 || isempty(dim)
    m = max(x);
    y = m + log(sum(exp(x-m)));
else
    m = max(x, [], dim);
    y = m + log(sum(exp(x-m), dim));
end
end