% ======================================================================================= % 
% This file studies the colored fading channel model. 
% We need to find a good estimate for this channel. 
% 
%   The PSWF-CE is a two-step algorithm. 
% 
%   First step: Bandwidth estimation. We need to add some noise into the bandwidth 
%               to see its impact on the NMSE performance. 
%   Second step: channel estimation 
% 
% ======================================================================================== % 
clear all; close all; clc;

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


%% Simulate channel estimators. 

rng(0); 
L                       = 0.5;
% L           = 3; 
N_test                  = 200; 
P                       = 100;
sigma2                  = 1; 
% channelGenerationMethod = "CDL"; 
channelGenerationMethod = "ColoredAngularCorrelation"; 

% Setup bandwidths. 
reg.Gamma       = 0.05; 
RegionsTx       = cell(1,1);
RegionsTx{1}    = [0, 0, 0.15, reg.Gamma];   % Centered wavenumber domain channel. 
RegionsRx       = cell(1,1); 
RegionsRx{1}    = [0, 0, 0.15, reg.Gamma];   
reg.RegionsTx   = RegionsTx; 
reg.RegionsRx   = RegionsRx; 

w_true = 0.15; 
w_bias_arr = (-0.1:0.02:0.2).';  



% MIMO = setup("SIMOULAwithFixedK", fc, L, reg);     
MIMO = setup("SymmetricULAwithFixedK", fc, L, reg);  

MIMO.P                          = P; 
MIMO.sigma2                     = sigma2; 
MIMO.normalizeChannel           = true;         % Normalize the average power of each row of the channel matrix to 1. 
MIMO.forceGetCovariance         = true;
MIMO.memorizeCorrelation        = true; 
% MIMO.Estimators.BWest.rcomb_mode = "external"; 
MIMO.Estimators.BWest.rcomb_mode = "auto"; 
% MIMO.Estimators.BWest.Nrcomb = 3; 
MIMO.Estimators.BWest.w_mode    = "fixed"; 

NRx             = MIMO.RxArray.N; 

% Dual-parameter scanning 
Ncombiner_arr   = (30:10:80).'; 

Ns1 = length(w_bias_arr);  
Ns2 = length(Ncombiner_arr); 
Nsweep = Ns1*Ns2; 

% Ncombiner_arr = (34:2:46).'; 
NMSE            = zeros(Nsweep, 7); 
StdNMSE         = zeros(Nsweep, 7); 
StdNMSEdB       = zeros(Nsweep, 7); 
BWs             = zeros(Nsweep, 1); 
StdBWs          = zeros(Nsweep, 1); 


fprintf('Simulation with SNR = %.2f dB with Nantenna = %d\n', pow2db(P/sigma2), NRx); 
fprintf('Channel Generation Method = %s\n', channelGenerationMethod); 
fprintf('+============+========================================================================================================+\n');
fprintf('|            |                                                 NMSE(dB)                                               |\n');
fprintf('| widx/Ncomb +--------------------------------------------------------------------------------------------------------+\n');
fprintf('|            | PSWF w/  Stat | PSWF w/o Stat |  RandComb  |  AMP  |   SBAR   | BT-PSWF | BWest-PSWF | Time Elapsed (s)|\n');
fprintf('+------------+---------------+---------------+------------+-------+----------+---------+------------+-----------------+\n');

for i_sweep = 1:Nsweep
    % rng(0); 
    i1 = mod(i_sweep-1, Ns1)+1; 
    i2 = floor((i_sweep-1)/Ns1)+1; 
    
    Ncombiner       = Ncombiner_arr(i2); 
    w_bias          = w_bias_arr(i1); 

    MIMO.Ncombiner  = Ncombiner; 
    
    inner_NMSE1     = zeros(N_test, 1); 
    inner_NMSE2     = zeros(N_test, 1); 
    inner_NMSE3     = zeros(N_test, 1); 
    inner_NMSE4     = zeros(N_test, 1); 
    inner_NMSE5     = zeros(N_test, 1); 
    inner_NMSE6     = zeros(N_test, 1); 
    inner_NMSE7     = zeros(N_test, 1); 
    inner_BWs       = zeros(N_test, 1); 
    
    tic; 
    parfor i_test = 1:N_test
        rng(i_test); 

        w = w_true + w_bias;  

        [H, Corr] = getChannel(MIMO, channelGenerationMethod); 
            
        NTx                     = MIMO.TxArray.N; 
        Nh                      = NRx*NTx; 
        H                       = H * sqrt(Nh);
        nvec                    = sqrt(sigma2)*(randn([Ncombiner, 1]) + 1j*randn([Ncombiner, 1]))/sqrt(2); 

        % [Method 7] Bandwidth Estimating PSWF
        inner_MIMO = MIMO; 
        inner_MIMO.Estimators.BWest.w = w; 
        [inner_NMSE7(i_test), inner_BWs(i_test)]  = AdaptivePSWFestimator(H, P, nvec, inner_MIMO, "MIMO-BWest-PSWF"); 
    end
    timeElapsed = toc(); 
    
    % Place 7 variables to enable parfor. 
    inner_NMSE      = [inner_NMSE1, inner_NMSE2, inner_NMSE3, inner_NMSE4, inner_NMSE5, inner_NMSE6, inner_NMSE7]; 

    BWs(i_sweep)           = mean(inner_BWs); 
    NMSE(i_sweep, :)       = mean(inner_NMSE, 1);
    StdNMSE(i_sweep, :)    = std(inner_NMSE, 0, 1); 
    StdNMSEdB(i_sweep, :)  = std(pow2db(inner_NMSE), 0, 1); 
    StdBWs(i_sweep)        = std(inner_BWs); 

    fprintf('|  %3d/%3d   |     % 6.2f    |    % 6.2f     |   % 6.2f   |% 6.2f |  % 6.2f  | % 6.2f  |   %6.2f   |      %6.3f     |\n', ...
        i1, Ncombiner,...
        pow2db(NMSE(i_sweep, 1)), ...
        pow2db(NMSE(i_sweep, 2)), ...
        pow2db(NMSE(i_sweep, 3)), ...
        pow2db(NMSE(i_sweep, 4)), ...
        pow2db(NMSE(i_sweep, 5)), ...
        pow2db(NMSE(i_sweep, 6)), ...
        pow2db(NMSE(i_sweep, 7)), ...
        timeElapsed); 
end
fprintf('+------------+---------------+---------------+------------+-------+----------+---------+------------+-----------------+\n');
fprintf('Simulation complete...\n');


%% Visualization
close all; 
set(0,'DefaultLineMarkerSize',  6);
set(0,'DefaultTextFontSize',    14);
set(0,'DefaultAxesFontSize',    12);
set(0,'DefaultLineLineWidth',   1.4);
set(0,'defaultfigurecolor',     'w');

fig1 = figure('color', [1,1,1]);  
% fig1.Position = [300 100 700 550]; 
% figure('Renderer', 'painters', 'Position', [300 100 700 550], 'Color', [1 1 1]);  % format: (x, y, width, height) in pixels. 
% figure('color', [1 1 1]); 

if MIMO.isSIMO
    methodSel = [3, 4, 5, 6, 7, 2, 1]; 
else
    methodSel = [3, 4, 7, 2, 1]; 
end

plotMethod  = "errorBar"; 
saveFigs    = true; 
colors      = othercolor('BrBG6', Ns2); 

if plotMethod == "plain"
    % plot(Nrcomb_arr, NMSE(:,3), 'Color','b', 'Marker','v'); hold on; grid on; box on;  
    % plot(Nrcomb_arr, NMSE(:,4), 'Color', [0 0.6 0.6] , 'Marker','<'); 
    % plot(Nrcomb_arr, NMSE(:,5), 'Color', [0.8, 0.2, 0.8], 'Marker','o'); 
    % plot(Nrcomb_arr, NMSE(:,6), 'Color', [1, 0.3, 0], 'Marker', 'square'); 
    plot(1-Nrcomb_ratio/Ncombiner, NMSE(:,7), 'Color', [1, 0, 0], 'Marker','diamond'); 
    % plot(Nrcomb_arr, NMSE(:,2), 'Color', [0.7, 0, 0], 'LineStyle','-', 'Marker', 'pentagram'); 
    % plot(Nrcomb_arr, NMSE(:,1), 'Color', 'k', 'LineStyle','--'); 
    set(gca, 'yscale', 'log');
    xlabel('PSWF combiner ratio'); 
    ylabel('NMSE'); 

elseif plotMethod == "errorBar"
    
    for i2 = 1:Ns2
        sel = (i2-1)*Ns1+1:i2*Ns1; 
        errorbar(w_bias_arr, pow2db(NMSE(sel,7)), StdNMSEdB(sel,7)/sqrt(N_test), 'Color', colors(i2, :), 'Marker','diamond', 'LineWidth', 1); hold on; grid on; box on; 

    end

    xlabel('Bandwidth bias $\Delta W$', 'Interpreter', 'latex');   
    ylabel('NMSE (dB)', 'Interpreter', 'latex');  
    xlim([min(w_bias_arr), max(w_bias_arr)]); 
    % ylim([-20, -2]); 

else
    error('plotMethod %s not defined.\n', plotMethod); 

end

% line([36, 36], [1e-3, 1e0], 'linestyle', ':', 'Color', 'k'); 
% methodNames = { 'PSWF w/ Stat CSI MMSE (Oracle)', ...
%                 'PSWF w/o Stat CSI MMSE', ...
%                 'Random Comb. MMSE', ...
%                 'Random Comb. AMP', ...
%                 'SBAR', ...
%                 'Adaptive PSWF', ...
%                 'BWest PSWF', ...
%                 }; 

methodNames = {}; 
for i2 = 1:Ns2
    methodNames{i2} = sprintf("Ncomb=%d", Ncombiner_arr(i2)); 
end
legend(methodNames, "FontSize", 10, 'Location', 'best'); 

if saveFigs
    exportgraphics(fig1, 'results/Response/MIMO_CE_BandwidthBias.pdf', 'ContentType', 'vector');
    fprintf('File saved.\n'); 
end

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

