colored_noise = true;
I_shannon=[];
for cmplx = [0,1]
    cnt = 1;
    for M_PAM = [4,8,16]


        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);

        if colored_noise
            for i_SNR = 1:n_SNR
                % Transmitted signal for the given indices
                signal = X_(:, idx_tx);
    
                % Generate white Gaussian noise of the same size as the signal
                noise_white = randn(size(signal));
    
                % Define the filter coefficient for colored noise (adjust as needed)
                a = 0;  % A higher value gives more correlation
    
                % Filter the white noise to create colored noise.
                % Here we filter each row (dimension) independently.
                noise_colored = zeros(size(signal));
                noise_colored(1,:) = filter(1, [1, -a], noise_white(1,:));
                noise_colored(2,:) = filter(1, [1, -a], noise_white(2,:));
    
                % Compute signal and unscaled noise power for proper scaling.
                signal_power = mean(abs(signal(:)).^2);
                noise_power = mean(abs(noise_colored(:)).^2);
    
                % Convert desired SNR from dB to linear scale.
                SNR_linear = 10^(SNR_dB_values(i_SNR)/10);
    
                % Scale the colored noise so that signal_power / noise_power equals SNR_linear.
                scaling_factor = sqrt(signal_power / (SNR_linear * noise_power));
                noise_colored_scaled = scaling_factor * noise_colored;
    
                % Received signal is the sum of the signal and the scaled colored noise.
                r_colored = signal + noise_colored_scaled;
    
                % For SNR measurement, compute the noise actually added.
                measured_noise = r_colored - signal;
                snr_meas_(i_SNR) = snr(signal(1,:) + 1i*signal(2,:), ...
                    measured_noise(1,:) + 1i*measured_noise(2,:));
    
                % Plotting the transmitted and received signal
                clf;
                hold on;
                xlim([-4, 4]);
                ylim([-4, 4]);
                scatter(signal(1,:), signal(2,:), 3, "black", 'o', 'DisplayName','Tx Constellation');
                scatter(r_colored(1,:), r_colored(2,:), 1, "red", 'x', 'DisplayName', 'Colored Noise Mapping');
                drawnow;
    
                % Compute Mutual Information (or any other metric) with the new channel
                MI_awgn_(i_SNR) = air(X_, r_colored, idx_tx);
            end
        else


            for i_SNR=1:n_SNR
    
    
                s_ = sqrt(SNR_values(i_SNR)*var_w) * X_(:, idx_tx);
                % r_awgn_ = s_ + w;
    
                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,:));
    
                %%%%%%%%%%%%
                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');
    
    
                scatter(r_awgn_(1,:),r_awgn_(2,:),1,"red",'x','DisplayName','Matlab Mapping');
    
                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
        end

        figure(3);
        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


        cols = [    0.9047    0.1918    0.1988
            0.2941    0.5447    0.7494
            0.3718    0.7176    0.3612
            1.0000    0.5482    0.1000
            0.8650    0.8110    0.4330
            0.6859    0.4035    0.2412];
        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');

        cnt = cnt+1;

    end
end