function saveFeatures(saveFile, options)
% estimate and save spectral features
%
% INPUTS
% saveFile: name of '.mat' file containing X and labels
%   variables; and to which the processed data will be saved.
% options: structure of options for running this function
% FIELDS
%   featureList: (Optional) Cell array of strings indicating which
%     features to calculate. The following strings are available options:
%     'directedSpectrum', 'pwDirectedSpectrum', 'pdc', 'dtf', 'psi',
%     'fft', 'granger', 'c_granger'.
%   parCores: (Optional) integer indicating number of cores to use for
%     parallel computing. If 0, all tasks executed in serial.
%   window: integer or vector
%     If an integer, gives length of Hamming subwindows used in
%     Welch's method for estimating a power spectrum.
%     If a vector, indicates the window that should be used in Welch's
%     method. Can also be left as an empty array '[]' to use the
%     matlab defaule window size.
%   overlap: integer
%     Indicate the overlap between sucessive windows when using
%     Welch's method to estimate a power spectrum. If left as empty
%     (i.e. []) then the matlab default is used.
%   mvgcFolder: needed if featureList includes 'granger'. String
%     indicating folder containing MVGC toolbox.
%
% LOADED VARIABLES (from saveFile)
% X: Preprocessed (filtered, averaged, checked for saturation) data. NxAxW
%       array. A is the # of areas. N=number of frequency points per
%       window. W=number of time windows.
% labels: Structure containing labeling infomation for data
%   FIELDS USED
%   area: cell array of labels for each area corresponding to
%       the second dimension of xFft
%   fs: sampling frequency of processed data (Hz)
%   windowLength: length of windows (s)
%
% SAVED VARIABLES
% granger: PxFxW array to store granger causality values. P
%     iterates over directed pairs of regions, F iterates over
%     frequencies, W iterates over windows.
% c_granger: PxFxW array to store conditional granger causality values. P
%     iterates over directed pairs of regions, F iterates over
%     frequencies, W iterates over windows.
% instant: PxFxW array to store instantaneous causality values. P
%     iterates over undirected pairs of regions, F iterates over
%     frequencies, W iterates over windows.
% directedSpectrum: PxFxW array to store 'full' model directed spectrum
%     features. P iterates over directed pairs of regions, F iterates over
%     frequencies, W iterates over windows.
% pwDirectedSpectrum: PxFxW array to store pairwise directed spectrum
%     features. P iterates over directed pairs of regions, F iterates over
%     frequencies, W iterates over windows.
% psi: CxCxFxW array to store phase-slope index values. C iterates over
%     regions, F iterates over frequencies, W iterates over windows.
% dtf: CxCxFxW array to store directed transfer function values. C iterates
%     over regions, F iterates over frequencies, W iterates over windows.
% pdc:CxCxFxW array to store partial directed coherence values. C iterates
%     over regions, F iterates over frequencies, W iterates over windows.
% fft: fourier transform of X
% labels: See above, with
%   ADDED FIELDS
%   f: integer frequency of processed data
%   gcFeatures: PxF array of string labels describing the
%       features represented in granger. P iterates over
%       directed pairs of regions, F iterates over frequencies.
%   cgcFeatures: PxF array of string labels describing the
%       features represented in c_granger. P iterates over
%       directed pairs of regions, F iterates over frequencies.
%   instFeatures: PxF array of string labels describing the
%       features represented in instArray. P iterates over
%       undirected pairs of regions, F iterates over frequencies.
%   dsFeatures: PxF array of string labels describing the
%       features represented in directedSpectrum and pwDirectedSpectrum
%       P iterates over directed pairs of regions, F iterates over
%       frequencies.
%   psiFeatures: CxCxF array of string labels describing the
%       features represented in psi. C iterates over
%       regions, F iterates over frequencies.
%   pdFeatures: CxCxF array of string labels describing the
%       features represented in pdc and dtf. C iterates over
%       regions, F iterates over frequencies.
%
% For Granger causality calculation, we use the Multivariate Granager
% Causality toolbox: https://users.sussex.ac.uk/~lionelb/MVGC/
% Please install that toolbox first if you wish to use this function to
% calculate Granger causality values for comparison. The install location
% should be passed to this function as options.mvgcFolder.

%% load data and prep for feature generation
load(saveFile,'X', 'labels')
fs = labels.fs;

if nargin < 2
    % fill with default parameters
    options=[];
end
options=fillDefaultOpts(options);

% evaluate at every integer frequency up to nyquist
f = 1:floor(fs/2);
nFreq = numel(f);
labels.f = f;

[N,C,W] = size(X);

fStrings = compose('%d', f)';
    
%% Get Granger causality features
if any(ismember('granger', options.featureList))
    % start multivariate granger causality toolbox
    mvgcStartupScript = [options.mvgcFolder '/startup.m'];
    run(mvgcStartupScript)

    % generate Granger causality values matrix in the form MxPxQ, where M
    % iterates over pairs of brain regions, P is frequency, and Q is time
    % window.
    [granger, gcFeatures, instant, instFeatures] = g_causality(double(X), labels.area, fs, ...
                                                               options);
    granger = single(granger);
    labels.gcFeatures = string(gcFeatures);
    instant = single(instant);
    labels.instFeatures = string(instFeatures);

    save(saveFile, 'granger', 'instant', '-append')
end

%% Get conditional Granger causality features
if any(ismember('c_granger', options.featureList))
    mvgcStartupScript = [options.mvgcFolder '/startup.m'];
    run(mvgcStartupScript)

    % generate Granger causality values matrix in the form MxPxQ, where M
    % iterates over pairs of brain regions, P is frequency, and Q is time
    % window.
    [c_granger, gcFeatures] = cond_g_causality(double(X), labels.area, fs, ...
                                                               options);
    c_granger = single(c_granger);
    labels.cgcFeatures = string(gcFeatures);

    save(saveFile, 'c_granger', '-append')
end


%% Get directed spectrum features
% Check if full or pairwise directedSpectrum features are to be calculated.
directionFeatures = ismember({'directedSpectrum', 'pwDirectedSpectrum'}, ...
    options.featureList);
if any(directionFeatures)
    X = double(X);
    
    % generate directed spectrum values matrix in the form MxPxQ, where M
    % iterates over pairs of brain regions, P is frequency, and Q is time
    % window.
    [directedSpectrum, dsFeatures, S] = directed_spectrum(X, labels.area, fs, f,...
        directionFeatures, options);
    
    labels.dsFeatures = string(dsFeatures);
    
    % Save calculated features
    if sum(directionFeatures) == 2
        pwDirectedSpectrum = single(directedSpectrum{2});
        directedSpectrum = single(directedSpectrum{1});
        save(saveFile, 'directedSpectrum', 'pwDirectedSpectrum', '-append')
    elseif directionFeatures(1)
        directedSpectrum = single(directedSpectrum);
        save(saveFile, 'directedSpectrum','-append')
    else
        pwDirectedSpectrum = single(directedSpectrum);
        save(saveFile, 'pwDirectedSpectrum','-append')
    end  
    
    % Check if power values should be saved from the CPSD generated in ds
    % calculations
    if any(ismember('power', options.featureList))
        power = zeros(nFreq, C, W);
        for k =1:C
            power(:,k,:) = S(k,k,:,:);
        end

        % generate/save matrix where elements name corresponding
        % feature in power matrix
        freqMat = repmat(fStrings, [1 C]);
        areaMat = repmat(labels.area', [nFreq 1]);
        labels.powerFeatures = string(cellfun(@(x,y) [x ' ' y], areaMat, freqMat, ...
        'UniformOutput', false));

        save(saveFile, 'power', '-append')
    end
end

% phase slope index
if ismember('psi', options.featureList)
    segleng = fs; % 1 Hz frequency resolution
    fBins = [f(1:end-1)', (f(1:end-1)'+1), (f(1:end-1)'+2)];
    
    % create feature labels
    psiFeatures = cell(C,C,nFreq-1);
    psiFeatures(:) = {''};
    areaList = labels.area;
    for c1 = 1:C
        for c2 = 1:C
            if c1==c2, continue, end
            % save feature labels for this pair of regions
            psiFeatures(c1,c2,:) = cellfun(@(x) [areaList{c1} '~>' areaList{c2} ' ' x], ...
                fStrings(1:end-1), 'UniformOutput', false);
        end
    end
    
    labels.psiFeatures = string(psiFeatures);
    
    if options.parCores, pp = parpool([2 options.parCores]); end
    
    % calculate psi for each window
    psi = zeros(C,C,nFreq-1,W, 'single');
    parfor (w = 1:W, options.parCores)
        thisData = X(:,:,w);
        psi(:,:,:,w) = data2psi(thisData, segleng, [], fBins);
    end
    if options.parCores, delete(pp), end
    
    save(saveFile, 'psi', '-append')
end

% partial directed coherence and directed transfer function
if any(ismember({'pdc','dtf'}, options.featureList))
    % generate directed spectrum values matrix in the form MxPxQ, where M
    % iterates over pairs of brain regions, P is frequency, and Q is time
    % window.
    X = double(X);
    
    [pdc, dtf, pdFeatures] = pdc_dtf(X, labels.area, fs, f, options);
    
    labels.pdFeatures = string(pdFeatures);
    
    save(saveFile, 'pdc', 'dtf', '-append')
end


%% Take Fourier transform of data
if any(ismember('fft', options.featureList))
    Ns = ceil(N/2);
    scale = 1/N;
    xFft = scale*fft(double(X));
    xFft = 2*(xFft(2:Ns+1,:,:));

    save(saveFile,'xFft','-append');
end

%% Save labels to JSON
save(saveFile,'labels','-append');
saveJson(saveFile, labels)

end

function opts = fillDefaultOpts(opts)
    if ~isfield(opts,'featureList')
        opts.featureList = {'directedSpectrum', 'pwDirectedSpectrum', 'granger',...
            'c_granger', 'pdc', 'dtf', 'psi'};
    end
    if ~isfield(opts,'parCores'), opts.parCores = 0; end
    if ~isfield(opts, 'window'), opts.window = 200; end
    if ~isfield(opts, 'overlap'), opts.overlap = 175; end
    if ~isfield(opts,'mvgcFolder'), opts.mvgcFolder = '~/mvgc'; end
end
