function [doa, params] = mclachlan2021(template,target,varargin)
%MCLACHLAN2021 - A dynamic ideal-observer model of human sound localization
%   Usage: [results,template,target] = mclachlan2014(template,target,'num_exp',20,'sig_S',4.2);
%
%   Input parameters:
%
%     template.fs      : sampling rate (Hz)
%     template.fc      : ERB frequency channels (Hz)
%     template.itd0    : itd computed for each hrir (samples)
%     template.H       : Matrix containing absolute values of HRTFS for all grid points
%     template.coords  : Matrix containing cartesian coordinates of all grid points, normed to radius 1m
%     template.T       : angular template for each coordinate
%     target.fs      : sampling rate
%     target.fc      : ERB frequency channels
%     target.itd0    : itd corresponding to source position
%     target.S       : sound source spectrum
%     target.H       : Matrix containing absolute values of HRTFS for all
%                      source directions
%
%     target.coords  : Matrix containing cartesian coordinates of all
%                      source positions to be estimated, normed to radius 1m
%
%     target.T       : angular template for each coordinate
%
%   Output parameters:
%
%     doa               : directions of arrival in spherical coordinates
%
%       .est            : estimated [num_sources, num_repetitions, 3]
%       .real           : actual    [num_sources, 3]
%
%     params            : additional model's data computerd for estimations
%
%       .est_idx        : Indices corresponding to template direction where
%                         the maximum probability density for each source
%                         position is found
%
%     	.est_loglik     : Log-likelihood of each estimated direction
%
%       .post_prob      : Maximum posterior probability
%
%                         density for each target source
%
%       .freq_channels  : number of auditory channels
%
%       .T_template     : Struct with template data elaborated by the model
%
%       .T_target       : Struct with target data elaborated by the model
%
%     	.Tidx           : Helper with indexes to parse
%                         the features from T and X
%
%   MCLACHLAN2021 accepts the following optional parameters:
%
%     'num_exp',num_exp Set the number of localization trials.
%                       Default is num_exp = 500.
%
%     'SNR',SNR         Set the signal to noise ratio corresponding to
%                       different sound source intensities.
%                       Default value is SNR = 75 [dB]
%
%     'dt',dt           Time between each acoustic measurement in seconds.
%                       Default value is dt = 0.005.
%
%     'sig_itd0',sig    Set standard deviation for the noise on the initial
%                       itd. Default value is sig_itd0 = 0.569.
%
%     'sig_itdi',sig    Set standard deviation for the noise on the itd
%                       change per time step. Default value is sig_itdi = 1.
%
%     'sig_I',sig       Set standard deviation for the internal noise.
%                       Default value is sig_I = 3.5.     
%
%     'sig_S',sig       Set standard deviation for the variation on the 
%                       source spectrum. Default value is sig_S = 3.5.
%
%     'rot_type',type   Set rotation type. Options are 'yaw', 'pitch' and
%                       'roll'. Default value is 'yaw'.
%
%     'rot_size',size   Set rotation amount in degrees. Default value is 
%                       rot_size = 0.
%
%     'stim_dur',dur    Set stimulus duration in seconds. Default value is
%                       stim_dur = 0.1.
%
%   Further, cache flags (see amt_cache) can be specified.
%
%   Url: http://amtoolbox.sourceforge.net/amt-0.10.0/doc/models/reijniers2014.php
%
%   Url: http://amtoolbox.org/amt-1.1.0/doc/models/mclachlan2021.php

% Copyright (C) 2009-2021 Piotr Majdak, Clara Hollomey, and the AMT team.
% This file is part of Auditory Modeling Toolbox (AMT) version 1.1.0
%
% 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/>.

% Copyright (C) 2009-2020 Piotr Majdak and the AMT team.
% This file is part of Auditory Modeling Toolbox (AMT) version 0.10.0
%
% 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/>.

%
%   Description: 
%   ------------
%
%   `mclachlan2021(...)` is a dynamic ideal-observer model of human sound 
%   localization, by which we mean a model that performs optimal 
%   information processing within a Bayesian context. The model considers
%   all available spatial information contained within the acoustic
%   signals encoded by each ear over a specified hear rotation. Parameters 
%   for the optimal Bayesian model are determined based on psychoacoustic 
%   discrimination experiments on interaural time difference and sound 
%   intensity.
%
%
%   Requirements: 
%   -------------
%
%   1) SOFA API v1.1  or higher from 
%      http://sourceforge.net/projects/sofacoustics for Matlab (e.g. in 
%      thirdparty/SOFA)
%
%   See also: exp_reijniers2014 plot_reijniers2014 reijniers2014_preproc
%   reijniers2014_metrics
%
%   References: reijniers2014 barumerli2020aes barumerli2020forum

%   AUTHOR: Michael Sattler and Roberto Barumerli (adapted from code provided by Jonas Reijniers)



%% Check input options
definput.import={'amt_cache'};
definput.flags.type = {'fig2'};

definput.keyvals.num_exp = 500;
definput.keyvals.SNR = 75;
definput.keyvals.dt = 0.005;
definput.keyvals.rot_type = 'yaw';
definput.keyvals.rot_size = 0.1;
definput.keyvals.stim_dur = 0.1;

% parameters of the model computed in the supplementary material
definput.keyvals.sig_itd0 = 0.569;
definput.keyvals.sig_itdi = 2*0.569;
definput.keyvals.sig_I = 3.5;
definput.keyvals.sig_S = 3.5; 
definput.keyvals.sig = 5; 

[flags,kv]  = ltfatarghelper({'num_exp','SNR','dt','rot_type','rot_size',...
                             'stim_dur','sig_itd0','sig_itdi','sig_I',...
                             'sig_S','sig'}, definput, varargin);

% prevent infinite matrix if rot_size==0                        
if kv.rot_size == 0
    kv.rot_size = 0.1;
end

%% derived parameters
rot_speed = kv.rot_size/kv.stim_dur;    % rotation speed
steps = floor(kv.stim_dur/kv.dt);       % amount of measurements steps during stimulus
alpha = linspace(0,kv.rot_size,steps);  % vector with angles at each measurement step                                               

%% define templates 
% T_template has size [Q x (2xfc+1)], where Q is number of sampled points
% and fc = number of frequency channels
T_template=[template.itd; template.ditd;...
    squeeze(template.H(:,:,1)-template.H(:,:,2)); ...
    0.5.*squeeze(template.H(:,:,1)+template.H(:,:,2))]'; 

T_target=[target.itd; target.ditd; ...
    squeeze(target.H(:,:,1)-target.H(:,:,2)); ...
    0.5.*squeeze(target.H(:,:,1)+target.H(:,:,2))]'; 


%% define error covariance matrix
sig_itd0 = kv.sig_itd0; %0.569;
sig_itdi = kv.sig_itdi; %still unknown, currently set to 1
sig_I = kv.sig_I; % 3.5; Intensity discrimination for broadband signal
sig_S = kv.sig_S; %3.5; Source's template error
sig = kv.sig; % Expected variance on the source strength - interchannel noise communication
sig_i = [sig_itd0, repmat(sig_itdi,1,steps-1)];   % var on ITD at each time step

% create M_beta covariance matrix
W = diag(1./(sig_i.^2));    % weight matrix
X = ones(steps,2);          % each column a slope of a parameter beta
X(:,2) = alpha;

M_beta = inv(X.'*W*X); % covariance matrix of ITD0 and ITDd

Sig = blkdiag(M_beta, 2*sig_I^2*eye(length(template.fc)), ...
    (sig_I^2/2 + sig_S^2)*eye(length(template.fc)) + sig^2); % full covariance matrix

%% simulate num_exp experimental trials 
num_exp = kv.num_exp;
invSig = inv(Sig);
num_targ = size(target.coords,1);
num_temp = size(template.coords,1);
log_lik = zeros(num_targ, num_exp);
doa_idx = zeros(num_targ, num_exp);
post_prob = zeros(num_targ, num_exp, num_temp);
doa_estimations = zeros(num_targ, num_exp, 3);
entropy = zeros(num_targ,1); % entropy in bits
Entropy = zeros(num_targ,1);
if nargout > 1
    X_all = zeros(num_targ, num_exp, size(T_target, 2));
end

for e = 1:num_exp
    amt_disp(sprintf('experiment %i', e))
    X = mvnrnd(T_target,Sig);
    if nargout > 1
        X_all(:,e,:) = X;
    end
    
    for s = 1:num_targ
        for d = 1:num_temp
            % Formula R
            u_diff = (X(s,:)-T_template(d,:));
            post_prob(s,e,d) = abs(exp(-0.5* u_diff*invSig*transpose(u_diff)));
        end
        % normalize
        post_prob(s,e,:) = post_prob(s,e,:)/sum(post_prob(s,e,:) + eps); 
        % maximum a posteriori
        [log_lik(s,e), doa_idx(s,e)] = max(post_prob(s,e,:));
        entropy(s)= - squeeze(post_prob(s,e,:))'*log2(squeeze(post_prob(s,e,:)) + eps);
        doa_estimations(s,e,:) = template.coords(doa_idx(s,e), :);
    end
    Entropy = Entropy + entropy; %cumulative entropy over experiments
end

Entropy = Entropy/num_exp;
Information = log2(num_temp) - Entropy; 

%% results
doa.est = doa_estimations;
doa.real = target.coords;

% user required more than the estimations
if nargout > 1
    params.template_coords = template.coords;
    params.post_prob = post_prob;
    params.entropy = Entropy;
    params.information = Information;
    params.est_idx = doa_idx;
    params.est_loglik = log_lik;
    params.X = X_all;
    params.T_template = T_template;
    params.T_target = T_target;
    params.freq_channels = template.fc;
    params.Tidx.itd = 1;
    assert(length(target.fc)==length(template.fc))
    params.Tidx.Hp = params.Tidx.itd + (1:length(target.fc));
    params.Tidx.Hm = params.Tidx.Hp(end) + (1:length(target.fc));
else
    clear X_all post_prob doa_idx log_lik
end

end

function Y_N = SH(N, dirs)
% calculate spherical harmonics up to order N for directions dirs [azi ele;...] (in radiant)
% 
    N_dirs = size(dirs, 1);
    N_SH = (N+1)^2;
	dirs(:,2) = pi/2 - dirs(:,2); % convert to inclinations

    Y_N = zeros(N_SH, N_dirs);

	  % n = 0
	Lnm = legendre(0, cos(dirs(:,2)'));
	Nnm = sqrt(1./(4*pi)) * ones(1,N_dirs);
	CosSin = zeros(1,N_dirs);
	CosSin(1,:) = ones(1,size(dirs,1));
	Y_N(1, :) = Nnm .* Lnm .* CosSin;
	
	  % n > 0
	idx = 1;
    for n=1:N
        
        m = (0:n)';            

		Lnm = legendre(n, cos(dirs(:,2)'));
		condon = (-1).^[m(end:-1:2);m] * ones(1,N_dirs);
		Lnm = condon .* [Lnm(end:-1:2, :); Lnm];
		
		mag = sqrt( (2*n+1)*factorial(n-m) ./ (4*pi*factorial(n+m)) );
		Nnm = mag * ones(1,N_dirs);
		Nnm = [Nnm(end:-1:2, :); Nnm];
		
		CosSin = zeros(2*n+1,N_dirs);
			% m=0
		CosSin(n+1,:) = ones(1,size(dirs,1));
			% m>0
		CosSin(m(2:end)+n+1,:) = sqrt(2)*cos(m(2:end)*dirs(:,1)');
			% m<0
		CosSin(-m(end:-1:2)+n+1,:) = sqrt(2)*sin(m(end:-1:2)*dirs(:,1)');

		Ynm = Nnm .* Lnm .* CosSin;
        Y_N(idx+1:idx+(2*n+1), :) = Ynm;
        idx = idx + 2*n+1;
    end
    
    Y_N = Y_N.';
    
end