%% Train random forest for distinguishing rhythms
%
% Joel Karel
% Maastricht University
% 16-03-2017
data_dir = ['training2017' filesep];
%% load labels training set

reffile = [data_dir 'REFERENCE.csv'];

fid = fopen(reffile, 'r');
if(fid ~= -1)
    Ref = textscan(fid,'%s %s','Delimiter',',');
else
    error(['Could not open ' reffile ' for scoring. Exiting...'])
end
fclose(fid); 

RECORDS = Ref{1};
target  = Ref{2};
N       = length(RECORDS);


% find indices of different groups in REFERENCE file
idx1 = ismember(Ref{2},'N');
idxN=find(idx1);
idx2 = ismember(Ref{2},'A');
idxA=find(idx2);
idx3 = ismember(Ref{2},'O');
idxO=find(idx3);
idx4 = ismember(Ref{2},'~');
idxnoise=find(idx4);

%% generate parameters
nrec=size(RECORDS,1)
x=zeros(nrec,81);
parfor cnt=1:nrec
cnt
    [tm,ecg,fs,siginfo]=rdmat([data_dir RECORDS{cnt}]);
    [B,A] = cheby2(3,20,[0.5/(fs/2) 100/(fs/2)]);
    nf =getNoiseFeatures(ecg); % get the features for the noise
    ecgf = filtfilt(B,A,ecg);

    QRS = findQRScomplexes(ecgf,fs);

    AFfeatures=extractParameters(ecgf,fs,QRS);
    nv=[AFfeatures nf];
    if(size(nv,2)~=81)
	AFfeatures
        nf
    end
    x(cnt,:)=nv;
%    x=[x;[nf]];
end
save huppa
    
%load('~/Documents/Research/CinCChallenge2017/workinprogress/newresults.mat') 
%load parameters
%newparameters(:,end)=[];
%x=[newparameters noiseParameters];


Parameters = {'P_LF'; 'f_LF'; 'P_HF'; 'f_HF'; 'meanRR'; 'SDRR'; 'rMSSDRR'; 'pNN10'; 'pNN20'; ...
    'pNN30'; 'pNN40'; 'pNN50'; 'pNN60'; 'pNN70'; 'pNN80';'pNN90'; 'pNN100'; ...
    'V0'; 'V1'; 'LV2'; 'UV2'; 'numbercomponents95percQRST'; 'SimQRST1'; 'SimQRST2'; ...
    'SimQRST3'; 'SimQRST4'; 'SimQRST5'; 'numbercomponents95percP'; 'SimP1'; 'SimP2'; ...
    'SimP3'; 'SimP4'; 'SimP5'; 'SCN'; 'DF'; 'fWP'; 'fWPMAW'; 'rhe';  'SampEn'; 'SampEnMAW'; ... 
    'OI'; 'SE'; 'FWA'; 'differenceMediaRRmeanRR';...
    'wlvarsc1';'wlvarsc2';'wlvarsc3';'wlvarsc4';'wlvarsc5';...
    'wlvarsc6';'locMaxSing1';'locMaxSing2';'locMaxSing3';'dh1';'dh2';'dh3';'dh4';'dh5';'dh6';'dh7';'dh8';'dh9';'dh10';'dh11';'dh1';'dh2';'dh3';'dh4';'dh5';'dh6';'dh7';'dh8';'dh9';'dh10';'dh11';'maxAutocorr';...
    'autocorr95quant';'autocorr90quant';'avAutocorr';'mad';'var'};


%% convert labels to numerical values
Labels = char(Ref{2});
y = double(Labels);

y(y==78) = 1;
y(y==79) = 3;
y(y==65) = 2;
y(y==126) = 4;


%% Split in training and validation set - 90% vs 10%

indx=rand(size(x,1),1)<0.90;
train.x=x(indx,:);
train.y=y(indx);
train.Labels=Labels(indx);
val.x=x(~indx,:);
val.y=y(~indx);
val.Labels=Labels(~indx);

rng(1); % For reproducibility


%% Learn ensembles with RUSBoost to handle class imbalance
% Note that the results vary from training to training

% see how much each class is represented
tabulate(train.Labels)
tabulate(val.Labels)
tabulate(Labels)

% fit the ensemble with deep trees
t = templateTree('MinLeafSize',5);
tic
rusTree = fitensemble(train.x,train.Labels,'RUSBoost',1000,t,...
    'LearnRate',0.1,'nprint',100);
toc

% Plot the classification error against the number of members in the ensemble
figure;
tic
plot(loss(rusTree,val.x,val.Labels,'mode','cumulative'));
toc
grid on;
xlabel('Number of trees');
ylabel('Test classification error');

% Examine the confusion matrix for each class as a percentage of the true class.
tic
Yfit = predict(rusTree,val.x);
toc
tab = tabulate(val.y);
bsxfun(@rdivide,confusionmat(int8(val.Labels),int8(Yfit)),tab(:,2))*100

% The ensemble is large. Remove the data using the compact method.
cmpctRus = compact(rusTree);

sz(1) = whos('rusTree');
sz(2) = whos('cmpctRus');
[sz(1).bytes sz(2).bytes]

% Remove half the trees from cmpctRus. This action is likely to have minimal effect on the predictive performance, based on the observation that 400 out of 1000 trees give nearly optimal accuracy.
cmpctRus = removeLearners(cmpctRus,[500:1000]);
sz(3) = whos('cmpctRus');
sz(3).bytes

% The reduced compact ensemble takes about a quarter the memory of the full ensemble. Its overall loss rate is under 24%:
tic
Yfit = predict(cmpctRus,val.x);
toc
tab = tabulate(val.y);
confusionmat(int8(val.Labels),int8(Yfit))

% compute scoring on validation subset of training set
AA=zeros(4,4);
n=sum(indx==0);
valRecords=find(indx==0);
Answers=zeros(size(n,1));
% classify the validation set
parfor cnt=1:n
    Answers(cnt) = char(challenge([data_dir RECORDS{valRecords(cnt)}]));
end
% score the answers
for cnt=1:n
    this_answer=char(Answers(cnt));
    target=char(Labels(valRecords(cnt)));
    switch target
        case 'N'
            if strcmp(this_answer,'N')
                AA(1,1) = AA(1,1)+1;
            elseif strcmp(this_answer,'A')
                AA(1,2) = AA(1,2)+1;
            elseif strcmp(this_answer,'O')
                AA(1,3) = AA(1,3)+1;
            elseif strcmp(this_answer,'~')
                AA(1,4) = AA(1,4)+1;
            end
        case 'A'
            if strcmp(this_answer,'N')
                AA(2,1) = AA(2,1)+1;
            elseif strcmp(this_answer,'A')
                AA(2,2) = AA(2,2)+1;
            elseif strcmp(this_answer,'O')
                AA(2,3) = AA(2,3)+1;
            elseif strcmp(this_answer,'~')
                AA(2,4) = AA(2,4)+1;
            end
        case 'O'
            if strcmp(this_answer,'N')
                AA(3,1) = AA(3,1)+1;
            elseif strcmp(this_answer,'A')
                AA(3,2) = AA(3,2)+1;
            elseif strcmp(this_answer,'O')
                AA(3,3) = AA(3,3)+1;
            elseif strcmp(this_answer,'~')
                AA(3,4) = AA(3,4)+1;
            end
        case '~'
            if strcmp(this_answer,'N')
                AA(4,1) = AA(4,1)+1;
            elseif strcmp(this_answer,'A')
                AA(4,2) = AA(4,2)+1;
            elseif strcmp(this_answer,'O')
                AA(4,3) = AA(4,3)+1;
            elseif strcmp(this_answer,'~')
                AA(4,4) = AA(4,4)+1;
            end
    end

end

F1n=2*AA(1,1)/(sum(AA(1,:))+sum(AA(:,1)));
F1a=2*AA(2,2)/(sum(AA(2,:))+sum(AA(:,2)));
F1o=2*AA(3,3)/(sum(AA(3,:))+sum(AA(:,3)));
F1p=2*AA(4,4)/(sum(AA(4,:))+sum(AA(:,4)));
F1=(F1n+F1a+F1o+F1p)/4;


str = ['F1 measure for Normal rhythm:  ' '%1.4f\n'];
fprintf(str,F1n)
str = ['F1 measure for AF rhythm:  ' '%1.4f\n'];
fprintf(str,F1a)
str = ['F1 measure for Other rhythm:  ' '%1.4f\n'];
fprintf(str,F1o)
str = ['F1 measure for Noisy recordings:  ' '%1.4f\n'];
fprintf(str,F1p)
str = ['Final F1 measure:  ' '%1.4f\n'];
fprintf(str,F1)

%save model cmpctRus
save ModelResultsAug

%% predictor importance
imp = predictorImportance(cmpctRus);
[B,I] = sort(imp,'descend');
% create sorted parameter labels
ParametersSorted=cell(length(I),1);
for cnt=1:length(I)
    ParametersSorted{cnt}=Parameters{I(cnt)};
end
bar(B);
ax=gca;
ax.XTick=1:length(imp);
ax.XTickLabel=ParametersSorted;
ax.XTickLabelRotation=60;

save extrastuff