clear, close all, clc;

% https://www.mathworks.com/help/stats/spectralcluster.html#mw_4cf06f77-e003-4d30-bce9-b255daa7aa4b

%% Read In Data

patient = 'mitdb/118';

[signal, Fs, tm] = rdsamp(patient);
[RR, tms] = ann2rr(patient, 'atr');
[ann, anntype, subtype, chan, num, comments] = rdann(patient, 'atr');
[hr] = tach(patient,'atr', [], [], size(tms,1)+1);

%% Split into beats of size 300 samples

beat_size = 300;
beats = split_into_beats(signal, tms, ann, anntype, hr, beat_size);

% Extract lead_1 & lead_2 data
lead_1 = zeros(beat_size, length(beats));
lead_2 = zeros(beat_size, length(beats));
for i=1:length(beats)
    temp_mat = cell2mat(beats(i,1));
    lead_1(:, i) = temp_mat(:,1);
    lead_2(:, i) = temp_mat(:,2);
end

% Filter lead_1 and lead_2
for i=1:length(beats)
    lead_1(:, i) = sgolayfilt(lead_1(:, i), 0, 5);
    lead_2(:, i) = sgolayfilt(lead_2(:, i), 0, 5);
end

% TODO extract outlier beats! 

% Plot the overlapped leads side by side
figure;
subplot(1,2,1), plot_ensemble(lead_1), title("Lead 1");
subplot(1,2,2), plot_ensemble(lead_2), title("Lead 2");


%% Spectral Clustering Time

dist_temp = pdist(lead_1');           % Here Rows are samples and columns are entries
dist = squareform(dist_temp);
S = exp(-dist.^2);
issymmetric(S)

k = 3; % Number of clusters
rng('default') % For reproducibility
idx = spectralcluster(S,k,'Distance','precomputed','LaplacianNormalization','symmetric');

% gscatter(X(:,1),X(:,2),idx);
plot_ensemble_colored(lead_1, idx);


%% Functions

function plot_ensemble(EKG)
    hold on
    for k1 = 1:size(EKG,2)
        plot(EKG(:,k1))
    end
    hold off
    xlabel('Time (ms)')
    ylabel('Amplitude (mV)')
    title('Ensemble EKG')
end

function plot_ensemble_colored(EKG, lbl)
    clr = ['r', 'g', 'b'];
    hold on
    for k1 = 1:size(EKG,2)
        plot(EKG(:,k1), Color=clr(lbl(k1)));
    end
    hold off
    xlabel('Time (ms)')
    ylabel('Amplitude (mV)')
    title('Ensemble EKG')
end