% ============================== % 
% This file studies the DoF/Capacity of the proposed colored fading channel
% model. 
% 
%  Regime: H-MIMO. 
% ============================== % 
clear all; close all; clc; 

addpath('algs/'); 
addpath('utils/');
fc      = 3.5e9; 
c       = physconst('lightspeed'); 
lambda  = c/fc; 
k0      = 2*pi/lambda; 

Gamma   = 0.05; 

RegionsTx = cell(1, 1);
RegionsTx{1} = [-0.4, 0, 0.45, Gamma];   % format: (nkx, nky, nkradius, gammaCorrelationFactor). 
% RegionsTx{2} = [0.3, 0, 0.15, Gamma];
% RegionsTx{3} = [0.7, 0, 0.15, Gamma]; 

RegionsRx = cell(1,1); 
RegionsRx{1} = [-0.15, 0, 0.45, Gamma]; 
% RegionsRx{2} = [-0.6, 0, 0.15, Gamma];
% RegionsRx{3} = [0.52, 0, 0.15, Gamma];

reg.RegionsTx = RegionsTx; 
reg.RegionsRx = RegionsRx; 
reg.Gamma = Gamma; 

P       = 10;
sigma2  = 1; 
L       = 1;      

%% Get channel and precode. 
fprintf('Simulation with SNR = %.2f dB\n', pow2db(P/sigma2)); 

SpacingArray    = flip(logspace(-1, log10(2), 20)); 
N_scan          = length(SpacingArray); 
N_test          = 300;                      % Usually this number is bigger than the XL-MIMO case. 

maxRate         = zeros(N_scan, 1); 
maxRateStd      = zeros(N_scan, 1); 
lbRate          = zeros(N_scan, 1); 
lbRateStd       = zeros(N_scan, 1); 
avgDoF          = zeros(N_scan, 1);
avgDoFStd       = zeros(N_scan, 1); 

% Bound evaluation
AntDensity = (1./SpacingArray).';               % Number of antennas within one λ  
theoreticalRelativeMI   = zeros(N_scan, 2); 

OmegaTx = 0; 
for idx = 1:length(RegionsTx)
    OmegaTx = OmegaTx + 2*RegionsTx{idx}(3); 
end

OmegaRx = 0; 
for idx = 1:length(RegionsRx)
    OmegaRx = OmegaRx + 2*RegionsRx{idx}(3); 
end

[ECpswf, ECa, aDoF, NyqNumber]  = getPSWFBound(OmegaTx, L/lambda, Gamma, P/sigma2); 
theoreticalDoF                  = repmat(aDoF, [N_scan, 1]); 
theoreticalRelativeMI(:,1)      = repmat(ECa, [N_scan, 1]);
theoreticalRelativeMI(:,2)      = repmat(ECpswf, [N_scan, 1]); 

fprintf('PSWF-Computed DoF (Nyquist number) = %.2f\n', NyqNumber); 
fprintf('PSWF capacity bound (TIT`14 formula) = %.2f\n', ECpswf); 


% Ergodic capacity simulation
fprintf('+============+===============+========================================+===================+\n');
fprintf('|            |               |               Rate (bps/Hz)            |                   |\n');
fprintf('|  Progress  | Ant. Spacing  +----------------------------------------+  Elapsed Time (s) |\n');
fprintf('|            |   (/λ, Nx)    | max Rate | lower bound |  average DoF  |                   |\n');
fprintf('+------------+---------------+----------+-------------+---------------+-------------------+\n');


for i_scan = 1:N_scan
    rng(0);

    MIMO_params = setup("SymmetricULAwithFixedK", fc, L, reg, SpacingArray(i_scan)*lambda);           % MIMO setup
    MIMO_params.memorizeCorrelation     = true; 
    MIMO_params.FreqShift               = false; 
    MIMO_params.FreqShiftFactor         = 20; 
    MIMO_params.sigma2                  = sigma2; 
    MIMO_params.P                       = P;
    MIMO_params.debug                   = false; 
    
    inner_maxRate                       = zeros(N_test, 1); 
    inner_lbRate                        = zeros(N_test, 1); 
    inner_avgDoF                        = zeros(N_test, 1); 
    
    tic; 
    for i_test = 1:N_test
        H = getChannel(MIMO_params, "ColoredAngularCorrelation"); 
        
        H = H*(min(1/2, SpacingArray(i_scan))); 
        % [NMSE, ~] = estimateChannel(H, MIMO_params, "PSWF"); 

        % Precode for MIMO channel H. 
        [W, U, p] = waterFilling(H, P, sigma2); 
        
        inner_maxRate(i_test) = log(real(det(eye(MIMO_params.RxArray.N)+(1/sigma2)*(H*W*diag(p)*W'*H'))))/log(1+MIMO_params.P/sigma2); 
        inner_lbRate(i_test) = log(real(det(eye(MIMO_params.RxArray.N)+(MIMO_params.P/(MIMO_params.TxArray.N*sigma2))*(H*H'))))/log(1+MIMO_params.P/sigma2); 
        inner_avgDoF(i_test) = sum(p>0); 
    end
    timeElapsed = toc(); 
    
    maxRate(i_scan) = mean(inner_maxRate); 
    lbRate(i_scan) = mean(inner_lbRate);
    avgDoF(i_scan) = mean(inner_avgDoF); 
    
    maxRateStd(i_scan) = std(inner_maxRate); 
    lbRateStd(i_scan) = std(inner_lbRate); 
    avgDoFStd(i_scan) = std(inner_avgDoF); 

    fprintf('|  %3d/%3d   | (%5.2f, %4d) |  % 6.2f  |   % 6.2f    |    % 6.2f     |      % 8.3f     |\n', ...
        i_scan, N_scan, ...
        SpacingArray(i_scan),  MIMO_params.TxArray.N, ...
        maxRate(i_scan), ...
        lbRate(i_scan), ...
        avgDoF(i_scan), ...
        timeElapsed); 
end
fprintf('+------------+---------------+----------+-------------+---------------+-------------------+\n');
fprintf('Simulation successfully completed.\n'); 



%% Visualization 
close all; 

saveFigs = false; 
set(0,'DefaultLineMarkerSize',  6);
set(0,'DefaultTextFontSize',    14);
set(0,'DefaultAxesFontSize',    12);
set(0,'DefaultLineLineWidth',   1.4);
set(0,'defaultfigurecolor',     'w');

% figure(1);

% plot(AntDensity, maxRate); grid on; hold on;
% plot(AntDensity, lbRate); 
% set(gca, 'xscale', 'log');
% xlabel('Antenna Density ($N_{\rm ant}$ within $\lambda_c$)', 'Interpreter','latex'); 
% ylabel('Information (nat/s/Hz)'); 
% legend({'Optimal SVD+WaterFilling', 'MI with equi-power (lower bound)'}); 

% Plot Rate

fig1 = figure(1);

plot(AntDensity, theoreticalRelativeMI(:,2), 'LineStyle', '--', 'Color', 'k'); grid on; hold on;
plot(AntDensity, maxRate, 'Marker','square', 'color', [0 0 1]); 
plot(AntDensity, lbRate, 'Marker','v', 'color', [1, 0, 0]); 

set(gca, 'xscale', 'log');
xlabel('Antenna Density ($N_{\rm ant}$ per $\lambda_c$)', 'Interpreter','latex'); 
ylabel('Normalized Capacity'); 
legend({'Ergodic Capacity (PSWF bound)', ...
        'Capacity SVD+WaterFilling', ...
        'Capacity Equi-power',}, 'FontSize', 10, 'Location','best'); 

if saveFigs
    exportgraphics(fig1, 'results/HMIMOregime/HMIMO_capacity.pdf', 'ContentType','vector'); 
    fprintf('HMIMO: capacity figure saved.\n'); 
end


% Plot Errorbar Std
fig2 = figure(2); 
plot(AntDensity, theoreticalRelativeMI(:,2), 'LineStyle', '--', 'Color', 'k'); hold on; grid on; box on;
errorbar(AntDensity, maxRate, maxRateStd, 'Marker','square', 'color', [0 0 1], 'LineWidth',1.4);    
errorbar(AntDensity, lbRate, lbRateStd, 'Marker','v', 'color', [1, 0, 0], 'LineWidth',1.4); 
line([2, 2], [0, 20], 'linestyle', ':', 'color', 'k', 'LineWidth', 1); 


set(gca, 'xscale', 'log');
xlabel('Antenna Density ($N_{\rm ant}$ per $\lambda_c$)', 'Interpreter','latex'); 
ylabel('Normalized Capacity'); 
legend({'Ergodic Capacity (PSWF bound)', ...
    'Capacity SVD+WaterFilling', ...
    'Capacity Equi-power'}, 'Fontsize', 10, 'Location','best'); 

if saveFigs
    exportgraphics(fig2, 'results/HMIMOregime/HMIMO_capacityWithErrorBars.pdf', 'ContentType','vector'); 
    fprintf('HMIMO: capacity figure with error bar saved.\n'); 
end

% Plot DoF
figure(3); 
errorbar(AntDensity, avgDoF, avgDoFStd, 'Marker','<', 'color', [0, 0.7, 0.7], 'LineWidth', 1.4);  hold on; grid on; 
plot(AntDensity, theoreticalRelativeMI(:,1), 'Linestyle', '--', 'Color', 'k'); 
set(gca, 'xscale', 'log');
xlabel('Antenna Density ($N_{\rm ant}$ per $\lambda_c$)', 'Interpreter','latex'); 
ylabel('Effective DoF'); 
legend({'Simulated', 'PSWF-DoF upper bound'}, 'FontSize', 10, 'Location','best');

if saveFigs
    exportgraphics(gcf, 'results/HMIMOregime/HMIMO_DoF.pdf', 'ContentType','vector'); 
    save('results/HMIMOregime/data.mat'); 
    fprintf('Files saved at results/HMIMOregime/. \n'); 
end


fprintf('End of this script.\n'); 

