rng('default');
numFrames = 6;                                      % Number of frames
numOFDMSymPerFrame = 140;                           % Number of OFDM symbols per frame
maxLenChEstiPerSym = 14400;                         % Maximum length of channel estimates per symbol
numSubCarPerSym = 72;                               % Number of subcarriers per OFDM symbol
hEstLen = numSubCarPerSym * numOFDMSymPerFrame;     % Channel estimate length
totNumOFDMSymbols = numFrames * numOFDMSymPerFrame; % Total number of OFDM symbols

dataInGrid = zeros(numSubCarPerSym,totNumOFDMSymbols);
for subCarCount = 0:numSubCarPerSym-1
    for numOFDMSymCount = 0:totNumOFDMSymbols-1
        realXgain = 1 + .2*sin(2*pi*subCarCount/numSubCarPerSym);
        realYgain = 1 + .5*sin(2*pi*numOFDMSymCount/numOFDMSymPerFrame);
        imagXgain = 1 + .3*sin(2*pi*subCarCount/numSubCarPerSym);
        imagYgain = 1 + .4*sin(2*pi*numOFDMSymCount/numOFDMSymPerFrame);
        dataInGrid(subCarCount+1,numOFDMSymCount+1) = realXgain*realYgain + 1i*(imagXgain*imagYgain);
    end
end
validIn = true(1,length(dataInGrid(:)));

% Normalize data subcarriers to make signal power unity
dataInGrid = dataInGrid./sqrt(mean(abs(dataInGrid).^2,'all'));

figure(1);
surf(real(dataInGrid))
xlabel('OFDM Symbols')
ylabel('Subcarriers')
zlabel('Magnitude')
title('Input Data Grid (Real Part)')

figure(2);
surf(imag(dataInGrid))
xlabel('OFDM Symbols')
ylabel('Subcarriers')
zlabel('Magnitude')
title('Input Data Grid (Imaginary Part)')

numOFDMSymToBeAvg = 1;                                            % Number of OFDM symbols to be averaged
interpolFac = 1;                                                  % Interpolation factor
dataInForChannelEsti = dataInGrid(:,1:numOFDMSymPerFrame);
validInForChanEsti = validIn(1:numSubCarPerSym*numOFDMSymPerFrame);

numScPerSymIn = numSubCarPerSym*true(1,length(dataInForChannelEsti(:)));

refDataIn = randsrc(size(dataInForChannelEsti(:),1),size(dataInForChannelEsti(:),2),[1 1]);
refValidIn = boolean(zeros(1,numOFDMSymPerFrame*numSubCarPerSym));
startRefValidIndex = randi(interpolFac,1,1);
for numOFDMSymCount = 1:numOFDMSymPerFrame
    refValidIn(startRefValidIndex+(numOFDMSymCount-1)*numSubCarPerSym:interpolFac:numSubCarPerSym*numOFDMSymCount) = true;
end

dataOut1 = channelEstReferenceForEqualizer( ...
    numOFDMSymToBeAvg,interpolFac,numSubCarPerSym,numOFDMSymPerFrame, ...
    dataInForChannelEsti(:),validInForChanEsti,refDataIn,refValidIn,numScPerSymIn);
matlabOut = dataOut1(:);
hEstIn = zeros(numel(matlabOut)*numSubCarPerSym*numOFDMSymToBeAvg,1);
for ii= 1:numel(matlabOut)
    loadArray = [matlabOut(ii).dataOut; repmat(matlabOut(ii).dataOut,[numOFDMSymToBeAvg-1 1]); zeros((length(hEstIn)-numSubCarPerSym*numOFDMSymToBeAvg),1)];
    shiftArray = circshift(loadArray,(ii-1)*numSubCarPerSym*numOFDMSymToBeAvg);
    hEstIn =  hEstIn + shiftArray;
end

% Repeat hEstIn for dataIn generation
hEstInForDataIn = repmat(hEstIn,numFrames,1);

% Normalize channel estimates to make signal power unity
hEstIn = hEstIn./sqrt(mean(abs(hEstIn).^2,'all'));
hEstIn = [hEstIn; hEstIn(end)*ones((hEstLen*((totNumOFDMSymbols/numOFDMSymPerFrame)-1)),1)];

% Generate noise samples
n = (1/sqrt(2))*(randn(length(dataInGrid(:)),1)+1i*randn(length(dataInGrid(:)),1)); % white gaussian noise, variance=1, mean=0;

SNR = 40;
% Calculate noise variance
nVar = (10^(-SNR/10));
noiseVarIn = (10^(-SNR/10))*ones(1,length(dataInGrid(:)));

modelname = 'genhdlOFDMEqualizerModel';
open_system(modelname);

EqMdUsed = get_param('genhdlOFDMEqualizerModel/HDL Algorithm/OFDM Equalizer','EqualizationMethod');
if strcmp(EqMdUsed,'ZF')
    % ZF equalization
    dataIn = hEstInForDataIn.*dataInGrid(:);
else
    % MMSE equalization
    dataIn = hEstInForDataIn.*dataInGrid(:) + (n.*(sqrt(nVar)))./(sqrt(var(n)));
end
% Generate signal with channel estimate length per symbol
hEstLenIn = hEstLen*true(1,length(dataInGrid(:)));
loadhEst = logical([1 zeros(1,length(dataInGrid(:))-1)]);
resetSig = false(1,length(dataInGrid(:)));

out = sim(modelname);

simOut = out.dataOut.Data(out.validOut.Data);
N = length(simOut)-mod(length(simOut),numSubCarPerSym);
temp = simOut(1:N);
EqualizerSimOut = reshape(temp,numSubCarPerSym,length(temp)/numSubCarPerSym);

figure(3);
surf(real(EqualizerSimOut))
xlabel('OFDM Symbols')
ylabel('Subcarriers')
zlabel('Magnitude')
title('OFDM Equalizer Output (Real Part)')

figure(4);
surf(imag(EqualizerSimOut))
xlabel('OFDM Symbols')
ylabel('Subcarriers')
zlabel('Magnitude')
title('OFDM Equalizer Output (Imaginary Part)')

if strcmp(EqMdUsed, 'ZF')
    % ZF equalization
    matOut = dataIn./hEstInForDataIn;
else
    % MMSE equalization
    matOut = (1./(conj(hEstInForDataIn).*hEstInForDataIn+nVar)).*(conj(hEstInForDataIn)).*dataIn;
end

figure('units','normalized','outerposition',[0 0 1 1])
subplot(2,1,1)
plot(real(matOut(:)));
hold on;
plot(real(simOut(:)));
grid on
legend('MATLAB reference output','Simulink block output')
xlabel('Sample Index')
ylabel('Magnitude')
title('Comparison of Simulink Block and MATLAB Function (Real Part)')

subplot(2,1,2)
plot(imag(matOut(:)));
hold on;
plot(imag(simOut(:)));
grid on
legend('MATLAB reference output','Simulink block output')
xlabel('Sample Index')
ylabel('Magnitude')
title('Comparison of Simulink Block and MATLAB Function (Imaginary Part)')

sqnrRealdB = 10*log10(double(var(real(simOut(:)))/abs(var(real(simOut(:)))-var(real(matOut(:))))));
sqnrImagdB = 10*log10(double(var(imag(simOut(:)))/abs(var(imag(simOut(:)))-var(imag(matOut(:))))));

fprintf('\n OFDM Equalizer \n SQNR of real part: %.2f dB',sqnrRealdB);
fprintf('\n SQNR of imaginary part: %.2f dB\n',sqnrImagdB);