%% Does the actual reliability calculations. Ruff removed HRTF's have less 
% directions, so downsample normal files. Makes things go faster too.
% Currently at 10 iterations
clear 
cd('F:\Shadron_Pena_2022\')
%cd('C:\Users\penalab2\Documents\MATLAB\Brian_Reliability')
addpath('F:\Shadron_Pena_2022')
% % Load HRIRs

owl_folder{1}='Merry';
owl_folder{2}='Owl19';
owl_folder{3}='Owl39';
owl_folder{4}='Ugly';
owl_folder{5}='Vespucci';

haircut_id{1}='normal';
haircut_id{2}='ohne'; % maximum cut
date_folder = '10-24-22 Data'; %change to current date
%% This anlaysis is similar to Figure 1 analysis, but no maskers are used.

for o=1:5
    for h=1:2
        clearvars -except h haircut_id o owl_folder date_folder
        load(['HRTF_Wagner_Lab\results\' owl_folder{o} '\' owl_folder{o} '__' haircut_id{h} '__hrir_result.mat'])
        
        newazimuths = -160:20:160;
        newelevations = -60:20:80;
        
        dir=NaN(2,length(azimuths)*length(elevations));
        TF1=NaN(length(azimuths)*length(elevations),size(hrir_l,3));
        TF2=NaN(length(azimuths)*length(elevations),size(hrir_l,3));
        count=0;
        for a=1:length(azimuths)
            for e=1:length(elevations)
                count=count+1;
                dir(:,count)=[elevations(e);azimuths(a)];
                TF1(count,:)=hrir_l(a,e,:);
                TF2(count,:)=hrir_r(a,e,:);
            end
        end        
                
        newdir=NaN(2,length(newazimuths)*length(newelevations));
        count=0;
        for a=1:length(newazimuths)
            for e=1:length(newelevations)
                count=count+1;
                newdir(:,count)=[newelevations(e);newazimuths(a)];
            end
        end                
        
        dir = dir';
        newdir = newdir';
        downsample = ismember(dir, newdir, 'rows');
        TF1 = TF1(downsample, :);
        TF2 = TF2(downsample, :);
        dir = newdir';
        azimuth = newazimuths;
        elevation = newelevations;
        
        clear newdir newazimuth newelevation
        
        % % Analysis parameters
        
        %center frequencies of cochlear filters (Hz)
        cF = 400:200:10000;
        %cF = linspace(1000,9000, 9);
        n_cF=length(cF);
        
        %Sampling frequency for HRTFs (Hz)
        Fs = samplingrate;
        
        %Factor by which you will upsample
        Factor = 5;
        %Number of trials for each condition
        Nt=10;
        % % Get cochlear filters
        %Upsampled frequency
        FS=Fs*Factor;
        %Get filters
        fcoefs=getFilterCoefs(cF,FS);


        % % Number of directions and initialize
        n_dir = size(TF1,1);

        IPDm_all = zeros(n_dir, n_cF, Nt);
        IPDs_all = zeros(n_dir, n_cF, Nt);
        %         ITDm=zeros(n_dir,n_cF);
        %         ITDs=zeros(n_dir,n_cF);
        ILDm_all = zeros(n_dir, n_cF, Nt);
        ILDs_all = zeros(n_dir, n_cF, Nt);

        gainr_all = zeros(n_dir, n_cF, Nt);
        gainl_all = zeros(n_dir, n_cF, Nt);


        % %
        tic

        TF1_res = resample(TF1',Factor,1)';
        TF2_res = resample(TF2',Factor,1)';

        parfor idir = 1:n_dir % TODO parfor

            %Resample the HRIRs
            hL=TF1_res(idir,:);
            hR=TF2_res(idir,:);
            %Run over trials

            for tt=1:Nt

                %Generate a target sound
                s=genSignal(100,1/30,1,2,2*pi*12);
                s1=resample(s,Factor,1);

                %Convolve target with HRIR
                sL=conv(hL,s1);
                sR=conv(hR,s1);

                ITDp=zeros(n_cF);

                %cochlear (gammmatone) filterbank
                vL=ERBFilterBank(sL,fcoefs);
                vR=ERBFilterBank(sR,fcoefs);

                %Run over frequency and compute cross correlation
                for icF = 1:n_cF
                    [itd,l]=xcorr(vL(icF,:),vR(icF,:));
                    [~,j]=max(itd);
                    ITDp(icF)=l(j)*1/Fs*1e6/Factor;
                end

                %Compute ILD
                R=rms(vR,2);
                L=rms(vL,2);
                Z=20*log10(R./L);
                GR=20*log10(R);
                GL=20*log10(L);
                
                %Run over frequency and compute mean and SD
                for n=1:n_cF
                    IPD_all(idir,n,tt) =ITDp(n)*cF(n)/1e6*2*pi; %#ok<PFBNS>
                end

                %compute mean and SD for ILD.
                ILD_all(idir,:,tt) = Z;

                %compute gain
                gainr_all(idir,:,tt) =GR;
                gainl_all(idir,:,tt) =GL;


            end
        end
        IPDs = std(IPD_all,0,3);
        IPDm = mean(IPD_all, 3);
        ITDs = (1e6 * IPDs) ./ (2 * pi * cF);
        ITDm = (1e6 * IPDm) ./ (2 * pi * cF);
        ILDs = std(ILD_all,0,3);
        ILDm = sum(ILD_all, 3)/Nt;
        gainl = sum(gainl_all, 3)/Nt;
        gainr = sum(gainr_all, 3)/Nt;

        toc
        save(['Shadron_Pena_2022\', date_folder, '\', owl_folder{o} '__' haircut_id{h} '__hrtf_fig5.mat'])
    end
end
%% Some minor editing

for o = 1:5
    for h = 1:2
        clearvars -except h haircut_id o owl_folder date_folder
        load(['Shadron_Pena_2022\', date_folder, '\', owl_folder{o} '__' haircut_id{h} '__hrtf_fig5.mat'])        
        
        %have to wrap ITD's to make things make sense
        %Earlier had to convert 17x8x563 to a 136x563, so now have to revert
        %back to 17x8x20.
        
        itd = size(hrir_l, 3);
        y = size(cF, 2);
            
            IxD.IPDm = zeros(17,8,y);
            IxD.IPDs = zeros(17,8,y);
            IxD.ITDm = zeros(17,8,y);
            IxD.ITDs = zeros(17,8,y);
            IxD.ILDm = zeros(17,8,y);
            IxD.ILDs = zeros(17,8,y);
            IxD.gainl = zeros(17,8,y);
            IxD.gainr = zeros(17,8,y);
            
            for n = 1:y
                for a = 1:17
                    for e = 1:8
                        counter = (a-1)*8 + e;
                        IxD.IPDm(a,e,n) = IPDm(counter,n);
                        IxD.IPDs(a,e,n) = IPDs(counter,n);
                        IxD.ITDm(a,e,n) = ITDm(counter,n);
                        IxD.ITDs(a,e,n) = ITDs(counter,n);
                        IxD.ILDm(a,e,n) = ILDm(counter,n);
                        IxD.ILDs(a,e,n) = ILDs(counter,n);
                        IxD.gainl(a,e,n) = gainl(counter,n);
                        IxD.gainr(a,e,n) = gainr(counter,n);
                    end
                end
            end

        clear n a e count Factor Fs FS idir n n_cF n_dir Nt tt itd ILDm ILDs IPDm IPDp IPDs ITDp ITDs ITDm
        clear samplingrate s1 s2 sL1 sR1 TF1 TF2 Z ax counter fcoef dir s fcoefs hL hR hrir_l hrir_r 
        
        if size(IxD.IPDm,2) > 8 %Loop occurs if spatial locations are > 17 azimuth and > 9 elevation positions
            %Resamples to have locations that are
            %-160:20:160 in azimuth and -60:20:80 in elevation
            for n = 1:y
                for a = 1:2:34
                    for e = 2:2:16
                        IxD2.IPDm((a+1)/2, e/2, n) = IxD.IPDm(a,e,n);
                        IxD2.IPDs((a+1)/2, e/2, n) = IxD.IPDs(a,e,n);
                        IxD2.ITDm((a+1)/2, e/2, n) = IxD.ITDm(a,e,n);
                        IxD2.ITDs((a+1)/2, e/2, n) = IxD.ITDs(a,e,n);
                        IxD2.ILDm((a+1)/2, e/2, n) = IxD.ILDm(a,e,n);
                        IxD2.ILDs((a+1)/2, e/2, n) = IxD.ILDs(a,e,n);
                    end
                end
            end
            
            
            clear IxD
            IxD = IxD2;
            clear IxD2
            
            azimuths = -160:20:160;
            elevations = -60:20:80;
        end
        
        
        save(['Shadron_Pena_2022\', date_folder, '\', owl_folder{o} '__' haircut_id{h} '__stats_fig5.mat'])
        
    end
end
%% Compile all data into one struct
clearvars -except h haircut_id o owl_folder date_folder

condition = [];
for o = [1 2 3 4 5]
    for h = 1:2
        clearvars -except h haircut_id o owl_folder condition date_folder
        load(['Shadron_Pena_2022\', date_folder, '\', owl_folder{o} '__' haircut_id{h} '__stats_fig5.mat'])
        condition.(owl_folder{o}).(haircut_id{h}) = IxD;
         
    end
end
            
clear IxD name

%average the owls
for h=1:2
for o = [1 2 3 4 5]
        if o == 1
            condition.(haircut_id{h}).avg.IPDm = condition.(owl_folder{o}).(haircut_id{h}).IPDm;
            condition.(haircut_id{h}).avg.IPDs = condition.(owl_folder{o}).(haircut_id{h}).IPDs;
            condition.(haircut_id{h}).avg.ITDm = condition.(owl_folder{o}).(haircut_id{h}).ITDm;
            condition.(haircut_id{h}).avg.ITDs = condition.(owl_folder{o}).(haircut_id{h}).ITDs;
            condition.(haircut_id{h}).avg.ILDm = condition.(owl_folder{o}).(haircut_id{h}).ILDm;
            condition.(haircut_id{h}).avg.ILDs = condition.(owl_folder{o}).(haircut_id{h}).ILDs;
            condition.(haircut_id{h}).avg.gainl = condition.(owl_folder{o}).(haircut_id{h}).gainl;
            condition.(haircut_id{h}).avg.gainr = condition.(owl_folder{o}).(haircut_id{h}).gainr;

        else
            condition.(haircut_id{h}).avg.IPDs = condition.(haircut_id{h}).avg.IPDs + condition.(owl_folder{o}).(haircut_id{h}).IPDs;
            condition.(haircut_id{h}).avg.ITDm = condition.(haircut_id{h}).avg.ITDm + condition.(owl_folder{o}).(haircut_id{h}).ITDm;
            condition.(haircut_id{h}).avg.ITDs = condition.(haircut_id{h}).avg.ITDs + condition.(owl_folder{o}).(haircut_id{h}).ITDs;
            condition.(haircut_id{h}).avg.ILDm = condition.(haircut_id{h}).avg.ILDm + condition.(owl_folder{o}).(haircut_id{h}).ILDm;
            condition.(haircut_id{h}).avg.ILDs = condition.(haircut_id{h}).avg.ILDs + condition.(owl_folder{o}).(haircut_id{h}).ILDs;
            condition.(haircut_id{h}).avg.gainl = condition.(haircut_id{h}).avg.gainl + condition.(owl_folder{o}).(haircut_id{h}).gainl;
            condition.(haircut_id{h}).avg.gainr = condition.(haircut_id{h}).avg.gainr + condition.(owl_folder{o}).(haircut_id{h}).gainr;
        end
end
condition.(haircut_id{h}).avg.IPDm = condition.(haircut_id{h}).avg.IPDm / 5;
condition.(haircut_id{h}).avg.IPDs = condition.(haircut_id{h}).avg.IPDs / 5;
condition.(haircut_id{h}).avg.ITDm = condition.(haircut_id{h}).avg.ITDm / 5;
condition.(haircut_id{h}).avg.ITDs = condition.(haircut_id{h}).avg.ITDs / 5;
condition.(haircut_id{h}).avg.ILDm = condition.(haircut_id{h}).avg.ILDm / 5;
condition.(haircut_id{h}).avg.ILDs = condition.(haircut_id{h}).avg.ILDs / 5;
condition.(haircut_id{h}).avg.gainl = condition.(haircut_id{h}).avg.gainl/5;
condition.(haircut_id{h}).avg.gainr = condition.(haircut_id{h}).avg.gainr/5;
end



%% Code for constructing Figure 5, model neurons

%model neuron input parameters 
freq = 2000:1000:7000;
width = length(freq);
sigma = length(freq);
itd = -800:800;
ild = -20:0.01:20;

%tuning for model neurons
itd_peak = 0;
ild_peak = 0;

%max spike count 
A = 10;

cF_ind = 9:5:34;

normalITD = zeros(length(azimuths), length(elevations), length(freq));
normalILD = zeros(length(azimuths), length(elevations), length(freq));
ruffcutITD = zeros(length(azimuths), length(elevations), length(freq));
ruffcutILD = zeros(length(azimuths), length(elevations), length(freq));

for p = 1:length(freq)
    
    %model ITD curve
    itd_y = A* (exp(cos( (2*pi*freq(p)/1000000) *(itd - itd_peak))) - exp(-1)      ) / (exp(1) - exp(-1));
    
    %model ILD curve
    ild_y = 10* exp( -(8)^-1* (ild - ild_peak).^2);

    %determine spike count given the ITD from HRTFs
    for q = 1:numel(condition.normal.avg.ITDm(:,:,p))
        [row, col] = ind2sub([17,8], q);
        normalITD(row,col,p) = itd_y(itd == round(condition.normal.avg.ITDm(row,col, cF_ind(p) )  )  );
        normalILD(row,col,p) = ild_y(round(ild,2) == round(condition.normal.avg.ILDm(row,col,cF_ind(p)  ),2));
        ruffcutITD(row,col,p) = itd_y(itd == round(condition.ohne.avg.ITDm(row,col, cF_ind(p) )  )  );
        ruffcutILD(row,col,p) = ild_y(round(ild,2) == round(condition.ohne.avg.ILDm(row,col,cF_ind(p)  ),2));

    end
    
end
normaltun = (normalITD.*normalILD)*.01;
ruffcuttun = (ruffcutITD.*ruffcutILD)*.01;

%% model neuron tuning across itds and freqs, normal HRTF

normalmap = figure;
set(gcf,'Position',[100 100 900 300]);

sgtitle('Normal', FontSize=20)


for plotID = 2:length(freq)
    subplot(1,length(freq), plotID-1)
    imagesc(azimuths(6:12), elevations(2:6), normaltun(6:12,2:6,plotID)'); axis xy;  caxis([0 1]); colormap jet;
    title([num2str(freq(plotID)/1000), ' kHz'], FontSize=16);
    xticks([azimuths(6:3:12)])
    yticks([elevations(2:2:6)])

    ax = gca;
    ax.FontSize = 16;
    if plotID == 2
        ylabel('Elevation (°)', FontSize=16)
        e = text(-200,90, 'b');
        e.FontSize = 20;
    end
    if plotID == 2
        xlabel('Azimuth (°)', FontSize=16);

    end
end

subplot(1,length(freq), length(freq))
cc = imagesc(azimuths, elevations, normaltun(:,:,plotID)'); axis xy;  caxis([0 1]); colormap jet;
c = colorbar('manual', 'Position', [.82 0.25 0.02, 0.5]);
c.FontSize = 16;
d = text(0, -90,'Norm. Spike Count');
d.FontSize = 16; d.Rotation = 90; 
set(cc,'Visible', 'off')
set(get(cc,'Children'),'Visible','off');


axis off
%saveas(normalmap, 'F:\Various Files\My Papers\Shadron and Pena 2022\eLife Figures\normalmap.png')


%% model neuron tuning across itds and freqs, ruff-removed HRTF

ruffcutmap = figure;

set(gcf,'Position',[100 100 900 300]);

sgtitle('Ruff-Removed', FontSize=20)


for plotID = 2:length(freq)
    subplot(1,length(freq), plotID-1)
    imagesc(azimuths(6:12), elevations(2:6), ruffcuttun(6:12,2:6,plotID)'); axis xy;  caxis([0 1]); colormap jet;
    title([num2str(freq(plotID)/1000), ' kHz'], FontSize=16);
    xticks([azimuths(6:3:12)])
    yticks([elevations(2:2:6)])

    ax = gca;
    ax.FontSize = 16;
    if plotID == 2
        ylabel('Elevation (°)', FontSize=16)
        e = text(-200,90, 'c');
        e.FontSize = 20;
    end
    if plotID == 2
        xlabel('Azimuth (°)', FontSize=16);

    end
end

subplot(1,length(freq), length(freq))
cc = imagesc(azimuths, elevations, ruffcuttun(:,:,plotID)'); axis xy;  caxis([0 1]); colormap jet;
c = colorbar('manual', 'Position', [.82 0.25 0.02, 0.5]);
c.FontSize = 16;
d = text(0, -90,'Norm. Spike Count');
d.FontSize = 16; d.Rotation = 90; 
set(cc,'Visible', 'off')
set(get(cc,'Children'),'Visible','off');

axis off
%saveas(ruffcutmap, 'F:\Various Files\My Papers\Shadron and Pena 2022\eLife Figures\ruffcutmap.png')

%% Other subplots for Figure 5

plotID = 5;
n = 29;

ITDmap = figure;
set(gcf,'Position',[100 100 400 300]);
imagesc(azimuths(6:12), elevations(2:6), condition.normal.avg.ITDm((6:12),(2:6),n)'); axis xy;  caxis([-150 150]);
title(['Frequency: ', num2str(cF(n)/1000), ' kHz'], FontSize=16); 
c = colorbar('XTick',[-150 0 150]); c.FontSize = 12; c.Label.FontSize = 18; colormap jet;
ax = gca; xticks(azimuths(6:12)); yticks(elevation(2:6)); ax.FontSize = 14; ylabel('Elevation (°)'); xlabel('Azimuth (°)')
set(c.XLabel,{'String','Rotation','Position'},{'ITD (μs)',90,[2 0]})
%saveas(ITDmap, 'F:\Various Files\My Papers\Shadron and Pena 2022\Matlab\ITDmap.png')

ILDmap = figure;
set(gcf,'Position',[100 100 400 300]);
imagesc(azimuths(6:12), elevations(2:6), condition.normal.avg.ILDm((6:12),(2:6),n)'); axis xy;  caxis([-16 16]);
title(['Frequency: ', num2str(cF(n)/1000), ' kHz'], FontSize=16); 
c = colorbar('XTick',[-15 0 15]); c.FontSize = 14; c.Label.FontSize = 18; colormap jet;
ax = gca; xticks(azimuths(6:12)); yticks(elevation(2:6)); ax.FontSize = 14; ylabel('Elevation (°)'); xlabel('Azimuth (°)');
set(c.XLabel,{'String','Rotation','Position'},{'ILD (dB)',90,[2 0]})
%saveas(ILDmap, 'F:\Various Files\My Papers\Shadron and Pena 2022\Matlab\ILDmap.png')

ITDcurve = figure;
set(gcf,'Position',[100 100 200 150]);
plot(itd, itd_y, LineWidth=3, Color = 'black');
xlim([-200,200]); ax = gca; ax.FontSize = 12; ylabel('Spike Count'); xlabel('ITD (μs)')
%saveas(ITDcurve, 'F:\Various Files\My Papers\Shadron and Pena 2022\Matlab\ITDcurve.png')

ILDcurve = figure;
set(gcf,'Position',[100 100 200 150]);
plot(ild, ild_y, LineWidth=3, Color = 'black');
xlim([-16,16]); ax = gca; ax.FontSize = 12; ylabel('Spike Count'); xlabel('ILD (dB)');
%saveas(ILDcurve, 'F:\Various Files\My Papers\Shadron and Pena 2022\Matlab\ILDcurve.png')

sixkhz = figure;
set(gcf,'Position',[100 100 400 300]);
imagesc(azimuths(6:12), elevations(2:6), normaltun(6:12,2:6,plotID)'); axis xy; caxis([0 1]); colormap jet;
title('Spatial Tuning', FontSize=16); ax = gca; ax.FontSize = 16;
c = colorbar; c.Label.FontSize = 16; 
ylabel('Elevation (°)'); xlabel('Azimuth (°)'); xticks(azimuths(6:12)); yticks(elevation(2:6));
set(c.XLabel,{'String','Rotation','Position'},{'Norm. Spike Count',90,[2.5 .5]})
%saveas(sixkhz, 'F:\Various Files\My Papers\Shadron and Pena 2022\Matlab\spatialtun.png')