%% METODI STATISTICI PER LA BIOINGEGNERIA - LABORATORIO 10   

clear all
close all
clc

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% PRIMA PARTE KAPLAN_MEIER & LOG-RANK TEST %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 1. CARICAMENTO DEI DATI E CREAZIONE DELLE VARIABILI CORRISPONDENTI AI DUE GRUPPI %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

load(fullfile('data_lab10.mat'))

% Creare variabile G per identificare i gruppi
% GRUPPO 1 – BMI > threshold​
% GRUPPO 2 – BMI <= threshold

G = BMI > threshold;

% Creazione variabili T1, T2, C1, C2

T1 = SurvTime(G);
T2 = SurvTime(~G); 
C1 = C(G);
C2 = C(~G);

patients_1=(1:1:length(T1));
patients_2=(1:1:length(T2));


%% %%%%%%%%%%%%%%%%%%%%%%
% 2. VISUALIZZAZIONE    %
%%%%%%%%%%%%%%%%%%%%%%%%%

% Visualizzare tramite stem plot i due gruppi e plottare eventi e censure
figure
subplot(211)
stem(patients_1,T1,  'Marker','none')
hold on
plot(patients_1(C1==1),T1(C1==1), 'ro', 'LineWidth',4)
plot(patients_1(C1==0),T1(C1==0),'ko', 'LineWidth',4)
hold off
view(90, 90)
legend('Survival time', 'event', 'censored')
ylabel('Survival Time [years]')
xlabel('Patients')
title(['Group BMI>' num2str(threshold)])

subplot(212)
stem(patients_2, T2, 'Marker','none')
hold on
plot(patients_2(C2==1),T2(C2==1),'ro', 'LineWidth',4)
plot(patients_2(C2==0),T2(C2==0),'ko', 'LineWidth',4)
hold off
view(90, 90)
legend('Survival time', 'event', 'censored')
ylabel('Survival Time [years]')
xlabel('Patients')
title(['Group BMI<=' num2str(threshold)])
sgtitle('Group Survival Time')
set(gcf, 'WindowState', 'maximized')


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 3. INIZIALIZZAZIONE KAPLAN-MEIER %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Impostare la griglia dei tempi di Kaplan-meier​
T = unique([T1(C1==1); T2(C2==1)]);
T = [0; T];

% Inizializzare S1 ed S2
S1 = zeros(length(T),1); 
S1(1) = 1;

S2 = zeros(length(T),1); 
S2(1) = 1;

O1 = zeros(length(T)-1,1); 
O2 = zeros(length(T)-1,1); 

E1 = zeros(length(T)-1,1); 
E2 = zeros(length(T)-1,1); 


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 3. Implementazione Kaplan-Meier per ciascun tempo t
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 

for k = 2:length(T)
    
    % Tempo corrente
    t=T(k);
    
    % Numerosità soggetti a rischio
    N1_t = sum(T1(:)>=t);
    N2_t = sum(T2(:)>=t);
    N_t = N1_t + N2_t;
    
    % Valori osservati al tempo t
    O1_t = length(find(T1==t & C1==1));
    O2_t = length(find(T2==t & C2==1));
    O_t = O1_t + O2_t;
    
    % Stime Kaplan-Meier della survival 
    S1(k) = S1(k-1)*(1-O1_t/N1_t);
    S2(k) = S2(k-1)*(1-O2_t/N2_t);
    
    % Mortalità attesa nel gruppo 1 e nel gruppo 2
    E1_t = (O_t/N_t)*N1_t;
    E2_t = (O_t/N_t)*N2_t;
    
    % O1, E1, O2, E2
    O1(k-1) = O1_t;
    O2(k-1) = O2_t;
    E1(k-1) = E1_t;
    E2(k-1) = E2_t;

end


%% %%%%%%%%%%%%%%%%
% 4. LOGRANK TEST %
%%%%%%%%%%%%%%%%%%%

disp(abs(sum(E1)-sum(O1)));
disp(abs(sum(E2)-sum(O2)));
% sono uguali, prendiamo come riferimento O1

Z = (sum(E1)-sum(O1))^2/sum(E1.*E2); %4.1970 > 3.84 rifiuto H0

P = chi2cdf(Z,1,'upper'); %0.0405 <0.05 rifiuto H0

%% %%%%%%%%%%%%%%%%
% 5. SURVIVAL PLOT  %
%%%%%%%%%%%%%%%%%%%

figure(2)
stairs(T, S1)
hold on
stairs(T, S2)
ylim([0 1])
xlim([0 20])
title(['BMI (p: ' num2str(P) ')'])
xlabel('Time [years]')
ylabel('S(t)')

% Marker per i pazienti censurati
pos_cens1 = interp1(T, S1, T1(C1==0),'previous');
pos_cens2 = interp1(T, S2, T2(C2==0),'previous');
scatter(T1(C1==0), pos_cens1,'k+','MarkerEdgeColor',[0, 0.4470-0.1, 0.7410-0.1],'LineWidth',1)
scatter(T2(C2==0), pos_cens2,'k+','MarkerEdgeColor',[0.8500-0.1, 0.3250-0.1, 0.0980],'LineWidth',1)
legend({['Group 1: >' num2str(threshold)], ['Group 2: <=' num2str(threshold)], 'censored G1', 'censored G2'}, 'Location','southwest')

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 3. COX PROPORTIONAL-HAZARD MODEL %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% 3.1 Fit del cox-model --> solo per BMI?
[b,~,~,stats] = coxphfit(double(G),SurvTime,'censoring', ~C); % not C perchè cox considera 1=censored, 0=event

% 3.2 Hazard ratio
HR = exp(b);