% Analyze differences between high, medium and low recurrence clusters on a
% patient basis using results generated with NMPC data with 3 clusters with
% mine_cdi_model and analyze_cdi_model_samples_3clusters
% Michael A, Henson, 5/21/2020

% Find number of recurrent patients in high, medium and low recurrence
% patient groups and all patient group allowing overlap between groups
numPatMax = length(patMax);
numPatMed = length(patMed);
numPatMin = length(patMin);
recurNumPat = zeros(numCl+1,1);
for i = 1:numPatMax
    idx = find(strcmp(patMax(i),all16S(2,:)));
    if strcmp(all16S(6,idx(1)),'recurrent') == 1 || strcmp(all16S(6,idx(1)),'reinfection') == 1
        recurNumPat(1) = recurNumPat(1)+1;
    end
end
for i = 1:numPatMed
    idx = find(strcmp(patMed(i),all16S(2,:)));
    if strcmp(all16S(6,idx(1)),'recurrent') == 1 || strcmp(all16S(6,idx(1)),'reinfection') == 1
        recurNumPat(2) = recurNumPat(2)+1;
    end
end
for i = 1:numPatMin
    idx = find(strcmp(patMin(i),all16S(2,:)));
    if strcmp(all16S(6,idx(1)),'recurrent') == 1 || strcmp(all16S(6,idx(1)),'reinfection') == 1
        recurNumPat(3) = recurNumPat(3)+1;
    end
end
patTot = union(union(patMin,patMed,'stable'),patMax,'stable');
numPatTot = length(patTot);
for i = 1:numPatTot
    idx = find(strcmp(patTot(i),all16S(2,:)));
    if strcmp(all16S(6,idx(1)),'recurrent') == 1 || strcmp(all16S(6,idx(1)),'reinfection') == 1
        recurNumPat(4) = recurNumPat(4)+1;
    end
end

% Test significance of number of recurrent patients in each patient group
% allowing overlap between groups
numPat = [numPatMax; numPatMed; numPatMin; numPatTot];
for i = 1:numCl+1
    for j = i:numCl+1
        occurData = [recurNumPat(i) numPat(i)-recurNumPat(i); recurNumPat(j) numPat(j)-recurNumPat(j)];
        [h,pPat(i,j),stats] = fishertest(occurData);
        if pPat(i,j) < 0.05
            display(['Recurrent samples for patient groups ' num2str(i), ',' num2str(j) ' significant with ' num2str(pPat(i,j))])
        end
    end
end

% Plot number of recurrent, norecurrent and total patients in each patient 
% group allowing overlap between groups
figure(25)
bar([recurNumPat numPat-recurNumPat numPat])
xLabels7 = { 'High'; 'Medium'; 'Low'; 'All' };
set(gca,'XTick',1:length(xLabels7),'XTickLabel',xLabels7);
xlabel('Recurrence cluster')
ylabel('Number of patients')
legend('Recurrent','Nonrecurrent','Total')
title('Recurrence in Patients Based on Clustered Samples')
set(gca,'FontSize',12)

% Find post-index samples belonging to patients with at least one
% post-index sample in the high recurrence cluster
idxPatMax = [];
patSampMax = [];
for i = 1:numPatMax
    temp = find(strcmp(patMax(i),text16S(2,:)))-1;
    idxPatMax = [idxPatMax temp];
    nS = length(temp);
    sampLabel = num2str([0:nS-1]');
    for j = 1:nS
        patSampMax = [patSampMax strcat(strcat(patMax(i),'-'),num2str(sampLabel(j)))];
    end
end

% Find indices, abundances and labels for post-index samples only to define
% the high recurrence patient group
idx = find(data16S(6,idxPatMax) > 0);
idxPatMaxP = idxPatMax(idx);
abdPatMaxP = abd16S(:,idxPatMaxP);
patSampMax = erase(patSampMax,'EIN00');
patSampMaxP = patSampMax(idx);

% Plot top taxa for post-index samples from patients belonging to the high
% recurrence group as Figure S8 in paper
figure(26)
numPltTaxa = [1:7];
colormap(parula)
h1 = heatmap(abdPatMaxP(numPltTaxa,:),'Colormap',parula);
h1.YData = xLabels1(numPltTaxa);
h1.XData = patSampMaxP;
h1.XLabel = 'Patient samples';

% Partition samples in high recurrence group into three subsets: Prior to
% first sample appearing in high recurrence cluster; During stay in high
% recurrence cluster; After emergence from high recurrence cluster
idxPatMaxPr = [];
idxPatMaxDu  = [];
idxPatMaxAf = [];
for i = 1:numPatMax
    temp = find(strcmp(patMax(i),text16S(2,:)))-1;
    nS = length(temp);
    idxPatSampMaxI(i,1:nS) = temp;
    clFlag = 0;
    for j = 1:nS
        if ~isempty(intersect(temp(j),idxClMax))
            clFlag = 1;
            idxPatMaxDu = [idxPatMaxDu temp(j)];
        else
            if clFlag == 0
                idxPatMaxPr = [idxPatMaxPr temp(j)];
            else
                idxPatMaxAf = [idxPatMaxAf temp(j)];
            end
        end
    end
end

% Assign a cluster number to each sample for every patient: 1 = high
% recurrence cluster; 2 = medium recurrence cluster; 3 = low recurrence
% cluster; 4 - not previously clustered due to low total abundance
numCol = 6;
clPatSamp = zeros(numPatMax,numCol);
for i = 1:numPatMax
    temp = nonzeros(idxPatSampMaxI(i,2:end));
    for j = 1:length(temp)
        if ~isempty(intersect(temp(j),idxClMax))
            clPatSamp(i,numCol-j+1) = 1;
        elseif ~isempty(intersect(temp(j),idxClMed))
            clPatSamp(i,numCol-j+1) = 2;
        elseif ~isempty(intersect(temp(j),idxClMin))
            clPatSamp(i,numCol-j+1) = 3;
        else
            clPatSamp(i,numCol-j+1) = 4;
        end
    end
    for j = 1:6
        if clPatSamp(i,j) == 0
          clPatSamp(i,j) = NaN;
        end
    end
end

% Plot the progression of samples between the three partitioned patient
% sample groups for each patient
figure(27)
colormap(parula)
h2 = heatmap(clPatSamp','Colormap',parula);
h2.XData = erase(patMax,'EIN00');
h2.XLabel = 'Patient';
h2.YData = {'6','5','4','3','2','1'};
h2.YLabel = 'Sample number';
h2.ColorbarVisible = 'off';
h2.Title = 'Cluster Progression in High Recurrence Patient Group';
set(gca,'FontSize',12)

% Plot average abundances of selected taxa in the three partitioned
% patient sample groups
figure(28)
bar([mean(abd16S(numPltTaxa,idxPatMaxPr)'); mean(abd16S(numPltTaxa,idxPatMaxDu)'); mean(abd16S(numPltTaxa,idxPatMaxAf)')]')
set(gca,'XTick',1:length(numPltTaxa),'XTickLabel',xLabels1(numPltTaxa));
xtickangle(25)
ylabel('Average abundance')
legend("Prior, " + length(idxPatMaxPr) + " samples",...
    "During, " + length(idxPatMaxDu) + " samples",...
    "After, " + length(idxPatMaxAf) + " samples")
title('Patients Entering and Leaving High Recurrence Cluster')
set(gca,'FontSize',12)

% Find significant differences between selected taxa in the three
% partitioned patient sample groups
abdMaxSig = [];
for i = 1:3
    for j = i:3
        for k = numPltTaxa
            if i == 1
                numSet1 = abd16S(k,idxPatMaxPr);
            elseif i == 2
                numSet1 = abd16S(k,idxPatMaxDu);
            else
                numSet1 = abd16S(k,idxPatMaxAf);
            end
            if j == 1
                numSet2 = abd16S(k,idxPatMaxPr);
            elseif j == 2
                numSet2 = abd16S(k,idxPatMaxDu);
            else
                numSet2 = abd16S(k,idxPatMaxAf);
            end
            pTaxaMax(i,j,k) = ranksum(numSet1,numSet2);
        end
    end
end

% Need to add code for calculate FDR for abdMaxSig!!!

% Calculate geometric means of all samples using only sufficiently positive
% abundances
spAbd = abd16S';
numSamp = size(spAbd,1);
numTaxa = size(spAbd,2);
clr = zeros(numSamp,numTaxa);
for i = 1:numSamp
    xt1 = spAbd(i,:);
    xt2 = xt1(xt1>1e-3);
    nDim(i) = length(xt2);
    gmean(i) = prod(xt2)^(1/nDim(i));
    clr(i,:) = spAbd(i,:)/gmean(i);
end

% Calculate rho correlation values for all taxa pairs for all samples in
% high recurrent group
idxSet = idxPatMax;
rho = zeros(numTaxa,numTaxa);
for i = 1:numTaxa
    for j = 1:numTaxa
        data1 = clr(idxSet,i);
        data2 = clr(idxSet,j);
        rho(i,j) = 1-var(data1-data2)/(var(data1)+var(data2));
    end
end

% Plot rho values values above a threshold for taxa of interest
figure(29)
iSp = [4];
for i = 1:length(iSp)
    idxP = find(abs(rho(iSp(i),:)) > 0.01);
    [rhos,idx] = sort(rho(iSp(i),idxP));
    barh(rhos(1:end-1))
    xlabel('\rho value')
    set(gca,'YTick',1:length(idxP)-1,'YTickLabel',xLabels1(idxP(idx(1:end-1))));
    title("Correlations with " + xLabels1(idxP(idx(end))))
end

% Perform PCA on NMPC data using all samples in high recurrence patient
% group
idxPatMaxPrt = [idxPatMaxPr idxPatMaxDu idxPatMaxAf];
MpPat = dataMp(:,idxPatMaxPrt)';
[coeff1,score1,latent1,tsquared1,explained1,mu1] = pca(MpPat);
nSampPat = size(MpPat,1);

% Calculate PCA components for all samples in high recurrence patient group
% both the three partitioned groups and the three clusters
numClPat = 3;
xPcaPat = zeros(nSampPat,numClPat+1);
yPcaPat = zeros(nSampPat,numClPat+1);
numPatMaxPr = length(idxPatMaxPr);
numPatMaxDu = length(idxPatMaxDu);
numPatMaxAf = length(idxPatMaxAf);
xPcaPat(1:numPatMaxPr,1) = score1(1:numPatMaxPr,1);
yPcaPat(1:numPatMaxPr,1) = score1(1:numPatMaxPr,2);
xPcaPat(1:numPatMaxDu,2) = score1(numPatMaxPr+1:numPatMaxPr+numPatMaxDu,1);
yPcaPat(1:numPatMaxDu,2) = score1(numPatMaxPr+1:numPatMaxPr+numPatMaxDu,2);
xPcaPat(1:numPatMaxAf,3) = score1(numPatMaxPr+numPatMaxDu+1:end,1);
yPcaPat(1:numPatMaxAf,3) = score1(numPatMaxPr+numPatMaxDu+1:end,2);

% Plot PCA results
figure(30)
scatter(xPcaPat(1:numPatMaxPr,1),yPcaPat(1:numPatMaxPr,1),150,'bs')
hold on
scatter(xPcaPat(1:numPatMaxDu,2),yPcaPat(1:numPatMaxDu,2),150,'ro')
scatter(xPcaPat(1:numPatMaxAf,3),yPcaPat(1:numPatMaxAf,3),150,[0.9290 0.6940 0.1250],'d')
hold off
xlabel("Principal Component 1 (" + explained1(1)+ "% variance)");
ylabel("Principal Component 2 (" + explained1(2)+ "% variance)")
ylim([-1200 700]);
legend("Prior, " + length(idxPatMaxPr) + " samples",...
    "During, " + length(idxPatMaxDu) + " samples",...
    "After, " + length(idxPatMaxAf) + " samples")
title('Patients Entering and Leaving High Recurrence Cluster')
set(gca,'FontSize',12)

% Find significant differences between NMPCs in the three partitioned
% patient sample groups
nmpcMaxSig = [];
for i = 1:3
    for j = i:3
        for k = 1:size(dataMp,1)
            if i == 1
                numSet1 = dataMp(k,idxPatMaxPr);
            elseif i == 2
                numSet1 = dataMp(k,idxPatMaxDu);
            else
                numSet1 = dataMp(k,idxPatMaxAf);
            end
            if j == 1
                numSet2 = dataMp(k,idxPatMaxPr);
            elseif j == 2
                numSet2 = dataMp(k,idxPatMaxDu);
            else
                numSet2 = dataMp(k,idxPatMaxAf);
            end
            [pMetMax(i,j,k),h] = ranksum(numSet1,numSet2,'Alpha',0.05);
            if h == 1 && (mean(numSet1) > 50  || mean(numSet2) > 50) && (mean(numSet1)/mean(numSet2) < 0.5 || mean(numSet1)/mean(numSet2) > 2)
                nmpcMaxSig = [nmpcMaxSig; [i j k pMetMax(i,j,k)]];
            end
        end
    end
end
nSigMax = length(nmpcMaxSig);

% Need to add code for calculate FDR for nmpcMaxSig!!!

% Plot significantly different NMPCs in the three partitioned patient
% sample groups as Figure S9 in paper
figure(31)
sigIdx = unique(nmpcMaxSig(:,3));
nPlt = round(length(sigIdx)/2);
tempLabels = erase(textMp(sigIdx+1,1),"EX_");
tempLabels = erase(tempLabels,"[fe]");
xLabels6 = strrep(tempLabels,'_','-');
dataMaxPr = mean(dataMp(sigIdx,idxPatMaxPr)')';
dataMaxDu = mean(dataMp(sigIdx,idxPatMaxDu)')';
dataMaxAf = mean(dataMp(sigIdx,idxPatMaxAf)')';
subplot(2,1,1)
bar([dataMaxPr(1:nPlt) dataMaxDu(1:nPlt) dataMaxAf(1:nPlt)]);
set(gca, 'XTick', 1:nPlt, 'XTickLabel', xLabels6(1:nPlt));
xtickangle(45)
ylabel('Maximal production (mmol/day)')
legend("Prior, " + length(idxPatMaxPr) + " samples",...
    "During, " + length(idxPatMaxDu) + " samples",...
    "After, " + length(idxPatMaxAf) + " samples")
set(gca,'FontSize',12)
subplot(2,1,2)
bar([dataMaxPr(nPlt+1:end) dataMaxDu(nPlt+1:end) dataMaxAf(nPlt+1:end)]);
set(gca, 'XTick', 1:length(xLabels6(nPlt+1:end)), 'XTickLabel', xLabels6(nPlt+1:end));
xtickangle(45)
ylabel('Maximal production (mmol/day)')
set(gca,'FontSize',12)

% Generate Figure 4 for paper
figure(32)

% Recurrence vs. patient group
subplot(2,2,1)
bar([recurNumPat numPat-recurNumPat numPat])
xLabels9 = { 'HEb group'; 'HBo group'; 'HEc group'; 'All' };
set(gca,'XTick',1:length(xLabels9),'XTickLabel',xLabels9);
ylabel('Number of patients')
legend('Recurrent','Nonrecurrent','Total')

% Taxa correlation with Enterobacteriaceae
subplot(2,2,2)
iSp = 4;
idxP = find(abs(rho(iSp,:)) > 0.031);
[rhos,idx] = sort(rho(iSp,idxP));
bar(rhos(1:end-1))
set(gca,'XTick',1:length(idxP)-1,'XTickLabel',xLabels1(idxP(idx(1:end-1))));
xtickangle(25)
ylabel('Proportionality with Enterobacteriaceae')

% Cluster progression in high Enterobacteriaceae group
subplot(2,2,3)
colormap(parula)
h2 = heatmap(clPatSamp','Colormap',parula);
h2.XData = erase(patMax,'EIN00');
h2.XLabel = 'Patient';
h2.YData = {'6','5','4','3','2','1'};
h2.YLabel = 'Sample number';
h2.ColorbarVisible = 'off';

% Taxa abundances vs. partitioned patient sample groups
subplot(2,2,4)
bar([mean(abd16S(numPltTaxa,idxPatMaxPr)'); mean(abd16S(numPltTaxa,idxPatMaxDu)'); mean(abd16S(numPltTaxa,idxPatMaxAf)')]')
set(gca,'XTick',1:length(numPltTaxa),'XTickLabel',xLabels1(numPltTaxa));
xtickangle(25)
ylabel('Average abundance')
legend("Prior, " + length(idxPatMaxPr) + " samples",...
    "During, " + length(idxPatMaxDu) + " samples",...
    "After, " + length(idxPatMaxAf) + " samples")