%*******************************************************************************
%     DBA with modified initial sequence selection, for EEG data analysis.
%     Copyright (C) 2016  Levi Davis
% 
%     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/>.
 %*****************************************************************************/ 
 
 % Contact: levi.davis@colorado.edu
 
 % If using this code or a derivation of it, please cite: 
 % Davis, L. (2016). Application of DTW Barycenter Averaging to Finding EEG 
 %   Consensus Sequences. Unpublished manuscript, Department of Psychology
 %   and Neuroscience, University of Colorado at Boulder.
 
 
 % Works referenced: 
%
%  Petitjean, F., Ketterlin, A., & Gançarski, P. (2011). A global averaging 
%    method for dynamic time warping, with applications to clustering. 
%    Pattern Recognition, 44(3), 678-693.
%  Kotas, M., Leski, J. M., & Moro?, T. (2015). Dynamic Time Warping Based 
%    on Modified Alignment Costs for Evoked Potentials Averaging. Advances 
%    in Intelligent Systems and Computing Man–Machine Interactions 4, 
%    305-314.
%  Sakoe, H., & Chiba, S. (1990). Dynamic Programming Algorithm 
%    Optimization for Spoken Word Recognition. Readings in Speech 
%    Recognition, 159-165.
%


function avg = dba_for_eeg(sequences, start_method, iterations, assoc_smooth, score_smooth)
    % This function returns the consensus of the input sequences, using
    %   one of several variations of the DBA algorithm described in Davis
    %   (2016). 
    
    % sequences: data to be averaged, assumed to be an array with first
    %   dimension corresponding to time and second dimension corresponding
    %   to sequence number.
    % start_method: an integer corresponding to the desired method of
    %   initial reference sequence selection. (default 0)
    %       0: DBA_0: arbitrary (first sequence used)
    %       1: DBA_ERP: ERP (average)
    %       2: DBA_DTW: input sequence most similar to ERP according to DTW
    %       3: DBA_LEN: input sequence most similar to ERP according to 
    %           length of DTW warping path
    %       4: DBA_DEV: input sequence most similar to ERP according to 
    %           integral of distance from diagonal in DTW warping path
    % iterations: number of DBA iterations to run (default 10)
    % assoc_smooth: value of v (Kotas et al., 2015) used when calculating
    %   vertex associations during DBA iteration
    % score_smooth: value of v (Kotas et al., 2015) used when selecting an
    %   initial reference sequence (start_method >= 2)


    % fill in optional arguments with default values
    if nargin < 2
        start_method = 0;
    end
    if nargin < 3
        iterations = 10;
    end
    if nargin < 4
        assoc_smooth = 3;
    end
    if nargin < 5
        score_smooth = 1;
    end

    % select initial sequence
    %display('SELECTING INITIAL SEQUENCE');
    if start_method == 0
        avg = sequences(:, 1);
    elseif start_method == 1
        avg = calc_erp(sequences);
    else 
        avg = select_reference(start_method, sequences, score_smooth);
    end
        
    
    % iteratively improve the sequence
    for i = 1:iterations
        %fprintf('iteration %i\n', i);
        avg = dba_iterate(avg, sequences, assoc_smooth);
    end
    
end

function result = dba_iterate(reference, sequences, smooth)
    
    % use modified DTW to associate reference verts w/ sequence verts
    assocTab = cell(1, size(reference, 1));
    for s = 1:size(sequences, 2)
        %fprintf('aligning with sequence %i\n', s);
        assocTab = dtw_assoc(reference, sequences(:, s), assocTab, smooth);
    end
    % generate a better reference by averaging all associated verts
    %display('calculating barycenters');
    result = zeros(size(reference, 1), 1);
    for i = 1:size(reference, 1)
        for j = 1:size(assocTab{i}, 2)
            result(i, 1) = result(i, 1) + assocTab{i}(j);
        end
        result(i, 1) = result(i, 1)/size(assocTab{i}, 2);
    end
    
end

function assoc = dtw_assoc(S, T, assocTab, smooth)

    % set up
    m = size(S, 1);
    n = size(T, 1);
    cost = zeros(m, n);
    path = zeros(m, n);
    leng = zeros(m, n);
    % top left
    cost(1, 1) = penalty(S, T, 1, 1, smooth);
    path(1, 1) = 0;
    leng(1, 1) = 1;
    % left
    for i = 2:m
        cost(i, 1) = cost(i - 1, 1) + penalty(S, T, i, 1, smooth);
        path(i, 1) = 1;
        leng(i, 1) = i;
    end
    % top
    for j = 2:n
        cost(1, j) = cost(1, j - 1) + penalty(S, T, 1, j, smooth);
        path(1, j) = 2;
        leng(1, j) = j;
    end
    % rest of matrix
    for i = 2:m
        for j = 2:n
            % diagonal
            if (cost(i - 1, j - 1) <= cost(i, j - 1) && cost(i - 1, j - 1) <= cost(i - 1, j))
                cost(i, j) = cost(i - 1, j - 1) + penalty(S, T, i, j, smooth);
                path(i, j) = 0;
                leng(i, j) = leng(i - 1, j - 1) + 1;
            % top
            elseif (cost(i - 1, j) <= cost(i, j - 1))
                cost(i, j) = cost(i - 1, j) + penalty(S, T, i, j, smooth);
                path(i, j) = 1;
                leng(i, j) = leng(i - 1, j) + 1;
            % left
            else
                cost(i, j) = cost(i, j - 1) + penalty(S, T, i, j, smooth);
                path(i, j) = 2;
                leng(i, j) = leng(i, j - 1) + 1;
            end
        end
    end
    
    % trace the optimal path in the matrix, recording the assoc pairings
    i = m;
    j = n;
    assoc = assocTab;
    while (i > 1 || j > 1)
        %make the association
        assoc{i}(end + 1) = T(j, 1);
        %step to the next location
        if (path(i, j) == 0)
            i = i - 1;
            j = j - 1;
        elseif (path(i, j) == 1)
            i = i - 1;
        elseif (path(i, j) == 2)
            j = j - 1;
        end
    end
    assoc{1}(end + 1) = T(1, 1);
    
end

function c = penalty(S, T, a, b, d)
% penalty function, modified from Kotas et al. (2016, pg308)
% S, T = sequences being compared
% a, b = indices being compared in each sequence
% d = the distance on either side a/b that should be considered (reduces effects of noise)
    
    c = 0;
    n = 0;
    
    for i = -d:d
        j = a + i;
        k = b + i;
        
        if (j > 0 && k > 0 && j <= size(S, 1) && k <= size(T, 1))
            n = n + 1;
            c = c + (S(j, 1) - T(k, 1))^2;
        end
        
    end
    
    c = sqrt(c)/n;
    
end

function erp = calc_erp(data)

    erp = zeros(size(data, 1), 1);

    for i = 1:size(data, 2)
        for j = 1:size(data, 1)
            erp(j, 1) = erp(j, 1) + data(j, i);
        end
    end

    for j = 1:size(data, 1)
        erp(j, 1) = erp(j, 1)/size(data, 2);
    end

end

function start = select_reference(start_method, sequences, smooth)

    %fprintf('calculating erp\n');
    avg = calc_erp(sequences);
    
    best_score = Inf(1);
    for i = 1:size(sequences, 2)
        %fprintf('aligning with sequence %i\n', i);
        
        if start_method == 2
            score = dtw(avg, sequences(:, i), smooth);
        elseif start_method == 3
            score = dtw_length(avg, sequences(:, i), smooth);
        elseif start_method == 4
            score = dtw_perturb(avg, sequences(:, i), smooth);
        end
        
        if (score < best_score)
            best_score = score;
            start = sequences(:, i);
        end
    end

end

function score = dtw(S, T, smooth)

    % set up
    m = size(S, 1);
    n = size(T, 1);
    cost = zeros(m, n);
    % top left
    cost(1, 1) = penalty(S, T, 1, 1, smooth);
    % left
    for i = 2:m
        cost(i, 1) = cost(i - 1, 1) + penalty(S, T, i, 1, smooth);
    end
    % top
    for j = 2:n
        cost(1, j) = cost(1, j - 1) + penalty(S, T, 1, j, smooth);
    end
    % rest of matrix
    for i = 2:m
        for j = 2:n
            % diagonal
            if (cost(i - 1, j - 1) <= cost(i, j - 1) && cost(i - 1, j - 1) <= cost(i - 1, j))
                cost(i, j) = cost(i - 1, j - 1) + penalty(S, T, i, j, smooth);
            % top
            elseif (cost(i - 1, j) <= cost(i, j - 1))
                cost(i, j) = cost(i - 1, j) + penalty(S, T, i, j, smooth);
            % left
            else
                cost(i, j) = cost(i, j - 1) + penalty(S, T, i, j, smooth);
            end
        end
    end
    score = cost(m, n);
    
end

function score = dtw_perturb(S, T, smooth)

    % set up
    m = size(S, 1);
    n = size(T, 1);
    cost = zeros(m, n);
    pert = zeros(m, n);
    % top left
    cost(1, 1) = penalty(S, T, 1, 1, smooth);
    pert(1, 1) = 0;
    % left
    for i = 2:m
        cost(i, 1) = cost(i - 1, 1) + penalty(S, T, i, 1, smooth);
        pert(i, 1) = pert(i - 1, 1) + abs(i - 1);
    end
    % top
    for j = 2:n
        cost(1, j) = cost(1, j - 1) + penalty(S, T, 1, j, smooth);
        pert(1, j) = pert(1, j - 1) + abs(1 - 1);
    end
    % rest of matrix
    for i = 2:m
        for j = 2:n
            % diagonal
            if (cost(i - 1, j - 1) <= cost(i, j - 1) && cost(i - 1, j - 1) <= cost(i - 1, j))
                cost(i, j) = cost(i - 1, j - 1) + penalty(S, T, i, j, smooth);
                pert(i, j) = pert(i - 1, j - 1) + abs(i - j);
            % top
            elseif (cost(i - 1, j) <= cost(i, j - 1))
                cost(i, j) = cost(i - 1, j) + penalty(S, T, i, j, smooth);
                pert(i, j) = pert(i - 1, j) + abs(i - j);
            % left
            else
                cost(i, j) = cost(i, j - 1) + penalty(S, T, i, j, smooth);
                pert(i, j) = pert(i, j - 1) + abs(i - j);
            end
        end
    end
    score = pert(m, n);
    
end

function score = dtw_length(S, T, smooth)

    % set up
    m = size(S, 1);
    n = size(T, 1);
    cost = zeros(m, n);
    leng = zeros(m, n);
    % top left
    cost(1, 1) = penalty(S, T, 1, 1, smooth);
    leng(1, 1) = 1;
    % left
    for i = 2:m
        cost(i, 1) = cost(i - 1, 1) + penalty(S, T, i, 1, smooth);
        leng(i, 1) = i;
    end
    % top
    for j = 2:n
        cost(1, j) = cost(1, j - 1) + penalty(S, T, 1, j, smooth);
        leng(1, j) = j;
    end
    % rest of matrix
    for i = 2:m
        for j = 2:n
            % diagonal
            if (cost(i - 1, j - 1) <= cost(i, j - 1) && cost(i - 1, j - 1) <= cost(i - 1, j))
                cost(i, j) = cost(i - 1, j - 1) + penalty(S, T, i, j, smooth);
                leng(i, j) = leng(i - 1, j - 1) + 1;
            % top
            elseif (cost(i - 1, j) <= cost(i, j - 1))
                cost(i, j) = cost(i - 1, j) + penalty(S, T, i, j, smooth);
                leng(i, j) = leng(i - 1, j) + 1;
            % left
            else
                cost(i, j) = cost(i, j - 1) + penalty(S, T, i, j, smooth);
                leng(i, j) = leng(i, j - 1) + 1;
            end
        end
    end
    score = leng(m, n);
    
end