function [type,beats,beatsuct] = beat_class(ecg,qrs,T_LENGTH)
% this function is used to contruct a template ecg based on the location of
% the R-peaks. A series of peaks that match with each other are stacked to
% build a template. This template can then be used for ecg morphological
% analysis or further processing. Note that the qrs location inputed must
% be as precise as possible. This approach for building the template ECG
% seems to be the best of the alternatives that were tested and leaves the
% freedom of having more than one mode (i.e. multiple ECG template can be
% built if there are different cycle morphology such as PVC)  but it is not
% particularly fast.
% The procedure for building the template is:
% 1. create average wrapped template
% 2. identify different modes present
%
% inputs
%   ecg:            the ecg channel(s)
%   qrs:            qrs location [number of samples]
%   T_LENGTH        template length
% 
% --
% ECG classification from single-lead segments using Deep Convolutional Neural 
% Networks and Feature-Based Approaches - December 2017
% 
% Released under the GNU General Public License
%
% Copyright (C) 2017  Fernando Andreotti, Oliver Carr
% University of Oxford, Insitute of Biomedical Engineering, CIBIM Lab - Oxford 2017
% fernando.andreotti@eng.ox.ac.uk
%
% 
% For more information visit: https://github.com/fernandoandreotti/cinc-challenge2017
% 
% Referencing this work
%
% Andreotti, F., Carr, O., Pimentel, M.A.F., Mahdi, A., & De Vos, M. (2017). 
% Comparing Feature Based Classifiers and Convolutional Neural Networks to Detect 
% Arrhythmia from Short Segments of ECG. In Computing in Cardiology. Rennes (France).
%
% Last updated : December 2017
% 
% This program is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
% 
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
% 
% You should have received a copy of the GNU General Public License
% along with this program.  If not, see <http://www.gnu.org/licenses/>.


% == manage inputs
if nargin<2; error('ecg_template_build: wrong number of input arguments \n'); end;

if size(ecg,1)>size(ecg,2), ecg = ecg';end

% == constants
% qrs(1) = []; qrs(end) = [];                                % remove extremity peaks
extremities = (qrs <= 2*T_LENGTH | qrs >= length(ecg)-2*T_LENGTH);        % test if there are peaks on the border that may lead to error
qrs = round(qrs(~extremities)); % remove extremity peaks

%% Phase wrap to get all within 2*pi
NB_BINS = 250; % number of bins for wrapping
[M,RES] = stackbeats(ecg,qrs,T_LENGTH,NB_BINS);
    

% beatsuct = median(M')';
% dem = round((length(beatsuct)-2*T_LENGTH-1)/2); 
% beats{1} =  beatsuct(dem+1:end-dem);
% type = ones(length(qrs),1);

%% Clustering
Nclus1 = 5;
Nclus2 = 2;
M = bsxfun(@times,M,hamming(size(M,1)));  % do it twice, it's the
M = bsxfun(@times,M,hamming(size(M,1)));  % magic sauce
M = bsxfun(@times,M,hamming(size(M,1)));  % magic sauce
M = M';
% [~, clusters, ~] = fkmeans(M',Nclus);
rng(1);
[~,clusters,~] = kmeans(M,Nclus1); % may need to provide centroid..
[~,clusters,~] = kmeans(clusters,Nclus2); % may need to provide centroid..
clusters = clusters';

% !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
% TODO: allow more than two clusters!!!!!!!!!!!!!!!!!!!!
% !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

corel = corrcoef(clusters(50:200,1),clusters(50:200,2));


%% In case beats are correlated to each other
if corel(1,2) > 0.7 
    type = ones(size(qrs));
    beats{1} = resample(mean([clusters(:,1), clusters(:,2)],2),RES,NB_BINS);    
    dem = (length(beats{1})-2*T_LENGTH-1)/2;   
    beatsuct = beats;
    beats = cellfun(@(x) x(dem+1:end-dem),beats,'UniformOutput',0);
    return;
end

%% In case one cluster has low amplitude

idx = var(clusters)<0.1*var(ecg); % remove clusters with variance less than 10% of signal
if any(idx)
   type = ones(size(qrs));
   beats{1} = resample(clusters(:,~idx),RES,NB_BINS);
   dem = round((length(beats{1})-2*T_LENGTH-1)/2);      
   beatsuct = beats;
   beats = cellfun(@(x) x(dem+1:end-dem),beats,'UniformOutput',0);
   return; 
end


%% Classifying beats

beats{1} = resample(clusters(:,1),RES,NB_BINS);
beats{2} = resample(clusters(:,2),RES,NB_BINS);
beatsuct = beats;
dem = (length(beats{1})-2*T_LENGTH-1)/2;
beats = cellfun(@(x) x(dem+1:end-dem),beats,'UniformOutput',0);
    
for i = 1:length(qrs)
    cor = corrcoef(ecg(qrs(i)-T_LENGTH:qrs(i)+T_LENGTH),beats{1});
    corel(1,i) = cor(1,2);
    cor = corrcoef(ecg(qrs(i)-T_LENGTH:qrs(i)+T_LENGTH),beats{2});
    corel(2,i) = cor(1,2);
end
[~,type]=max(corel);

if sum(type==2)>sum(type==1)
    type(type==2) = 3;
    type(type==1) = 2;
    type(type==3) = 2;
end

if all(type==2)
    type = ones(size(type));
end
type = type';
    
end