% This script calculates numerically the constrained capacity of different
% modulation schemes considering equiprobable signals, coherent detection
% and soft output, in an AWGN channel.
% The following modulation schemes are supported:
% OOK, BPSK, QPSK, 8-PSK, 4-PAM, 16-QAM, 64-QAM, 4-ASK (non-negative)
%
% The results are saved in file 'mutualInf-MOD.mat'

%
%   Author: André Noll Barreto, Universidade de Brasília
%   E-mail: andrebarreto@ene.unb.br
%   History:
%       v1.0 Jan 2014 - created
%
%   Copyright (C) 2014, André Noll Barreto
%
%    This program is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    This program is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with this program.  If not, see <http://www.gnu.org/licenses/>.

modTypes.ook = 0;
modTypes.bpsk = 1;
modTypes.qpsk = 2;
modTypes.psk8 = 3;
modTypes.qam16 = 4;
modTypes.pam4 = 5;
modTypes.qam64 = 7;
modTypes.ask4 = 8;

% parameters
mod = modTypes.ask4; % modulation scheme
EsN0dB = -10:20; % signal-to-noise ratio


integPoints = 4000; % number of integration points
integConf = .999; % integral limits are defined based on integConf ratio of
                  % noise samples

mutualInf = zeros(size(EsN0dB));

switch mod
    case modTypes.ook
        A = [0,1];
        modName = 'OOK';
        dimensions = 1;
        symmetry = false;
        Asym = A;
        Wsym = [1 1];
        
    case modTypes.ask4
        A = [0,1,2,3]/ sqrt(3.5);
        modName = '4-ASK';
        dimensions = 1;
        symmetry = false;
        Asym = A;
        Wsym = [1 1 1 1];
        
    case modTypes.bpsk
        A = [-1,1];
        modName = 'BPSK';
        dimensions = 1;
        symmetry = true;
        Asym = -1;
        Wsym = 2;
    case modTypes.qpsk
        A = [-1-1i, -1+1i, 1-1i, 1+1i] / sqrt(2);
        modName = 'QPSK';
        dimensions = 2;
        symmetry = true;
        Asym = [-1-1i]/ sqrt(2);
        Wsym = 4;
    case modTypes.pam4
        A = [-3, -1, 1, 3] / sqrt(5);
        modName = '4-PAM';
        dimensions = 1;
        symmetry = true;      
        Asym = [-3 -1] / sqrt(5);
        Wsym = [2 2];
    case modTypes.psk8
        A = exp(1i * 2 * pi * [0, 1/8, 2/8, 3/8, 4/8, 5/8, 6/8, 7/8 ]);
        modName = '8-PSK';
        dimensions = 2;
        symmetry = true;
        Asym = exp(1i * 2 * pi * 0);
        Wsym = 8;
    case modTypes.qam16
        A = [-1-1i, -1+1i, 1-1i, 1+1i, -1-3i, -1+3i, 1-3i, 1+3i, ...
            -3-1i, -3+1i, 3-1i, 3+1i, -3-3i, -3+3i, 3-3i, 3+3i]/ sqrt(10);
        modName = '16-QAM';
        dimensions = 2;
        symmetry = true;
        Asym = [-1-1i, -1-3i,  -3-3i]/ sqrt(10);
        Wsym = [4, 8, 4];

        
    case modTypes.qam64
        A = [-1-1i, -1+1i, 1-1i, 1+1i, -1-3i, -1+3i, 1-3i, 1+3i, ...
             -1-5i, -1+5i, 1-5i, 1+5i, -1-7i, -1+7i, 1-7i, 1+7i, ...
             -3-1i, -3+1i, 3-1i, 3+1i, -3-3i, -3+3i, 3-3i, 3+3i, ...
             -3-5i, -3+5i, 3-5i, 3+5i, -3-7i, -3+7i, 3-7i, 3+7i, ...
             -5-1i, -5+1i, 5-1i, 5+1i, -5-3i, -5+3i, 5-3i, 5+3i, ...
             -5-5i, -5+5i, 5-5i, 5+5i, -5-7i, -5+7i, 5-7i, 5+7i, ...
             -7-1i, -7+1i, 7-1i, 7+1i, -7-3i, -7+3i, 7-3i, 7+3i, ...
             -7-5i, -7+5i, 7-5i, 7+5i, -7-7i, -7+7i, 7-7i, 7+7i, ...
            ]/ sqrt(42);
        modName = '64-QAM';
        dimensions = 2;
        symmetry = true;
        Asym = [-1-1i, -1-3i, -1-5i, -1-7i, -3-3i, -3-5i, -3-7i, -5-5i, -5-7i, -7-7i] / sqrt(42);
        Wsym = [4,     8,     8,     8,     4,     8,     8,     4,     8,     4];
        
end


fileName = ['mutualInf-' modName '.mat'];

%scatter(real(A), imag(A))

Es = mean(abs(A).^2);
EsN0vec = 10.^(EsN0dB/10);

M = length(A);

maxr = max(real(A));
maxi = max(imag(A));
minr = min(real(A));
mini = min(imag(A));

fprintf('Calculating constrained capacity for %s:\n', modName)

for EsN0ind = 1:length(EsN0vec)
    
    fprintf('\tEs/N0 = %d dB\n', EsN0dB(EsN0ind));
    
    EsN0 = EsN0vec(EsN0ind);
    N0 = Es/EsN0;
    sigma = sqrt(N0/2);
    
    integLim = icdf('normal',integConf, 0, sigma);
    integStep = (max(maxr - minr, maxi - mini) + 2*integLim) ...
                 / (integPoints + 1);
    minrtmp = floor((-integLim+minr)/integStep)*integStep;
    maxrtmp = ceil((integLim+maxr)/integStep)*integStep;
    intervalr = minrtmp:integStep:maxrtmp;
    
    minitmp = floor((-integLim+mini)/integStep)*integStep;
    maxitmp = ceil((integLim+maxi)/integStep)*integStep;
    intervali = minitmp:integStep:maxitmp;
    
    if dimensions == 1
        intervali = 0;
    end
    
    prA = 1/M;
    
    auxConst = 1 / sqrt(2 * pi) / sigma;
       
    if dimensions == 2
        pyfun = @(yr, yi) ...
            sum ( auxConst *  exp(-(yr - real(A)).^2 /N0) .* ...
                  auxConst .* exp(-(yi - imag(A)).^2 /N0) )  / length(A);
        
        py_given_Ak_fun = @(yr, yi, Ak) ...
                          auxConst * exp(-(yr - real(Ak))^2 /N0) * ...
                          auxConst * exp(-(yi - imag(Ak))^2 /N0);
       
        
    else
        pyfun = @(yr, yi) ...
                sum(auxConst * exp(-(yr - real(A)).^2 /N0)) / length(A);
        
        py_given_Ak_fun = @(yr, yi, Ak) ...
                          auxConst * exp(-(yr - real(Ak)).^2 /N0);          

    end
    
    
    % pre-calculate p_y
    
    fprintf('\t\tpre-calculating p_y\n');
    tic
    nelr = length(intervalr);
    neli = length(intervali);
    py = zeros(nelr,neli);
    indr = 1;
    
    if symmetry
        for indr = 1:ceil(nelr/2)
            for indi = 1:ceil(neli/2)
                temp = pyfun(intervalr(indr),intervali(indi));
                py(indr, indi) = temp;
                py(nelr-indr+1,neli-indi+1) = temp;
                py(nelr-indr+1,indi) = temp;
                py(indr,neli-indi+1) = temp;
            end
        end
    else
        for indr = 1:nelr
            for indi = 1:neli
                py(indr, indi) = pyfun(intervalr(indr),intervali(indi));
            end
        end
    end
    toc  
       
    for m = 1:length(Asym)
        
        fprintf('\t\tcalculating mutual information for constellation point %d \n', m);
        fprintf('\t\t\tpre-calculating p_y | A_m\n');
        tic
        py_given_Ak = zeros (nelr, neli);
        for indr = 1:nelr
            yr = intervalr (indr);
            for indi = 1:neli
                yi = intervali(indi);
                py_given_Ak(indr, indi) = py_given_Ak_fun (yr, yi, Asym(m));
            end
        end
        toc
        
        fprintf('\t\t\tcalculating integral\n');
        tic
        
        tic
        if dimensions == 2
            % Trapezoid Integral
            y = py_given_Ak .* log2(py_given_Ak ./ py);
            y(isnan(y)) = 0;
            
            integ = sum(sum(y(2:nelr-1,2:neli-1))) * 4 + ...
                    y(1,1) + y(1,neli) + y(nelr,1) + y(nelr, neli) + ...
                    2 * ( sum(y(1, 2:neli-1)) + sum(y(nelr, 2:neli-1)) + ...
                          sum(y(2:nelr-1, 1)) + sum(y(2:nelr-1, neli)) );
           
            integ = integStep^2 * integ / 4;
        else
            y = py_given_Ak .* log2(py_given_Ak ./ py);
            y(isnan(y)) = 0;
            integ = trapz(intervalr, y);
        end
       
        toc
        
        
     
        mutualInf(EsN0ind) = mutualInf(EsN0ind) + prA * integ * Wsym(m);
        
    end
    
    
    if isnan(mutualInf)
        mutualInf(EsN0ind:end) = log2(M);
                break
    end   
        
    fprintf('\t mutual information is %d\n', mutualInf(EsN0ind) )
    
    save (fileName, 'EsN0dB', 'mutualInf')
end


