% Script to mine NMPC predictions 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
idxPindex = find(ismember(dataIndex > 0, 1));
idxHigh = intersect(idxPindex,idxHighAll);
X = abd16S(:,idxHigh);

% Read NMPC data
[dataMp,textMp,allMp] = xlsread('cdi_abd_nmpc.xlsx',3,'A1:HR412');
Mp = dataMp(:,idxHigh);

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

% Determine optimal cluster number and cluster NMPC data
% eva = evalclusters(Mp','linkage','DaviesBouldin','KList',[1:20]);
% numCl = eva.OptimalK;
numCl = 3;
idxSamp = zeros(size(Mp,2),numCl);
iCl = kmeans(Mp',numCl,'Replicates',1000);
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(15)
bar(abdCl(idxAbd,:))
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbd),'XTickLabel',xLabels1(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 Model 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 number of recurrent, nonrecurrent and total samples in each cluster
figure(16)
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 Model Post-index Samples')
set(gca,'FontSize',12)

% Perform PCA on NMPC data
[coeff,score,latent,tsquared,explained,mu] = pca(Mp');
nSamp = size(Mp,2);

% 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(17)
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, " +  nSampNrecur + " samples", "Recurrent, " +  nSampRecur + " samples")
title('PCA of Model Post-index Samples')
set(gca,'FontSize',12)

% Cluster NMPC data for 2-6 clusters
numCls = [2:6];
iCls = zeros(nSamp,length(numCls));
idxSampCls = zeros(length(numCls),numCls(end),size(X,2));
abdCls = zeros(length(numCls),numCls(end),size(X,1));
nSampCls = zeros(length(numCls),numCls(end)+1);
for i = 1:length(numCls)
    if numCls(i) == numCl
        iCls(:,i) = iCl;
    else
        iCls(:,i) = kmeans(Mp',numCls(i),'Replicates',1000);
    end
    for j = 1:numCls(i)
        idx = find(iCls(:,i) == j);
        idxSampCls(i,j,1:length(idx)) = idx;
        abdCls(i,j,:) = mean(X(:,idx)');
    end
    nSampCls(i,1:end-1) = sum(squeeze(idxSampCls(i,:,:))' ~= 0);
end

% Find numbers of recurrent samples in each cluster and entire dataset
recurNumCls = zeros(length(numCls),numCls(end)+1);
for i = 1:length(numCls)
    for j = 1:numCls(i)
        recurNumCls(i,j) = length(intersect(recurSamp,squeeze(idxSampCls(i,j,:))'));
    end
    recurNumCls(i,j+1) = length(recurSamp);
    nSampCls(i,j+1) = size(X,2);
end

% Test clustered recurrent samples for significant differences
pCls = zeros(length(numCls),numCls(end)+1,numCls(end)+1);
for i = 1:length(numCls)
    for j = 1:numCls(i)+1
        for k = j:numCls(i)+1
            occurData = [recurNumCls(i,j) nSampCls(i,j)-recurNumCls(i,j); recurNumCls(i,k) nSampCls(i,k)-recurNumCls(i,k)]; 
            [h,pCls(i,j,k),stats] = fishertest(occurData);
            if pCls(i,j,k) < 0.05
                display(['Recurrent samples for ' num2str(numCls(i)) ' clusters ' num2str(j), ',' num2str(k) ' significant with ' num2str(pCls(i,j,k))])
            end
        end
    end
end

% Find taxa for which average abundance in at least one cluster is above a
% threshold
threshAbds = 0.15;
idxAbds = [];
for i = 2:length(numCls)
    for j = 1:numCls(i)
        idxAbds = union(idxAbds,find(ismember(squeeze(abdCls(i,j,:)) > threshAbds, 1)));
    end
end

% Plot selected taxa abundances for 2-6 clusters for Figure S4 in paper
figure(18)
subplot(2,3,1)
bar(squeeze(abdCls(1,1:numCls(1),idxAbds))')
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbds),'XTickLabel',xLabels1(idxAbds));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCls(1,1) + " samples", "Cluster 2, " + nSampCls(1,2) + " samples")

subplot(2,3,2)
bar(squeeze(abdCls(2,1:numCls(2),idxAbds))')
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbds),'XTickLabel',xLabels1(idxAbds));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCls(2,1) + " samples", "Cluster 2, " + nSampCls(2,2) + " samples",...
    "Cluster 3, " +  nSampCls(2,3) + " samples")

subplot(2,3,3)
bar(squeeze(abdCls(3,1:numCls(3),idxAbds))')
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbds),'XTickLabel',xLabels1(idxAbds));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCls(3,1) + " samples", "Cluster 2, " + nSampCls(3,2) + " samples",...
    "Cluster 3, " +  nSampCls(3,3) + " samples", "Cluster 4, " +  nSampCls(3,4) + " samples")

subplot(2,3,4)
bar(squeeze(abdCls(4,1:numCls(4),idxAbds))')
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbds),'XTickLabel',xLabels1(idxAbds));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCls(4,1) + " samples", "Cluster 2, " + nSampCls(4,2) + " samples",...
    "Cluster 3, " +  nSampCls(4,3) + " samples", "Cluster 4, " +  nSampCls(4,4) + " samples",...
    "Cluster 5, " +  nSampCls(4,5) + " samples")

subplot(2,3,5)
bar(squeeze(abdCls(5,1:numCls(5),idxAbds))')
xLabels1 = erase(text16S(22:end,1),"pan");
set(gca,'XTick',1:length(idxAbds),'XTickLabel',xLabels1(idxAbds));
xtickangle(25)
ylabel('Average abundance')
legend("Cluster 1, " +  nSampCls(5,1) + " samples", "Cluster 2, " + nSampCls(5,2) + " samples",...
    "Cluster 3, " +  nSampCls(5,3) + " samples", "Cluster 4, " +  nSampCls(5,4) + " samples",...
    "Cluster 5, " +  nSampCls(5,5) + " samples", "Cluster 6, " +  nSampCls(5,6) + " samples")

% Eliminate "total" data and fraction of recurrent samples for all clusters
subplot(2,3,6)
tempData = recurNumCls./nSampCls;
for i = 1:length(numCls)
    tempData(i,i+2) = NaN;
end
bar(tempData(:,1:end-1))
xLabels3 = {'2', '3', '4', '5', '6'}; 
set(gca,'XTick',1:length(xLabels3),'XTickLabel',xLabels3);
xlabel('Number of clusters')
ylabel('Fraction recurrent samples')

% Determine how samples have moved between clusters
iSampCls = zeros(length(numCls)-1,numCls(end-1),numCls(end));
for i = 1:length(numCls)-1
    for j = 1:numCls(i)
        for k = 1:numCls(i+1)
            iSampCls(i,j,k) = sum(iCls(:,i) == j & iCls(:,i+1) == k);
        end
    end
end

% Generate figures for paper
figure(19)

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

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

% PCA vs. cluster
subplot(2,2,3)
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, " +  nSampNrecur + " samples", "Recurrent, " +  nSampRecur + " samples")

% Samples vs. cluster
subplot(2,2,4)

% % Find numbers of samples shared by 16S and model clusters for Figure 1 in paper
% % Must first generate 16S results with mine_cdi_data, set iCl16S = iCl and
% % do not clear workspace in this code
% iClMat = [];
% for i = 1:numCl16S
%     for j = 1:numCl
%         iClMat(i,j) = sum(iCl16S == i & iCl == j);
%     end
% end

% % Create heatmap of numbers of shared samples
% colormap(parula)
% h1 = heatmap(iClMat,'Colormap',parula);
% h1.XLabel = 'Model index cluster';
% h1.YLabel = '16S index cluster';

% % Plot explained variance for PCA of data and model data sets for Figure 2 in paper
% % Must first generate 16S results with mine_cdi_data, set explainedData = explained
% % and do not clear workspace in this code
% explainedModel = explained;
% stairs(cumsum(explainedData))
% hold on
% stairs(cumsum(explainedModel))
% hold off
% xlim([1 15])
% xlabel('Number of principal components')
% ylabel('Variance explained (%)')
% legend("Data, " +  length(explainedData) + " components", "Model, " +  length(explainedModel) + " components");