function [Hnm,G1,G2] = applyCovConst(Hnm,Hnm_ref,debug_plots,fs)
% Apply diffuse field and covariance constraint as described in [1] sec.
% 4.11.3, which refers to [2] and [3].
%
% Important: the method works better if the HRIRs are aligned at t=0. This
% is taken care of in the function 'toSH'.
%
% INPUT:
% Hnm = HRTF's SH coefficients (nfreqs x (N+1)^2 x 2 ears)
% Hnm_ref = reference (high-order) HRTF's SH coefficients
% debug_plots = whether to show some debug plots (def=false)
% fs = sampling rate in Hz
%
% OUTPUT:
% Hnm = equalised HRTF's SH coefficients
% G1 = global EQ filter (nfreqs x 2 ears)
% G2 = left/right mix (nfreqs x 2 ears)
%
% REFERENCES:
% [1] Zotter, Franz, and Matthias Frank. Ambisonics: A practical 3D audio
% theory for recording, studio production, sound reinforcement, and
% virtual reality. Springer Nature, 2019.
% [2] Zaunschirm, Markus, Christian SchÃ¶rkhuber, and Robert HÃ¶ldrich.
% "Binaural rendering of Ambisonic signals by head-related impulse
% response time alignment and a diffuseness constraint." The Journal
% of the Acoustical Society of America 143.6 (2018): 3616-3627.
% [3] Vilkamo, Juha, Tom BÃ¤ckstrÃ¶m, and Achim Kuntz. "Optimized
% covariance domain framework for timeâ€“frequency processing of
% spatial audio." Journal of the Audio Engineering Society 61.6
% (2013): 403-411.
%
% AUTHOR: Isaac Engel - isaac.engel(at)imperial.ac.uk
% February 2021

%% Check inputs
if ~exist('debug_plots','var')
debug_plots = 0;
end

%% Some parameters
nfreqs = size(Hnm,1);
nsh = size(Hnm,2);
nears = size(Hnm,3);
N = sqrt(nsh)-1;
G1 = zeros(ones,nears); % diffuse EQ (nfreqs x nears)
G2 = zeros(nfreqs,nears); % mix-channel EQ (nfreqs x nears)
if N<1
warning('Hnm order was 0. Covariance constraint not applied...')
return
end
if nears < 2
warning('Hnm was one ear only. Covariance constraint not applied...')
end

%% Covariance constraint ([2] sec.4.11.3)
fprintf('Applying covariance constraint')
Hnm_match = zeros(size(Hnm));
count = 0; % for text output
if debug_plots % variables for debugging
df_ref = zeros(nfreqs,2);
df_1 = zeros(nfreqs,2);
df_2 = zeros(nfreqs,2);
ic_ref = zeros(nfreqs,1);
ic_1 = zeros(nfreqs,1);
ic_2 = zeros(nfreqs,1);
EQprod = zeros(nfreqs,nears);
EQsum = zeros(nfreqs,nsh,nears);
end
for fi=1:nfreqs
Hnmref_k = squeeze(Hnm_ref(fi,:,:)); % reference HRTF at fi [ndirs x 2]
Hnm_k = squeeze(Hnm(fi,:,:)); % interpolated HRTF at fi [ndirs x 2]
Rref = Hnmref_k'*Hnmref_k/(4*pi);%size(Hnmref_k,1); % reference HRTF covariance matrix
R = Hnm_k'*Hnm_k/(4*pi);%size(Hnm_k,1); % interp. HRTF covariance matrix
Rref(1,1) = real(Rref(1,1)); % force diagonals to be real
Rref(2,2) = real(Rref(2,2)); % Imag part is due to numerical error
R(1,1) = real(R(1,1));
R(2,2) = real(R(2,2));
Xref = chol(Rref); % Xref'*Xref == Href_k'*Href_k
X = chol(R); % X'*X == H_k'*H_k
[U,S,V] = svd(X'*Xref); % [1] eq.A.72 and above
if ~isreal(S) || any(S(:)<0)
warning('The SVD decomposition yielded negative or complex values')
end
M = X\(V*U'*Xref);
Hnm_match(fi,:,:) = squeeze(Hnm(fi,:,:))*M; % [1] eq.4.63 and eq.A.72
G1(fi,:) = [M(1,1),M(2,2)]; % diffuse EQ at fi (L,R)
G2(fi,:) = [M(2,1),M(1,2)]; % mix-channel EQ at fi (L,R)

% Update text output
count = count+1;
if count >= nfreqs/10
count = 0;
fprintf('.');
end
if debug_plots
% Calculate covariance matrix for the corrected low-order HRTF
Hmatch_k = squeeze(Hnm_match(fi,:,:));
Rm = Hmatch_k'*Hmatch_k/(4*pi);
Rm(1,1) = real(Rm(1,1));
Rm(2,2) = real(Rm(2,2));
% Save data to plot later
df_ref(fi,:) = [Rref(1,1),Rref(2,2)];
df_1(fi,:) = [Rm(1,1),Rm(2,2)];
df_2(fi,:) = [R(1,1),R(2,2)];
ic_ref(fi) = abs(Rref(1,2))/sqrt(Rref(1,1)*Rref(2,2));
ic_1(fi) = abs(Rm(1,2))/sqrt(Rm(1,1)*Rm(2,2));
ic_2(fi) = abs(R(1,2))/sqrt(R(1,1)*R(2,2));
EQprod(fi,:) = [M(1,1),M(2,2)];
EQsum(fi,:,:) = cat(3,M(2,1)*squeeze(Hnm(fi,:,2)),M(1,2)*squeeze(Hnm(fi,:,1)));
end
end
fprintf('\n')

% Hnm_match2(:,:,1) = Hnm(:,:,1).*G1(:,1) + Hnm(:,:,2).*G2(:,1);
% Hnm_match2(:,:,2) = Hnm(:,:,2).*G1(:,2) + Hnm(:,:,1).*G2(:,2);
%
% Hnm_match3 = Hnm.*permute(G1,[1,3,2]);
% Hnm_match3(:,:,1) = Hnm_match3(:,:,1) + Hnm(:,:,2).*G2(:,1);
% Hnm_match3(:,:,2) = Hnm_match3(:,:,2) + Hnm(:,:,1).*G2(:,2);

Hnm = Hnm_match;

%% Plots
% Test diffuse field and IC compared to ref
% This is trying to reproduce Fig. 7 from [2]
if debug_plots
f = linspace(0,fs/2,nfreqs);
figure
subplot(2,2,1)
semilogx(f,db(abs(df_ref(:,1))),'k--','Linewidth',1.5), hold on
semilogx(f,db(abs(df_1(:,1))))
semilogx(f,db(abs(df_2(:,1))))
xlim([20 20000]), grid on, xlabel('f (Hz)'), ylabel('dB')
title('diffuse field energy left')
subplot(2,2,2)
semilogx(f,db(abs(df_ref(:,2))),'k--','Linewidth',1.5), hold on
semilogx(f,db(abs(df_1(:,2))))
semilogx(f,db(abs(df_2(:,2))))
xlim([20 20000]), grid on, xlabel('f (Hz)'), ylabel('dB')
title('diffuse field energy right')
subplot(2,2,3)
semilogx(f,ic_ref,'k--','Linewidth',1.5), hold on
semilogx(f,ic_1)
semilogx(f,ic_2)
xlim([20 20000]), ylim([0 1]), grid on, xlabel('f (Hz)'), ylabel('IACC')
legend({'Ref','matched conv const','not matched'},'position',[0.5972,0.2236,0.2632,0.1090])
title('Interaural coherence')
sgtitle(sprintf('N=%d',N))
end