function ecgc_res = ecgc(dat,maxlag,ensemble,silent,varargin)
% ECGC Ensemble Conditional Granger Causality
% The ecgc function calculates the conditional Granger causality. To allow 
% for an analysis of transient events an ensemble of time series can be 
% used to significantly shorten the required window size for a 
% time-resolved analysis.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
% ECGC_RES = ECGC(DAT,MAXLAG,ENSEMBLE,SILENT,VARARGIN) gives back a
% structure containing the Granger causality indices GC, their differences,
% the direction of interaction DOI, the probability values PRB, and the
% significant Granger causality indices SIG_GC using the method of
% conditional Granger causality on a data set DAT. The model order is
% determined by MAXLAG. The parameters ENSEMBLE and SILENT determine, 
% whether an ensemble approach is used and whether there will be a 
% graphical representation of the results, respectively. Optional
% parameters in VARARGIN comprise the following possibilites. PVAL 
% describes the significance level for the  Granger-Sargent significance 
% test. QUAD and MIX add pure and mixed quadratic terms to the model,
% respectively. ZVEC contains a list of variables for which lag-zero values
% are regarded in the model.
%
%   'dat'       - (matrix) [variables, observations, realizations]
%   'maxlag'    - (integer) maximum lag to consider in AR-model estimation
%   'ensemble'  - (boolean) '1' ensemble method, '0' classical
%   'silent'    - (boolean) '1' no plot, '0' plot
% varargin:
%   'pval'      - (double) significance threshold
%   'quad'      - (boolean) '1' use quadratic terms, '0' only linear
%   'mix'       - (boolean) '1' allow mixed terms, '0' only pure terms
%   'zvec'      - (vector) containing the variable indices, for which
%                  variables lag '0' should be considered
%
% output:
%
%   'ecgc_res'  - (struct)
%                   - .gc:  Granger causality values (column causes row)
%                   - .doi: direction of interaction (differences of gc)
%                   - .prb: probability values (if prb < significance level
%                           --> significant)
%                   - .sig_gc: only significant values of gc (column causes
%                                row)
% References:
%
% Seth AK (2010) A MATLAB toolbox for Granger causal connectivity analysis.
% J Neurosci Methods 186: 262–273.
%
% Mueller A, Riedl M, Penzel T, Bonnemeier H, Kurths J, et al. (2013) 
% Coupling analysis of transient cardiovascular dynamics. Biomed Tech 
% (Berl) 58: 131–139.
%
% Mueller A, Riedl M, Penzel T, Kurths J, Wessel N (2014) 
% Kardiorespiratorische Koordination und Ensemble-Kopplungsspuren zur 
% ereignisbasierten Charakterisierung kardiovaskulärer Interaktionen 
% während des Schlafes. Somnologie - Schlafforsch und Schlafmedizin 18: 
% 243–251.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% example: bivariate AR-model
% 
% len = 500;
% N = 100;
% M = 200;
% dat = zeros(2,N,M);
% noise = randn(2,len,M);
% for mm = 1:M
%     x = zeros(2,len);
%     x(:,1:2) = randn(2,2);
%     for ii = 3:len
%         x(1,ii) = 0.3*x(1,ii-1) + 0.7*x(2,ii-1) + 0.01*noise(1,ii,mm);
%         x(2,ii) = 0.3*x(2,ii-1) - 0.7*x(1,ii-2) + 0.01*noise(2,ii,mm);
%     end
%     dat(:,:,mm) = x(:,end-N+1:end);
% end
% 
% res = ecgc(dat,10,1,0,0.05);
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

pval = 0.05;
quad = 0;
mix = 0;
zvec = [];
zvec_l = 0;

if nargin>4
    if ~isempty(varargin{1})
        pval = varargin{1};
    end
    if nargin > 5
        if ~isempty(varargin{2})
            quad = varargin{2};
        end
        if nargin > 6
            if ~isempty(varargin{3})
                mix = varargin{3};
            end
            if nargin>7
                if ~isempty(varargin{4})
                    zvec = varargin{4};
                    zvec_l = length(zvec);
                end
            end
        end
    end
end

if numel(size(dat)) == 2
    [vars,N] = size(dat);
    M = 1;
    % remove sample means if present (no constant terms in this regression)
    ms = mean(dat,2);
    if(abs(sum(ms)) > 0.0001)
        msall = repmat(ms,1,N);
        dat = dat-msall;
    end
elseif numel(size(dat)) == 3
    [vars,N,M] = size(dat);
    
    % remove ensemble means if present (no constant terms in this regression)
    me = mean(dat,3);
    if(abs(sum(me)) > 0.0001)
        meall = repmat(me,[1 1 M]);
        dat = dat-meall;
    end
else
    error('incorrect size of data matrix (vars,N,M)')
end

% build lag-index-matrix

% without zvec
% indmat = ones(N-maxlag,1) * reshape(0:maxlag-1,1,[]) + ...
% reshape(1:N-maxlag,[],1) * ones(1,maxlag);
% with zvec
indmat = ones(N-maxlag,1) * reshape(0:maxlag,1,[]) + ...
    reshape(1:N-maxlag,[],1) * ones(1,maxlag+1);

if ~ensemble || M==1
    % prediction errors (unrestricted)
    uregerr = zeros(N-maxlag,vars);
    % sum of squared errors (unrestricted)
    uerrsum = zeros(vars,1);
    % variance of error (unrestricted)
    uerrvar = zeros(vars,1);
    % prediction errors (restricted)
    rregerr = zeros(N-maxlag,vars,vars);
    % sum of squared errors (restricted)
    rerrsum = zeros(vars,vars);
    % variance of error (restricted)
    rerrvar = zeros(vars,vars);
    % covariance of errors (restricted)
    rerrcov = cell(1,vars);
    
    for ii=1:vars
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % unrestricted regression
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % build unrestricted lag matrix
        if ismember(ii,zvec)
            uregmat = zeros(N-maxlag,vars*maxlag+zvec_l-1);
        else
            uregmat = zeros(N-maxlag,vars*maxlag+zvec_l);
        end
        ind_counter = 1;
        var_start_ind = zeros(1,vars);
        for jj=1:vars
            if ~ismember(jj,zvec) || jj==ii
                uregmat(:,ind_counter:ind_counter+maxlag-1) = ...
                    reshape(dat(jj,indmat(:,1:maxlag)),N-maxlag,maxlag);
                var_start_ind(jj) = ind_counter;
                ind_counter = ind_counter + maxlag;
            else
                uregmat(:,ind_counter:ind_counter+maxlag) = ...
                    reshape(dat(jj,indmat),N-maxlag,maxlag+1);
                var_start_ind(jj) = ind_counter;
                ind_counter = ind_counter + maxlag + 1;
            end
        end
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % quadratic and mixed terms
        if quad
            cols = size(uregmat,2);
            uregmat_quad = zeros(N-maxlag,cols);
            for jj=1:cols
                uregmat_quad(:,jj) = uregmat(:,jj) .* uregmat(:,jj);
            end
        else
            uregmat_quad = [];
        end
        if mix
            uregmat_mix = zeros(N-maxlag,vars*(vars-1)/2);
            counter = 1;
            for jj=1:(vars-1)
                for kk=(jj+1):vars
                    uregmat_mix(:,counter) = uregmat(:,var_start_ind(jj)) .* uregmat(:,var_start_ind(kk));
                    counter = counter + 1;
                end
            end
        else
            uregmat_mix = [];
        end
        uregmat = [uregmat uregmat_quad uregmat_mix];
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % unrestricted regression parameters (result)
        dat_dep = dat(ii,maxlag+1:N)';
        uregpar = uregmat\dat_dep;
        uregerr(:,ii) = dat_dep - uregmat*uregpar;
        uerrsum(ii) = sum(uregerr(:,ii).*uregerr(:,ii));
        uerrvar(ii) = var(uregerr(:,ii));
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % restricted regression
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % possible causal influences
        cause_ind = setdiff(1:vars,ii);
        for jj=1:vars-1
            % vars to include in restricted regression
            reg_ind = setdiff(1:vars,cause_ind(jj));
            % build restricted lag matrix
            regmatcol = (vars-1)*maxlag+zvec_l;
            if ~isempty(setdiff(zvec,reg_ind))
                regmatcol = regmatcol - 1;
            end
            % build restricted lag matrix
            rregmat = zeros(N-maxlag,regmatcol);
            ind_counter = 1;
            var_start_ind = zeros(1,vars-1);
            for kk=1:vars-1
                if ~ismember(reg_ind(kk),zvec) || reg_ind(kk)==ii
                    rregmat(:,ind_counter:ind_counter+maxlag-1) = ...
                        reshape(dat(reg_ind(kk),indmat(:,1:maxlag)),...
                        N-maxlag,maxlag);
                    var_start_ind(kk) = ind_counter;
                    ind_counter = ind_counter + maxlag;
                else
                    rregmat(:,ind_counter:ind_counter+maxlag) = ...
                        reshape(dat(reg_ind(kk),indmat),N-maxlag,maxlag+1);
                    var_start_ind(kk) = ind_counter;
                    ind_counter = ind_counter + maxlag + 1;
                end
            end
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            % quadratic terms
            if quad
                cols = size(rregmat,2);
                rregmat_quad = zeros(N-maxlag,cols);
                for kk=1:cols
                    rregmat_quad(:,kk) = rregmat(:,kk) .* rregmat(:,kk);
                end
            else
                rregmat_quad = [];
            end
            if mix
                rregmat_mix = zeros(N-maxlag,(vars-1)*(vars-2)/2);
                counter = 1;
                for kk=1:(vars-2)
                    for ll=(kk+1):(vars-1)
                        rregmat_mix(:,counter) = rregmat(:,var_start_ind(kk)) .* rregmat(:,var_start_ind(ll));
                        counter = counter + 1;
                    end
                end
            else
                rregmat_mix = [];
            end
            rregmat = [rregmat rregmat_quad rregmat_mix];
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            % restricted regression parameters (result)
            rregpar = rregmat\dat_dep;
            rregerr(:,ii,cause_ind(jj)) = dat_dep - rregmat*rregpar;
            rerrsum(ii,cause_ind(jj)) = sum(rregerr(:,ii,cause_ind(jj)).*...
                rregerr(:,ii,cause_ind(jj)));
            rerrvar(ii,cause_ind(jj)) = var(rregerr(:,ii,cause_ind(jj)));
        end
        rerrcov{ii} = cov(squeeze(rregerr(:,ii,:)));
    end
    uerrcov = cov(uregerr);
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % calculate Granger values
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    gc = ones(vars).*NaN;
    doi = ones(vars).*NaN;
    % do Granger f-tests
    prb = ones(vars).*NaN;
    sig_gc = ones(vars).*NaN;
    ftest = zeros(vars);
    % degrees of freedom = #predictable values - #used lagged values of all variables
    dof = N-(vars+1)*maxlag-zvec_l;
    for ii=1:vars-1
        for jj=ii+1:vars
            % causality jj->ii
            ftest(ii,jj) = dof*(rerrsum(ii,jj)-uerrsum(ii))/(maxlag*uerrsum(ii));
            prb(ii,jj) = 1 - cdf('f',ftest(ii,jj),maxlag,dof);
            % causality ii->jj
            ftest(jj,ii) = dof*(rerrsum(jj,ii)-uerrsum(jj))/(maxlag*uerrsum(jj));
            prb(jj,ii) = 1 - cdf('f',ftest(jj,ii),maxlag,dof);
            gc(ii,jj) = log(rerrvar(ii,jj)/uerrvar(ii));
            gc(jj,ii) = log(rerrvar(jj,ii)/uerrvar(jj));
            doi(ii,jj) = gc(ii,jj) - gc(jj,ii);
            doi(jj,ii) = gc(jj,ii) - gc(ii,jj);
        end
    end
    sig_gc((1-prb)>=(1-pval)) = 1;
    sig_gc((1-prb)<(1-pval)) = 0;
    sig_gc = sig_gc .* gc;
else
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % ensemble regression
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % unrestricted regression
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    gc = cell(1,N-maxlag);
    doi = cell(1,N-maxlag);
    prb = cell(1,N-maxlag);
    sig_gc = cell(1,N-maxlag);
    for ii=1:N-maxlag
        % prediction errors (unrestricted)
        uregerr = zeros(M,vars);
        % sum of squared errors (unrestricted)
        uerrsum = zeros(vars,1);
        % variance of error (unrestricted)
        uerrvar = zeros(vars,1);
        % prediction errors (restricted)
        rregerr = zeros(M,vars,vars);
        % sum of squared errors (restricted)
        rerrsum = zeros(vars,vars);
        % variance of error (restricted)
        rerrvar = zeros(vars,vars);
        for jj=1:vars
            if ismember(jj,zvec)
                uregmat = zeros(M,vars*maxlag+zvec_l-1);
            else
                uregmat = zeros(M,vars*maxlag+zvec_l);
            end
            ind_counter = 1;
            var_start_ind = zeros(1,vars);
            for kk=1:vars
                if ~ismember(kk,zvec) || kk==jj
                    uregmat(:,ind_counter:ind_counter+maxlag-1) = ...
                        squeeze(dat(kk,ii:ii+maxlag-1,:))';
                    var_start_ind(kk) = ind_counter;
                    ind_counter = ind_counter + maxlag;
                else
                    uregmat(:,ind_counter:ind_counter+maxlag) = ...
                        squeeze(dat(kk,ii:ii+maxlag,:))';
                    var_start_ind(kk) = ind_counter;
                    ind_counter = ind_counter + maxlag + 1;
                end
            end
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            % quadratic terms
            if quad
                cols = size(uregmat,2);
                uregmat_quad = zeros(M,cols);
                for kk=1:cols
                    uregmat_quad(:,kk) = uregmat(:,kk) .* uregmat(:,kk);
                end
            else
                uregmat_quad = [];
            end
            if mix
                uregmat_mix = zeros(M,vars*(vars-1)/2);
                counter = 1;
                for kk=1:(vars-1)
                    for ll=(kk+1):vars
                        uregmat_mix(:,counter) = uregmat(:,var_start_ind(kk)) .* uregmat(:,var_start_ind(ll));
                        counter = counter + 1;
                    end
                end
            else
                uregmat_mix =[];
            end
            uregmat = [uregmat uregmat_quad uregmat_mix];
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            % unrestricted regression parameters (result)
            dat_dep = squeeze(dat(jj,maxlag+ii,:));
            uregpar = uregmat\dat_dep;
            uregerr(:,jj) = dat_dep - uregmat*uregpar;
            uerrsum(jj) = sum(uregerr(:,jj).*uregerr(:,jj));
            uerrvar(jj) = var(uregerr(:,jj));
            
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            % restricted regression
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
            
            % possible causal influences
            cause_ind = setdiff(1:vars,jj);
            for kk=1:vars-1
                % vars to include in restricted regression
                reg_ind = setdiff(1:vars,cause_ind(kk));
                % build restricted lag matrix
                regmatcol = (vars-1)*maxlag+zvec_l;
                if ~isempty(setdiff(zvec,reg_ind))
                    regmatcol = regmatcol - 1;
                end
                % build restricted lag matrix
                rregmat = zeros(M,regmatcol);
                ind_counter = 1;
                var_start_ind = zeros(1,vars-1);
                for ll=1:vars-1
                    if ~ismember(reg_ind(ll),zvec) || reg_ind(ll)==jj
                        rregmat(:,ind_counter:ind_counter+maxlag-1) = ...
                            squeeze(dat(reg_ind(ll),ii:ii+maxlag-1,:))';
                        var_start_ind(ll) = ind_counter;
                        ind_counter = ind_counter + maxlag;
                    else
                        rregmat(:,ind_counter:ind_counter+maxlag) = ...
                            squeeze(dat(reg_ind(ll),ii:ii+maxlag,:))';
                        var_start_ind(ll) = ind_counter;
                        ind_counter = ind_counter + maxlag + 1;
                    end
                end
                %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
                % quadratic terms
                if quad
                    cols = size(rregmat,2);
                    rregmat_quad = zeros(M,cols);
                    for ll=1:cols
                        rregmat_quad(:,ll) = rregmat(:,ll) .* rregmat(:,ll);
                    end
                else
                    rregmat_quad = [];
                end
                if mix
                    rregmat_mix = zeros(M,(vars-1)*(vars-2)/2);
                    counter = 1;
                    for ll=1:(vars-2)
                        for mm=(ll+1):(vars-1)
                            rregmat_mix(:,counter) = rregmat(:,var_start_ind(ll)) .* rregmat(:,var_start_ind(mm));
                            counter = counter + 1;
                        end
                    end
                else
                    rregmat_mix =[];
                end
                rregmat = [rregmat rregmat_quad rregmat_mix];
                %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
                % restricted regression parameters (result)
                rregpar = rregmat\dat_dep;
                rregerr(:,jj,cause_ind(kk)) = dat_dep - rregmat*rregpar;
                rerrsum(jj,cause_ind(kk)) = sum(rregerr(:,jj,cause_ind(kk)).*...
                    rregerr(:,jj,cause_ind(kk)));
                rerrvar(jj,cause_ind(kk)) = var(rregerr(:,jj,cause_ind(kk)));
            end
        end
        
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % calculate Granger values
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        gc{ii} = ones(vars).*NaN;
        doi{ii} = ones(vars).*NaN;
        % do Granger f-tests
        prb{ii} = ones(vars).*NaN;
        sig_gc{ii} = ones(vars).*NaN;
        ftest = zeros(vars);
        % degrees of freedom = #predictable values - ...
        % #used lagged values of all variables
        dof = N-(vars+1)*maxlag-zvec_l;
        for jj=1:vars-1
            for kk=jj+1:vars
                % causality kk->jj
                ftest(jj,kk) = dof*(rerrsum(jj,kk)-uerrsum(jj))/(maxlag*uerrsum(jj));
                prb{ii}(jj,kk) = 1 - cdf('f',ftest(jj,kk),maxlag,dof);
                % causality jj->kk
                ftest(kk,jj) = dof*(rerrsum(kk,jj)-uerrsum(kk))/(maxlag*uerrsum(kk));
                prb{ii}(kk,jj) = 1 - cdf('f',ftest(kk,jj),maxlag,dof);
                
                gc{ii}(jj,kk) = log(rerrvar(jj,kk)/uerrvar(jj));
                gc{ii}(kk,jj) = log(rerrvar(kk,jj)/uerrvar(kk));
                doi{ii}(jj,kk) = gc{ii}(jj,kk) - gc{ii}(kk,jj);
                doi{ii}(kk,jj) = gc{ii}(kk,jj) - gc{ii}(jj,kk);
            end
        end
        sig_gc{ii}((1-prb{ii})>=(1-pval)) = 1;
        sig_gc{ii}((1-prb{ii})<(1-pval)) = 0;
        sig_gc{ii} = sig_gc{ii} .* gc{ii};
        sig_gc{ii}(sig_gc{ii}<0.05) = 0;
    end
end

ecgc_res = struct('gc',gc,'doi',doi,'prb',prb,'sig_gc',sig_gc);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% graphical output
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if ~silent
    if ensemble
        rescell = struct2cell(ecgc_res);
        resmat = cell2mat(rescell);
        [~,~,t_range] = size(resmat);
        
        % significant gc values
        n = vars*(vars-1)/2;
        ind = 0;
        for ii=1:(vars-1)
            for jj=(ii+1):vars
                ind = ind + 1;
                figure;
                plot(1:t_range,squeeze(resmat((ii+3*vars),jj,:)),'b',1:t_range,squeeze(resmat((jj+3*vars),ii,:)),'r')
                title(['sig gc: red: ',num2str(ii),'->',num2str(jj)])
            end
        end
    else
        figure;
        imagesc(gc,[0,1]);
    end
end

end % of function
