function displayWaveformLabels(ECG,labels)
% DISPLAYWAVEFORMLABELS displays the ECG signal data color coded based on
% the input label values.

% Copyright 2020 The MathWorks, Inc.

ECG = ECG(:)';
labels = labels(:)';

tecg = 1:length(ECG);

%% Find P regions
pidx = find(labels == 'P');
pidxsep = find(diff(pidx)~=1);
if ~isempty(pidxsep)
    pidxw = cell(1,length(pidxsep)+1);
    for ii = 1:length(pidxsep)+1
        if ii == 1
            % Catch first P region
            pidxw{ii} = pidx(1:pidxsep(ii));
        elseif ii <= length(pidxsep)
            pidxw{ii} = pidx(pidxsep(ii-1)+1:pidxsep(ii));
        else
            % Catch last P region
            pidxw{ii} = pidx(pidxsep(ii-1)+1:end);
        end
    end
else
    if ~isempty(pidx)
        % Special case of only a single region
        pidxw{1} = pidx;
    else
        % No P wave present
        pidxw = {};
    end
end

%% Find QRS Regions
qrsidx =  find(labels == 'QRS');
qrsidxsep = find(diff(qrsidx)~=1);

if ~isempty(qrsidxsep)
    qrsidxw = cell(1,length(qrsidxsep)+1);
    for ii = 1:length(qrsidxsep)+1
        if ii == 1
            % Catch first QRS region
            qrsidxw{ii} = qrsidx(1:qrsidxsep(ii));
        elseif ii <= length(qrsidxsep)
            qrsidxw{ii} = qrsidx(qrsidxsep(ii-1)+1:qrsidxsep(ii));
        else
            % Catch last QRS region
            qrsidxw{ii} = qrsidx(qrsidxsep(ii-1)+1:end);
        end
    end
else
    if ~isempty(qrsidx)
        % Special case of only a single region
        qrsidxw{1} = qrsidx;
    else
        % No QRS Present
        qrsidxw = {};
    end
end

%% Find T Regions
tidx =  find(labels == 'T');
tidxsep = find(diff(tidx)~=1);

if ~isempty(tidxsep)
    tidxw = cell(1,length(tidxsep)+1);
    for ii = 1:length(tidxsep)+1
        if ii == 1
            % Catch first T region
            tidxw{ii} = tidx(1:tidxsep(ii));
        elseif ii <= length(tidxsep)
            tidxw{ii} = tidx(tidxsep(ii-1)+1:tidxsep(ii));
        else
            % Catch last T region
            tidxw{ii} = tidx(tidxsep(ii-1)+1:end);
        end
    end
else
    if ~isempty(tidx)
        % Special case of only a single region
        tidxw{1} = tidx;
    else
        % No QRS Present
        tidxw = {};
    end
end

%% Find N/A regions
naidx =  find(labels == 'n/a');
naidxsep = find(diff(naidx)~=1);

if ~isempty(naidxsep)
    naidxw = cell(1,length(naidxsep)+1);
    for ii = 1:length(naidxsep)+1
        if ii == 1
            % Catch first N/A region
            naidxw{ii} = naidx(1:naidxsep(ii));
        elseif ii <= length(naidxsep)
            naidxw{ii} = naidx(naidxsep(ii-1)+1:naidxsep(ii));
        else
            % Catch last N/A region
            naidxw{ii} = naidx(naidxsep(ii-1)+1:end);
        end
    end
else
    if ~isempty(naidx)
        % Special case of only a single region
        naidxw{1} = naidx;
    else
        % No N/A Present
        naidxw = {};
    end
end

%% Plot labelled waveform

hAx = gca;
hold on

% Plot P wave regions
for ii = 1:length(pidxw)
    iidx = pidxw{ii};
    if iidx(1) ~= 1
        iidx = [iidx(1)-1 iidx]; %#ok<AGROW>
    end
    
    if iidx ~= length(tecg)
        iidx(end+1) = iidx(end)+1; %#ok<AGROW>
    end
    plot(tecg(iidx),ECG(iidx),'Color',[0 0.447 0.741],'LineWidth',2)
end

% Plot QRS regions
for ii = 1:length(qrsidxw)
    iidx = qrsidxw{ii};
    if iidx(1) ~= 1
        iidx = [iidx(1)-1 iidx]; %#ok<AGROW>
    end
    
    if iidx ~= length(tecg)
        iidx(end+1) = iidx(end)+1; %#ok<AGROW>
    end
    plot(tecg(iidx),ECG(iidx),'Color',[0.85 0.325 0.098],'LineWidth',2)
end

% Plot T wave regions
for ii = 1:length(tidxw)
    iidx = tidxw{ii};
    
    if iidx(1) ~= 1
        iidx = [iidx(1)-1 iidx]; %#ok<AGROW>
    end
    
    if iidx ~= length(tecg)
        iidx(end+1) = iidx(end)+1; %#ok<AGROW>
    end
    plot(tecg(iidx),ECG(iidx),'Color',[0.929 0.694 0.125],'LineWidth',2)
end

% Plot N/A regions
for ii = 1:length(naidxw)
    iidx = naidxw{ii};
    if iidx(1) ~= 1
        iidx = [iidx(1)-1 iidx]; %#ok<AGROW>
    end
    
    if iidx ~= length(tecg)
        iidx(end+1) = iidx(end)+1; %#ok<AGROW>
    end
    plot(tecg(iidx),ECG(iidx),'k','LineWidth',2)
end

% Setup plot options
grid on
box on
hold off
axis tight

cmap = [0 0 0;
    0.929 0.694 0.125;
    0.85 0.325 0.098;
    0 0.447 0.741];
classNames = {'N/A','T Wave','QRS','P Wave'};
colormap(hAx,cmap)

% Add colorbar to current figure
c = colorbar(hAx);

% Use class names for tick marks
c.TickLabels = classNames;
numClasses = size(cmap,1);

% Center tick labels
c.Ticks = 1/(numClasses*2):1/numClasses:1;

% Remove tick mark
c.TickLength = 0;
end
