clear;
I_shannon=[];
for cmplx = [1,0]
    cnt = 1;
    for M_PAM = [2,4,8,16]

        % close all;
        usegarcia = 0;

        complex_constellation = cmplx;

        %%% Simulation parameters
        % M_PAM=2; % QAM size
        var_w=1; % noise variance
        SNR_dB_values=(-5):1:35; % SNR values in dB
        N=8000; % number of Monte-Carlo points

        %%% Derived parameters
        n_SNR=length(SNR_dB_values);
        SNR_values=10.^(SNR_dB_values/10);

        %%% Generation of QAM constellation with power 1
        M=sqrt(M_PAM); % number of points per real dimension

        if complex_constellation
            X_ = qammod(0:M_PAM-1,M_PAM,"gray");
        else
            X_ = pammod(0:M_PAM-1,M_PAM,0,'gray');
        end

        X_ = X_ ./ rms(unique(X_));
        X_=[real(X_); imag(X_)];

        %%%%%%%%%%%%
        figure(10);
        clf
        hold on
        scatter(X_(1,:),X_(2,:),15,'red','x','DisplayName','Matlab');
        %%%%%%%%%%%%

        %%% Uniformly choose transmit indices
        idx_tx=randi(M_PAM, [1, N]);

        %%% AWGN noise
        w=sqrt(var_w)*(randn([2, N]));

        %%% Loop over the SNR values

        MI_awgn=zeros(1, n_SNR);
        fig = figure(2);
        for i_SNR=1:n_SNR

            if usegarcia


                %scale transmitted points acc. to SNR condition (this is not correct I think, the papaer also show diff results)
                % s=sqrt(SNR_values(i_SNR)*var_w)*X(:, idx_tx);

                %scale nosie acc. to snr condition
                noise_power = mean(abs(X_(1, idx_tx)+1i*X_(2, idx_tx)).^2) / SNR_values(i_SNR); % Noise power
                w_ = sqrt(noise_power) .* w / sqrt(2);

                % AWGN channel
                r_awgn = X(:, idx_tx)+w_;
                % snr_meas(i_SNR) = snr(s(1,:)+1i*s(2,:),w(1,:)+1i*w(2,:));
                snr_meas(i_SNR) = snr(X_(1, idx_tx)+1i*X_(2, idx_tx),w_(1,:)+1i*w_(2,:));
            end

            s_ = sqrt(SNR_values(i_SNR)*var_w) * X_(:, idx_tx);
            % r_awgn_ = s_ + w;

            if cmplx
                r_awgn_ = awgn(X_(:, idx_tx),SNR_dB_values(i_SNR),"measured",10);
                w_awgn_matlab = r_awgn_ - X_(:,idx_tx);
                snr_meas_(i_SNR) = snr(X_(1, idx_tx)+1i*X_(2, idx_tx) , w_awgn_matlab(1,:)+1i*w_awgn_matlab(2,:));
            else
                r_awgn_ = awgn(X_(1, idx_tx),SNR_dB_values(i_SNR),"measured",10);
                w_awgn_matlab = r_awgn_ - X_(1,idx_tx);
                snr_meas_(i_SNR) = snr(X_(1, idx_tx) , w_awgn_matlab(1,:));
            end

            %%%%%%%%%%%%
            clf
            hold on
            xlim([-4, 4]);
            ylim([-4, 4]);
            %received signal
            scatter(X_(1, idx_tx),X_(2, idx_tx),3,"black",'o','DisplayName','Tx Constellation');


            if cmplx
                scatter(r_awgn_(1,:),r_awgn_(2,:),1,"red",'x','DisplayName','Matlab Mapping');
            else
                scatter(r_awgn_,zeros(size(r_awgn_)),1,"red",'x','DisplayName','Matlab Mapping');
            end
            drawnow

            MI_awgn_(i_SNR)=air(X_, r_awgn_, idx_tx);
            % The following also works but is slower
            %     MI_awgn(i_SNR)=mi_cg(s, r_awgn);
        end

        figure(4);
        hold on

        if isempty(I_shannon)
            I_shannon=log2(1+db2pow(snr_meas_));
            plot(snr_meas_, I_shannon, '-', 'DisplayName', 'log_2 (1+SNR)','Color','black');
        end

        if usegarcia
            plot(SNR_dB_values, MI_awgn, '--', 'DisplayName', ['GARCIA ',num2str(M_PAM) '-QAM, AWGN']);
        end

        cols = linspecer(6);
        if complex_constellation
            plot(snr_meas_, MI_awgn_, ':', 'DisplayName', ['',num2str(M_PAM) '-QAM, AWGN'],'LineWidth',1,'Color',cols(cnt,:));
        else
            plot(snr_meas_, MI_awgn_, '-', 'DisplayName', ['',num2str(M_PAM) '-PAM, AWGN'],'LineWidth',1,'Color',cols(cnt,:));
        end

        hold off;
        ylim([0,8]);
        xlim([min(SNR_dB_values) max(SNR_dB_values)])
        xlabel('SNR (dB)'); ylabel('Achievable rate (bits/complex dimension)');
        title(['Achievable rate of ' num2str(M_PAM) '-QAM in AWGN']);
        legend('Location', 'NorthWest');
        yline(log2(M_PAM),'LineStyle',':','HandleVisibility','off','Color','black');

        autoArrangeFigures;
        cnt = cnt+1;

    end
end