% Script to mine adundance data and metadata for 225 CDI samples
% Michael A. Henson, 1/22/2020

clearvars

% Read unnormalized sample data and find samples with "high" unnormalized
% abundance
raw16S = xlsread('cdi_abd_nmpc.xlsx',1,'B23:HR70');
idxHighAll = find(ismember(sum(raw16S) > 0.9, 1));

% Read normalized sample data
[data16S,text16S,all16S] = xlsread('cdi_abd_nmpc.xlsx',2,'A2:HR70');
abd16S = xlsread('cdi_abd_nmpc.xlsx',2,'B23:HR70');

% Find samples of interest
dataIndex = data16S(6,:);
% dataIndex == 1 for index samples; dataIndex > 1 for post-index samples
idxIndex = find(ismember(dataIndex > 0, 1));
idxHigh = intersect(idxIndex,idxHighAll);
X = abd16S(:,idxHigh);

% Find recurrent + reinfection samples
recurSamp = find(strcmp(all16S(6,idxHigh+1),'recurrent'));
recurSamp = sort(union(recurSamp,find(strcmp(all16S(6,idxHigh+1),'reinfection'))));

% Cluster abundance data
% eva = evalclusters(X','linkage','DaviesBouldin','KList',[1:10]);
% numCl = eva.OptimalK;
numCl = 3;
idxSamp = zeros(size(X,2),numCl);
iCl = kmeans(X',numCl,'Replicates',1000);
numCl16S = numCl;
iCl16S = iCl;
for i = 1:numCl
    idx = find(ismember(iCl == i, 1));
    idxSamp(1:length(idx),i) = idx;
end
nSampCl = sum(idxSamp ~= 0);

% Average abundance data over each cluster
abdCl = zeros(numCl,size(X,1));
for i = 1:numCl
    idx = nonzeros(idxSamp(:,i));
    if length(idx) == 1
        abdCl(i,:) = X(:,idx)';
    else
        abdCl(i,:) = mean(X(:,idx)');
    end
end
abdCl = abdCl';

% Find taxa for which average abundance in at least one cluster is above a
% threshold
threshAbd = 0.05;
idxAbd = [];
for i = 1:numCl
    idxAbd = union(idxAbd,find(ismember(abdCl(:,i) > threshAbd, 1)));
end

% Plot selected taxa abundances for each cluster
figure(11)
bar(abdCl(idxAbd,:))
xLabels = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbd),'XTickLabel',xLabels(idxAbd));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCl(1) + " samples", "Cluster 2, " +  nSampCl(2) + " samples",...
    "Cluster 3, " +  nSampCl(3) + " samples")
title('Abundances of Clustered 16S Post-index Samples')
set(gca,'FontSize',12)

% Find numbers of recurrent samples in each cluster and entire dataset
for i = 1:numCl
    recurNum(i) = length(intersect(recurSamp,idxSamp(:,i)));
end
nSampCl(numCl+1) = size(X,2);
recurNum(numCl+1) = length(recurSamp);

% Test clustered recurrent samples for significant differences
for i = 1:numCl+1
    for j = i:numCl+1
        occurData = [recurNum(i) nSampCl(i)-recurNum(i); recurNum(j) nSampCl(j)-recurNum(j)]; 
        [h(i,j),p(i,j),stats] = fishertest(occurData);
        if p(i,j) < 0.05
            display(['Recurrent samples for ' num2str(numCl) ' clusters ' num2str(i), ',' num2str(j) ' significant with ' num2str(p(i,j))])
        end
    end
end

% Plot fraction of recurrent samples in each cluster
figure(12)
recurCl = [recurNum; nSampCl-recurNum; nSampCl]';
bar([recurCl])
xLabels2 = { 'Cluster 1'; 'Cluster 2'; 'Cluster 3'; 'All' };
set(gca,'XTick',1:length(xLabels2),'XTickLabel',xLabels2);
ylabel('Sample number')
legend('Recurrent','Nonrecurrent','Total')
title('Recurrence in Clustered 16S Post-index Samples')
set(gca,'FontSize',12)

% Perform PCA on 16S data
[coeff,score,latent,tsquared,explained,mu] = pca(X');
explainedData = explained;
nSamp = size(X,2);
nrecurSamp = setdiff([1:nSamp],recurSamp);

% Calculate PCA components for nonrecurrent, recurrent and clustered
% samples
xPca = zeros(nSamp,numCl+2);
yPca = zeros(nSamp,numCl+2);
nSampNrecur = length(nrecurSamp);
nSampRecur = length(recurSamp);
xPca(1:nSampNrecur,1) = score(nrecurSamp,1);
yPca(1:nSampNrecur,1) = score(nrecurSamp,2);
xPca(1:nSampRecur,2) = score(recurSamp,1);
yPca(1:nSampRecur,2) = score(recurSamp,2);
for i = 1:numCl
    xPca(1:nSampCl(i),i+2) = score(idxSamp(1:nSampCl(i),i),1);
    yPca(1:nSampCl(i),i+2) = score(idxSamp(1:nSampCl(i),i),2);
end

% Plot PCA results
figure(13)
scatter(xPca(1:nSampNrecur,1),yPca(1:nSampNrecur,1),175,'bs')
hold on
scatter(xPca(1:nSampRecur,2),yPca(1:nSampRecur,2),125,'ro')
text(xPca(1:nSampCl(1),3),yPca(1:nSampCl(1),3),'1','HorizontalAl','center','color','k')
text(xPca(1:nSampCl(2),4),yPca(1:nSampCl(2),4),'2','HorizontalAl','center','color','k')
text(xPca(1:nSampCl(3),5),yPca(1:nSampCl(3),5),'3','HorizontalAl','center','color','k')
hold off
xlabel("Principal Component 1 (" + explained(1)+ "% variance)");
ylabel("Principal Component 2 (" + explained(2)+ "% variance)")
legend('Nonrecurrent','Recurrent')
title('PCA of Clustered 16S Post-index Samples')
set(gca,'FontSize',12)

% Generate Figure S2 or S3 in paper
figure(14)

% Abundance vs. cluster plot
subplot(2,2,1)
bar(abdCl(idxAbd,:))
set(gca,'XTick',1:length(idxAbd),'XTickLabel',xLabels(idxAbd));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCl(1) + " samples", "Cluster 2, " +  nSampCl(2) + " samples",...
    "Cluster 3, " +  nSampCl(3) + " samples")

% Reucurrence vs. cluster plot
subplot(2,2,2)
bar([recurCl])
set(gca,'XTick',1:length(xLabels2),'XTickLabel',xLabels2);
ylabel('Sample number')
legend('Recurrent','Nonrecurrent','Total')

% PCA vs. cluster plot
subplot(2,2,[3 4])
scatter(xPca(1:nSampNrecur,1),yPca(1:nSampNrecur,1),175,'bs')
hold on
scatter(xPca(1:nSampRecur,2),yPca(1:nSampRecur,2),125,'ro')
text(xPca(1:nSampCl(1),3),yPca(1:nSampCl(1),3),'1','HorizontalAl','center','color','k')
text(xPca(1:nSampCl(2),4),yPca(1:nSampCl(2),4),'2','HorizontalAl','center','color','k')
text(xPca(1:nSampCl(3),5),yPca(1:nSampCl(3),5),'3','HorizontalAl','center','color','k')
hold off
xlabel("PC 1 (" + explained(1)+ "% variance)");
ylabel("PC 2 (" + explained(2)+ "% variance)")
legend('Nonrecurrent','Recurrent')