function varargout = process_fooof(varargin)
% PROCESS_FOOOF: Applies the "Fitting Oscillations and One Over F" (specparam) algorithm on a Welch's PSD
%
% REFERENCE: Please cite the original algorithm:
%    Donoghue T, Haller M, Peterson E, Varma P, Sebastian P, Gao R, Noto T,
%    Lara AH, Wallis JD, Knight RT, Shestyuk A, Voytek B. Parameterizing
%    neural power spectra into periodic and aperiodic components.
%    Nature Neuroscience (2020)

% @=============================================================================
% This function is part of the Brainstorm software:
% https://neuroimage.usc.edu/brainstorm
%
% Copyright (c) University of Southern California & McGill University
% This software is distributed under the terms of the GNU General Public License
% as published by the Free Software Foundation. Further details on the GPLv3
% license can be found at http://www.gnu.org/copyleft/gpl.html.
%
% FOR RESEARCH PURPOSES ONLY. THE SOFTWARE IS PROVIDED "AS IS," AND THE
% UNIVERSITY OF SOUTHERN CALIFORNIA AND ITS COLLABORATORS DO NOT MAKE ANY
% WARRANTY, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO WARRANTIES OF
% MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, NOR DO THEY ASSUME ANY
% LIABILITY OR RESPONSIBILITY FOR THE USE OF THIS SOFTWARE.
%
% For more information type "brainstorm license" at command prompt.
% =============================================================================@
%
% Authors: Luc Wilson, Francois Tadel, 2020-2022

eval(macro_method);
end

%% ===== GET DESCRIPTION =====
function sProcess = GetDescription() %#ok<DEFNU>
    % Description the process
    sProcess.Comment     = 'specparam: Fitting oscillations and 1/f';
    sProcess.Category    = 'Custom';
    sProcess.SubGroup    = 'Frequency';
    sProcess.Index       = 490;
    sProcess.Description = 'https://neuroimage.usc.edu/brainstorm/Tutorials/Fooof';
    % Definition of the input accepted by this process
    sProcess.InputTypes  = {'timefreq'};
    sProcess.OutputTypes = {'timefreq'};
    sProcess.nInputs     = 1;
    sProcess.nMinFiles   = 1;
    % Definition of the options
    % === FOOOF TYPE
    sProcess.options.implementation.Comment = {'Matlab', 'Python 3 (3.7 recommended)', 'specparam implementation:'; 'matlab', 'python', ''};
    sProcess.options.implementation.Type    = 'radio_linelabel';
    sProcess.options.implementation.Value   = 'matlab';
    sProcess.options.implementation.Controller.matlab = 'Matlab';
    sProcess.options.implementation.Controller.python = 'Python';
    % === FREQUENCY RANGE
    sProcess.options.freqrange.Comment = 'Frequency range for analysis: ';
    sProcess.options.freqrange.Type    = 'freqrange_static';   % 'freqrange'
    sProcess.options.freqrange.Value   = {[1 40], 'Hz', 1};
    % === POWER LINE
    sProcess.options.powerline.Comment = {'None', '50 Hz', '60 Hz', 'Ignore power line frequencies:'; '-5', '50', '60', ''};
    sProcess.options.powerline.Type    = 'radio_linelabel';
    sProcess.options.powerline.Value   = '60';
    sProcess.options.powerline.Class   = 'Matlab';
    % === PEAK TYPE
    sProcess.options.peaktype.Comment = {'Gaussian', 'Cauchy*', 'Best of both* (* experimental)', 'Peak model:'; 'gaussian', 'cauchy', 'best', ''};
    sProcess.options.peaktype.Type    = 'radio_linelabel';
    sProcess.options.peaktype.Value   = 'gaussian';
    sProcess.options.peaktype.Class   = 'Matlab';
    % === PEAK WIDTH LIMITS
    sProcess.options.peakwidth.Comment = 'Peak width limits (default=[0.5-12]): ';
    sProcess.options.peakwidth.Type    = 'freqrange_static';
    sProcess.options.peakwidth.Value   = {[0.5 12], 'Hz', 1};
    % === MAX PEAKS
    sProcess.options.maxpeaks.Comment = 'Maximum number of peaks (default=3): ';
    sProcess.options.maxpeaks.Type    = 'value';
    sProcess.options.maxpeaks.Value   = {3, '', 0};
    % === MEAN PEAK HEIGHT
    sProcess.options.minpeakheight.Comment = 'Minimum peak height (default=3): ';
    sProcess.options.minpeakheight.Type    = 'value';
    sProcess.options.minpeakheight.Value   = {3, 'dB', 1};
    % === PROXIMITY THRESHOLD
    sProcess.options.proxthresh.Comment = 'Proximity threshold (default=2): ';
    sProcess.options.proxthresh.Type    = 'value';
    sProcess.options.proxthresh.Value   = {2, 'stdev of peak model', 1};
    sProcess.options.proxthresh.Class   = 'Matlab';
    % === APERIODIC MODE
    sProcess.options.apermode.Comment = {'Fixed', 'Knee', 'Aperiodic mode (default=fixed):'; 'fixed', 'knee', ''};
    sProcess.options.apermode.Type    = 'radio_linelabel';
    sProcess.options.apermode.Value   = 'fixed';
    % === GUESS WEIGHT
    sProcess.options.guessweight.Comment = {'None', 'Weak', 'Strong', 'Guess weight (default=none):'; 'none', 'weak', 'strong', ''};
    sProcess.options.guessweight.Type    = 'radio_linelabel';
    sProcess.options.guessweight.Value   = 'none';
    sProcess.options.guessweight.Class   = 'Matlab';

    % === SORT PEAKS TYPE
    sProcess.options.sorttype.Comment = {'Peak parameters', 'Frequency bands', 'Sort peaks using:'; 'param', 'band', ''};
    sProcess.options.sorttype.Type    = 'radio_linelabel';
    sProcess.options.sorttype.Value   = 'param';
    sProcess.options.sorttype.Controller.param = 'Param';
    sProcess.options.sorttype.Controller.band = 'Band';
    sProcess.options.sorttype.Group   = 'output';
    % === SORT PEAKS PARAM
    sProcess.options.sortparam.Comment = {'Frequency', 'Amplitude', 'Std dev.', 'Sort by peak...'; 'frequency', 'amplitude', 'std', ''};
    sProcess.options.sortparam.Type    = 'radio_linelabel';
    sProcess.options.sortparam.Value   = 'frequency';
    sProcess.options.sortparam.Class   = 'Param';
    sProcess.options.sortparam.Group   = 'output';
    % === SORT FREQ BANDS
    DefaultFreqBands = bst_get('DefaultFreqBands');
    sProcess.options.sortbands.Comment = '';
    sProcess.options.sortbands.Type    = 'groupbands';
    sProcess.options.sortbands.Value   = DefaultFreqBands(:,1:2);
    sProcess.options.sortbands.Class   = 'Band';
    sProcess.options.sortbands.Group   = 'output';
end

%% ===== FORMAT COMMENT =====
function Comment = FormatComment(sProcess) %#ok<DEFNU>
    Comment = sProcess.Comment;
end

%% ===== RUN =====
function OutputFile = Run(sProcess, sInputs) %#ok<DEFNU>
    % Initialize returned list of files
    OutputFile = {};

    % Fetch user settings
    implementation = sProcess.options.implementation.Value;
    opt.freq_range          = sProcess.options.freqrange.Value{1};
    opt.peak_width_limits   = sProcess.options.peakwidth.Value{1};
    opt.max_peaks           = sProcess.options.maxpeaks.Value{1};
    opt.min_peak_height     = sProcess.options.minpeakheight.Value{1} / 10; % convert from dB to B
    opt.aperiodic_mode      = sProcess.options.apermode.Value;
    opt.peak_threshold      = 2;   % 2 std dev: parameter for interface simplification
    opt.border_threshold    = 1;   % 1 std dev: proximity to edge of spectrum, static in Python
    opt.return_spectrum     = 0;   % SPM/FT: set to 1
    % Matlab-only options
    opt.power_line          = sProcess.options.powerline.Value;
    opt.peak_type           = sProcess.options.peaktype.Value;
    opt.proximity_threshold = sProcess.options.proxthresh.Value{1};
    opt.guess_weight        = sProcess.options.guessweight.Value;
    opt.thresh_after        = true;   % Threshold after fitting always selected for Matlab (mirrors the Python FOOOF closest by removing peaks that do not satisfy a user's predetermined conditions)
    % Python-only options
    opt.verbose    = false;
    % Output options
    opt.sort_type  = sProcess.options.sorttype.Value;
    opt.sort_param = sProcess.options.sortparam.Value;
	opt.sort_bands = sProcess.options.sortbands.Value;

    % Check input frequency bounds
    if (any(opt.freq_range < 0) || opt.freq_range(1) >= opt.freq_range(2))
        bst_report('error','Invalid Frequency range');
        return
    end

    hasOptimTools = 0;
    if exist('fmincon') == 2 && strcmp(implementation,'matlab')
        hasOptimTools = 1;
        disp('Using constrained optimization, Guess Weight ignored.')
    end

    % Initialize returned list of files
    OutputFile = {};
    for iFile = 1:length(sInputs)
        bst_progress('text',['Standby: FOOOFing spectrum ' num2str(iFile) ' of ' num2str(length(sInputs))]);
        % Load input file
        PsdMat = in_bst_timefreq(sInputs(iFile).FileName);
        % Exclude 0Hz from the computation
        if (opt.freq_range(1) == 0) && (PsdMat.Freqs(1) == 0) && (length(PsdMat.Freqs) >= 2)
            opt.freq_range(1) = PsdMat.Freqs(2);
        end

        % === COMPUTE FOOOF MODEL ===
        % Switch between implementations
        switch (implementation)
            case 'matlab'   % Matlab standalone FOOOF
                [FOOOF_freqs, FOOOF_data] = FOOOF_matlab(PsdMat.TF, PsdMat.Freqs, opt, hasOptimTools);
            case 'python'
                opt.peak_type = 'gaussian';
                [FOOOF_freqs, FOOOF_data] = process_fooof_py('FOOOF_python', PsdMat.TF, PsdMat.Freqs, opt);
                % Remove unnecessary structure level, allowing easy concatenation across channels, e.g. for display.
                FOOOF_data = FOOOF_data.FOOOF;
            otherwise
                error('Invalid implentation.');
        end

        % === FOOOF ANALYSIS ===
        TFfooof = PsdMat.TF(:,1,ismember(PsdMat.Freqs,FOOOF_freqs));
        [ePeaks, eAperiodics, eStats] = FOOOF_analysis(FOOOF_data, PsdMat.RowNames, TFfooof, opt.max_peaks, opt.sort_type, opt.sort_param, opt.sort_bands);

        % === PREPARE OUTPUT STRUCTURE ===
        % Create file structure
        PsdMat.Options.FOOOF = struct(...
            'options',    opt, ...
            'freqs',      FOOOF_freqs, ...
            'data',       FOOOF_data, ...
            'peaks',      ePeaks, ...
            'aperiodics', eAperiodics, ...
            'stats',      eStats);
        % Comment: Add FOOOF
        if ~isempty(strfind(PsdMat.Comment, 'PSD:'))
            PsdMat.Comment = strrep(PsdMat.Comment, 'PSD:', 'specparam:');
        else
            PsdMat.Comment = strcat(PsdMat.Comment, ' | specparam');
        end
        % History: Computation
        PsdMat = bst_history('add', PsdMat, 'compute', 'specparam');

        % === SAVE FILE ===
        % Filename: add _fooof tag
        [fPath, fName, fExt] = bst_fileparts(file_fullpath(sInputs(iFile).FileName));
        NewFile = file_unique(bst_fullfile(fPath, [fName, '_specparam', fExt]));
        % Save file
        bst_save(NewFile, PsdMat, 'v6');
        % Add file to database structure
        db_add_data(sInputs(iFile).iStudy, NewFile, PsdMat);
        % Return new file
        OutputFile{end+1} = NewFile;
    end
end

%% ===================================================================================
%  ===== MATLAB FOOOF ================================================================
%  ===================================================================================

%% ===== MATLAB STANDALONE FOOOF =====
function [fs, fg] = FOOOF_matlab(TF, Freqs, opt, hOT)
    % Find all frequency values within user limits
    fMask = (round(Freqs.*10)./10 >= opt.freq_range(1)) & (round(Freqs.*10)./10 <= opt.freq_range(2)) & ~mod(sum(abs(round(Freqs.*10)./10-[1;2;3].*str2double(opt.power_line)) >= 2),3);
    fs = Freqs(fMask);
    spec = log10(squeeze(TF(:,1,fMask))); % extract log spectra
    nChan = size(TF,1);
    if nChan == 1, spec = spec'; end
    % Initalize FOOOF structs
    fg(nChan) = struct(...
            'aperiodic_params', [],...
            'peak_params',      [],...
            'peak_types',       '',...
            'ap_fit',           [],...
            'fooofed_spectrum', [],...
            'peak_fit',         [],...
            'error',            [],...
            'r_squared',        []);
    % Iterate across channels
    for chan = 1:nChan
        bst_progress('set', round(chan./nChan.*100));
        % Fit aperiodic
        aperiodic_pars = robust_ap_fit(fs, spec(chan,:), opt.aperiodic_mode);
        % Remove aperiodic
        flat_spec = flatten_spectrum(fs, spec(chan,:), aperiodic_pars, opt.aperiodic_mode);
        % Fit peaks
        [peak_pars, peak_function] = fit_peaks(fs, flat_spec, opt.max_peaks, opt.peak_threshold, opt.min_peak_height, ...
            opt.peak_width_limits/2, opt.proximity_threshold, opt.border_threshold, opt.peak_type, opt.guess_weight,hOT);
        if opt.thresh_after && ~hOT  % Check thresholding requirements are met for unbounded optimization
            peak_pars(peak_pars(:,2) < opt.min_peak_height,:)     = []; % remove peaks shorter than limit
            peak_pars(peak_pars(:,3) < opt.peak_width_limits(1)/2,:)  = []; % remove peaks narrower than limit
            peak_pars(peak_pars(:,3) > opt.peak_width_limits(2)/2,:)  = []; % remove peaks broader than limit
            peak_pars = drop_peak_cf(peak_pars, opt.border_threshold, opt.freq_range); % remove peaks outside frequency limits
            peak_pars(peak_pars(:,1) < 0,:) = []; % remove peaks with a centre frequency less than zero (bypass drop_peak_cf)
            peak_pars = drop_peak_overlap(peak_pars, opt.proximity_threshold); % remove smallest of two peaks fit too closely
        end
        % Refit aperiodic
        aperiodic = spec(chan,:);
        for peak = 1:size(peak_pars,1)
            aperiodic = aperiodic - peak_function(fs,peak_pars(peak,1), peak_pars(peak,2), peak_pars(peak,3));
        end
        aperiodic_pars = simple_ap_fit(fs, aperiodic, opt.aperiodic_mode);
        % Generate model fit
        ap_fit = gen_aperiodic(fs, aperiodic_pars, opt.aperiodic_mode);
        model_fit = ap_fit;
        for peak = 1:size(peak_pars,1)
            model_fit = model_fit + peak_function(fs,peak_pars(peak,1),...
                peak_pars(peak,2),peak_pars(peak,3));
        end
        % Calculate model error
        MSE = sum((spec(chan,:) - model_fit).^2)/length(model_fit);
        rsq_tmp = corrcoef(spec(chan,:),model_fit).^2;
        % Return FOOOF results
        aperiodic_pars(2) = abs(aperiodic_pars(2));
        fg(chan).aperiodic_params = aperiodic_pars;
        fg(chan).peak_params      = peak_pars;
        fg(chan).peak_types       = func2str(peak_function);
        fg(chan).ap_fit           = 10.^ap_fit;
        fg(chan).fooofed_spectrum = 10.^model_fit;
        fg(chan).peak_fit         = 10.^(model_fit-ap_fit);
        fg(chan).error            = MSE;
        fg(chan).r_squared        = rsq_tmp(2);
        if opt.return_spectrum
            fg(chan).power_spectrum = spec(chan,:);
        end
        %plot(fs', [fg(chan).ap_fit', fg(chan).peak_fit', fg(chan).fooofed_spectrum'])
    end
end

%% ===== GENERATE APERIODIC =====
function ap_vals = gen_aperiodic(freqs,aperiodic_params,aperiodic_mode)
%       Generate aperiodic values, from parameter definition.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%       	Frequency vector to create aperiodic component for.
%       aperiodic_params : 1x3 array
%           Parameters that define the aperiodic component.
%       aperiodic_mode : {'fixed', 'knee'}
%           Defines absence or presence of knee in aperiodic component.
%
%       Returns
%       -------
%       ap_vals : 1d array
%           Generated aperiodic values.

    switch aperiodic_mode
        case 'fixed'  % no knee
            ap_vals = expo_nk_function(freqs,aperiodic_params);
        case 'knee'
            ap_vals = expo_function(freqs,aperiodic_params);
        case 'floor'
            ap_vals = expo_fl_function(freqs,aperiodic_params);
    end
end

%% ===== CORE MODELS =====
function ys = gaussian(freqs, mu, hgt, sigma)
%       Gaussian function to use for fitting.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency vector to create gaussian fit for.
%       mu, hgt, sigma : doubles
%           Parameters that define gaussian function (centre frequency,
%           height, and standard deviation).
%
%       Returns
%       -------
%       ys :    1xn array
%       Output values for gaussian function.

    ys = hgt*exp(-(((freqs-mu)./sigma).^2) /2);

end

function ys = cauchy(freqs, ctr, hgt, gam)
%       Cauchy function to use for fitting.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency vector to create cauchy fit for.
%       ctr, hgt, gam : doubles
%           Parameters that define cauchy function (centre frequency,
%           height, and "standard deviation" [gamma]).
%
%       Returns
%       -------
%       ys :    1xn array
%       Output values for cauchy function.

    ys = hgt./(1+((freqs-ctr)/gam).^2);

end

function ys = expo_function(freqs,params)
%       Exponential function to use for fitting 1/f, with a 'knee' (maximum at low frequencies).
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Input x-axis values.
%       params : 1x3 array (offset, knee, exp)
%           Parameters (offset, knee, exp) that define Lorentzian function:
%           y = 10^offset * (1/(knee + x^exp))
%
%       Returns
%       -------
%       ys :    1xn array
%           Output values for exponential function.

    ys = params(1) - log10(abs(params(2)) +freqs.^params(3));

end

function ys = expo_nk_function(freqs, params)
%       Exponential function to use for fitting 1/f, without a 'knee'.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Input x-axis values.
%       params : 1x2 array (offset, exp)
%           Parameters (offset, exp) that define Lorentzian function:
%           y = 10^offset * (1/(x^exp))
%
%       Returns
%       -------
%       ys :    1xn array
%           Output values for exponential (no-knee) function.

    ys = params(1) - log10(freqs.^params(2));

end

function ys = expo_fl_function(freqs, params)

    ys = log10(f.^(params(1)) * 10^(params(2)) + params(3));

end

%% ===== FITTING ALGORITHM =====
function aperiodic_params = simple_ap_fit(freqs, power_spectrum, aperiodic_mode)
%       Fit the aperiodic component of the power spectrum.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency values for the power spectrum, in linear scale.
%       power_spectrum : 1xn array
%           Power values, in log10 scale.
%       aperiodic_mode : {'fixed','knee'}
%           Defines absence or presence of knee in aperiodic component.
%
%       Returns
%       -------
%       aperiodic_params : 1xn array
%           Parameter estimates for aperiodic fit.

%       Set guess params for lorentzian aperiodic fit, guess params set at init
    options = optimset('Display', 'off', 'TolX', 1e-4, 'TolFun', 1e-6, ...
        'MaxFunEvals', 5000, 'MaxIter', 5000);

    switch (aperiodic_mode)
        case 'fixed'  % no knee
            exp_guess = -(power_spectrum(end)-power_spectrum(1))./log10(freqs(end)./freqs(1));
            guess_vec = [power_spectrum(1), exp_guess];
            aperiodic_params = fminsearch(@error_expo_nk_function, guess_vec, options, freqs, power_spectrum);
        case 'knee'
            exp_guess = -(power_spectrum(end)-power_spectrum(1))./log10(freqs(end)./freqs(1));
            guess_vec = [power_spectrum(1),0, exp_guess];
            aperiodic_params = fminsearch(@error_expo_function, guess_vec, options, freqs, power_spectrum);
    end

end

function aperiodic_params = robust_ap_fit(freqs, power_spectrum, aperiodic_mode)
%       Fit the aperiodic component of the power spectrum robustly, ignoring outliers.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency values for the power spectrum, in linear scale.
%       power_spectrum : 1xn array
%           Power values, in log10 scale.
%       aperiodic_mode : {'fixed','knee'}
%           Defines absence or presence of knee in aperiodic component.
%
%       Returns
%       -------
%       aperiodic_params : 1xn array
%           Parameter estimates for aperiodic fit.

    % Do a quick, initial aperiodic fit
    popt = simple_ap_fit(freqs, power_spectrum, aperiodic_mode);
    initial_fit = gen_aperiodic(freqs, popt, aperiodic_mode);

    % Flatten power_spectrum based on initial aperiodic fit
    flatspec = power_spectrum - initial_fit;

    % Flatten outliers - any points that drop below 0
    flatspec(flatspec(:) < 0) = 0;

    % Use percential threshold, in terms of # of points, to extract and re-fit
    perc_thresh = bst_prctile(flatspec, 0.025);
    perc_mask = flatspec <= perc_thresh;
    freqs_ignore = freqs(perc_mask);
    spectrum_ignore = power_spectrum(perc_mask);

    % Second aperiodic fit - using results of first fit as guess parameters

    options = optimset('Display', 'off', 'TolX', 1e-4, 'TolFun', 1e-6, ...
        'MaxFunEvals', 5000, 'MaxIter', 5000);
    guess_vec = popt;

    switch (aperiodic_mode)
        case 'fixed'  % no knee
            aperiodic_params = fminsearch(@error_expo_nk_function, guess_vec, options, freqs_ignore, spectrum_ignore);
        case 'knee'
            aperiodic_params = fminsearch(@error_expo_function, guess_vec, options, freqs_ignore, spectrum_ignore);
    end
end

function spectrum_flat = flatten_spectrum(freqs, power_spectrum, robust_aperiodic_params, aperiodic_mode)
%       Flatten the power spectrum by removing the aperiodic component.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency values for the power spectrum, in linear scale.
%       power_spectrum : 1xn array
%           Power values, in log10 scale.
%       robust_aperiodic_params : 1x2 or 1x3 array (see aperiodic_mode)
%           Parameter estimates for aperiodic fit.
%       aperiodic_mode : 1 or 2
%           Defines absence or presence of knee in aperiodic component.
%
%       Returns
%       -------
%       spectrum_flat : 1xn array
%           Flattened (aperiodic removed) power spectrum.

spectrum_flat = power_spectrum - gen_aperiodic(freqs,robust_aperiodic_params,aperiodic_mode);

end

function [model_params,peak_function] = fit_peaks(freqs, flat_iter, max_n_peaks, peak_threshold, min_peak_height, gauss_std_limits, proxThresh, bordThresh, peakType, guess_weight,hOT)
%       Iteratively fit peaks to flattened spectrum.
%
%       Parameters
%       ----------
%       freqs : 1xn array
%           Frequency values for the power spectrum, in linear scale.
%       flat_iter : 1xn array
%           Flattened (aperiodic removed) power spectrum.
%       max_n_peaks : double
%           Maximum number of gaussians to fit within the spectrum.
%       peak_threshold : double
%           Threshold (in standard deviations of noise floor) to detect a peak.
%       min_peak_height : double
%           Minimum height of a peak (in log10).
%       gauss_std_limits : 1x2 double
%           Limits to gaussian (cauchy) standard deviation (gamma) when detecting a peak.
%       proxThresh : double
%           Minimum distance between two peaks, in st. dev. (gamma) of peaks.
%       peakType : {'gaussian', 'cauchy', 'both'}
%           Which types of peaks are being fitted
%       guess_weight : {'none', 'weak', 'strong'}
%           Parameter to weigh initial estimates during optimization (None, Weak, or Strong)
%       hOT : 0 or 1
%           Defines whether to use constrained optimization, fmincon, or
%           basic simplex, fminsearch.
%
%       Returns
%       -------
%       gaussian_params : mx3 array, where m = No. of peaks.
%           Parameters that define the peak fit(s). Each row is a peak, as [mean, height, st. dev. (gamma)].
    switch peakType
        case 'gaussian' % gaussian only
            peak_function = @gaussian; % Identify peaks as gaussian
            % Initialize matrix of guess parameters for gaussian fitting.
            guess_params = zeros(max_n_peaks, 3);
            % Save intact flat_spectrum
            flat_spec = flat_iter;
            % Find peak: Loop through, finding a candidate peak, and fitting with a guess gaussian.
            % Stopping procedure based on either the limit on # of peaks,
            % or the relative or absolute height thresholds.
            for guess = 1:max_n_peaks
                % Find candidate peak - the maximum point of the flattened spectrum.
                max_ind = find(flat_iter == max(flat_iter));
                max_height = flat_iter(max_ind);

                % Stop searching for peaks once max_height drops below height threshold.
                if max_height <= peak_threshold * std(flat_iter)
                    break
                end

                % Set the guess parameters for gaussian fitting - mean and height.
                guess_freq = freqs(max_ind);
                guess_height = max_height;

                % Halt fitting process if candidate peak drops below minimum height.
                if guess_height <= min_peak_height
                    break
                end

                % Data-driven first guess at standard deviation
                % Find half height index on each side of the center frequency.
                half_height = 0.5 * max_height;

                le_ind = sum(flat_iter(1:max_ind) <= half_height);
                ri_ind = length(flat_iter) - sum(flat_iter(max_ind:end) <= half_height)+1;

                % Keep bandwidth estimation from the shortest side.
                % We grab shortest to avoid estimating very large std from overalapping peaks.
                % Grab the shortest side, ignoring a side if the half max was not found.
                % Note: will fail if both le & ri ind's end up as None (probably shouldn't happen).
                short_side = min(abs([le_ind,ri_ind]-max_ind));

                % Estimate std from FWHM. Calculate FWHM, converting to Hz, get guess std from FWHM
                fwhm = short_side * 2 * (freqs(2)-freqs(1));
                guess_std = fwhm / (2 * sqrt(2 * log(2)));

                % Check that guess std isn't outside preset std limits; restrict if so.
                % Note: without this, curve_fitting fails if given guess > or < bounds.
                if guess_std < gauss_std_limits(1)
                    guess_std = gauss_std_limits(1);
                end
                if guess_std > gauss_std_limits(2)
                    guess_std = gauss_std_limits(2);
                end

                % Collect guess parameters.
                guess_params(guess,:) = [guess_freq, guess_height, guess_std];

                % Subtract best-guess gaussian.
                peak_gauss = gaussian(freqs, guess_freq, guess_height, guess_std);
                flat_iter = flat_iter - peak_gauss;

            end
            % Remove unused guesses
            guess_params(guess_params(:,1) == 0,:) = [];

            % Check peaks based on edges, and on overlap
            % Drop any that violate requirements.
            guess_params = drop_peak_cf(guess_params, bordThresh, [min(freqs) max(freqs)]);
            guess_params = drop_peak_overlap(guess_params, proxThresh);

            % If there are peak guesses, fit the peaks, and sort results.
            if ~isempty(guess_params)
                model_params = fit_peak_guess(guess_params, freqs, flat_spec, 1, guess_weight, gauss_std_limits,hOT);
            else
                model_params = zeros(1, 3);
            end

        case 'cauchy' % cauchy only
            peak_function = @cauchy; % Identify peaks as cauchy
            guess_params = zeros(max_n_peaks, 3);
            flat_spec = flat_iter;
            for guess = 1:max_n_peaks
                max_ind = find(flat_iter == max(flat_iter));
                max_height = flat_iter(max_ind);
                if max_height <= peak_threshold * std(flat_iter)
                    break
                end
                guess_freq = freqs(max_ind);
                guess_height = max_height;
                if guess_height <= min_peak_height
                    break
                end
                half_height = 0.5 * max_height;
                le_ind = sum(flat_iter(1:max_ind) <= half_height);
                ri_ind = length(flat_iter) - sum(flat_iter(max_ind:end) <= half_height);
                short_side = min(abs([le_ind,ri_ind]-max_ind));

                % Estimate gamma from FWHM. Calculate FWHM, converting to Hz, get guess gamma from FWHM
                fwhm = short_side * 2 * (freqs(2)-freqs(1));
                guess_gamma = fwhm/2;
                % Check that guess gamma isn't outside preset limits; restrict if so.
                % Note: without this, curve_fitting fails if given guess > or < bounds.
                if guess_gamma < gauss_std_limits(1)
                    guess_gamma = gauss_std_limits(1);
                end
                if guess_gamma > gauss_std_limits(2)
                    guess_gamma = gauss_std_limits(2);
                end

                % Collect guess parameters.
                guess_params(guess,:) = [guess_freq(1), guess_height, guess_gamma];

                % Subtract best-guess cauchy.
                peak_cauchy = cauchy(freqs, guess_freq(1), guess_height, guess_gamma);
                flat_iter = flat_iter - peak_cauchy;

            end
            guess_params(guess_params(:,1) == 0,:) = [];
            guess_params = drop_peak_cf(guess_params, bordThresh, [min(freqs) max(freqs)]);
            guess_params = drop_peak_overlap(guess_params, proxThresh);

            % If there are peak guesses, fit the peaks, and sort results.
            if ~isempty(guess_params)
                model_params = fit_peak_guess(guess_params, freqs, flat_spec, 2, guess_weight, gauss_std_limits,hOT);
            else
                model_params = zeros(1, 3);
            end
        case 'best' % best of both: model both fits and compare error, save best
            % Gaussian Fit
            guess_params = zeros(max_n_peaks, 3);
            flat_spec = flat_iter;
            for guess = 1:max_n_peaks
                max_ind = find(flat_iter == max(flat_iter));
                max_height = flat_iter(max_ind);
                if max_height <= peak_threshold * std(flat_iter)
                    break
                end
                guess_freq = freqs(max_ind);
                guess_height = max_height;
                if guess_height <= min_peak_height
                    break
                end
                half_height = 0.5 * max_height;
                le_ind = sum(flat_iter(1:max_ind) <= half_height);
                ri_ind = length(flat_iter) - sum(flat_iter(max_ind:end) <= half_height)+1;
                short_side = min(abs([le_ind,ri_ind]-max_ind));
                fwhm = short_side * 2 * (freqs(2)-freqs(1));
                guess_std = fwhm / (2 * sqrt(2 * log(2)));
                if guess_std < gauss_std_limits(1)
                    guess_std = gauss_std_limits(1);
                end
                if guess_std > gauss_std_limits(2)
                    guess_std = gauss_std_limits(2);
                end
                guess_params(guess,:) = [guess_freq, guess_height, guess_std];
                peak_gauss = gaussian(freqs, guess_freq, guess_height, guess_std);
                flat_iter = flat_iter - peak_gauss;
            end
            guess_params(guess_params(:,1) == 0,:) = [];
            guess_params = drop_peak_cf(guess_params, bordThresh, [min(freqs) max(freqs)]);
            guess_params = drop_peak_overlap(guess_params, proxThresh);
            if ~isempty(guess_params)
                gauss_params = fit_peak_guess(guess_params, freqs, flat_spec, 1, guess_weight, gauss_std_limits,hOT);
                flat_gauss = zeros(size(freqs));
                for peak = 1:size(gauss_params,1)
                    flat_gauss =  flat_gauss + gaussian(freqs,gauss_params(peak,1),...
                        gauss_params(peak,2),gauss_params(peak,3));
                end
                error_gauss = sum((flat_gauss-flat_spec).^2);
            else
                gauss_params = zeros(1, 3); error_gauss = 1E10;
            end

            % Cauchy Fit
            guess_params = zeros(max_n_peaks, 3);
            flat_iter = flat_spec;
            for guess = 1:max_n_peaks
                max_ind = find(flat_iter == max(flat_iter));
                max_height = flat_iter(max_ind);
                if max_height <= peak_threshold * std(flat_iter)
                    break
                end
                guess_freq = freqs(max_ind);
                guess_height = max_height;
                if guess_height <= min_peak_height
                    break
                end
                half_height = 0.5 * max_height;
                le_ind = sum(flat_iter(1:max_ind) <= half_height);
                ri_ind = length(flat_iter) - sum(flat_iter(max_ind:end) <= half_height)+1;
                short_side = min(abs([le_ind,ri_ind]-max_ind));
                fwhm = short_side * 2 * (freqs(2)-freqs(1));
                guess_gamma = fwhm/2;
                if guess_gamma < gauss_std_limits(1)
                    guess_gamma = gauss_std_limits(1);
                end
                if guess_gamma > gauss_std_limits(2)
                    guess_gamma = gauss_std_limits(2);
                end
                guess_params(guess,:) = [guess_freq(1), guess_height, guess_gamma];
                peak_cauchy = cauchy(freqs, guess_freq(1), guess_height, guess_gamma);
                flat_iter = flat_iter - peak_cauchy;
            end
            guess_params(guess_params(:,1) == 0,:) = [];
            guess_params = drop_peak_cf(guess_params, bordThresh, [min(freqs) max(freqs)]);
            guess_params = drop_peak_overlap(guess_params, proxThresh);
            if ~isempty(guess_params)
                cauchy_params = fit_peak_guess(guess_params, freqs, flat_spec, 2, guess_weight, gauss_std_limits,hOT);
                flat_cauchy = zeros(size(freqs));
                for peak = 1:size(cauchy_params,1)
                    flat_cauchy =  flat_cauchy + cauchy(freqs,cauchy_params(peak,1),...
                        cauchy_params(peak,2),cauchy_params(peak,3));
                end
                error_cauchy = sum((flat_cauchy-flat_spec).^2);
            else
                cauchy_params = zeros(1, 3); error_cauchy = 1E10;
            end
            % Save least-error model
                if min([error_gauss,error_cauchy]) == error_gauss
                    model_params = gauss_params;
                    peak_function = @gaussian;
                else
                    model_params = cauchy_params;
                    peak_function = @cauchy;
                end
    end

end

function guess = drop_peak_cf(guess, bw_std_edge, freq_range)
%       Check whether to drop peaks based on center's proximity to the edge of the spectrum.
%
%       Parameters
%       ----------
%       guess : mx3 array, where m = No. of peaks.
%           Guess parameters for peak fits.
%
%       Returns
%       -------
%       guess : qx3 where q <= m No. of peaks.
%           Guess parameters for peak fits.

    cf_params = guess(:,1)';
    bw_params = guess(:,3)' * bw_std_edge;

    % Check if peaks within drop threshold from the edge of the frequency range.

    keep_peak = abs(cf_params-freq_range(1)) > bw_params & ...
        abs(cf_params-freq_range(2)) > bw_params;

    % Drop peaks that fail the center frequency edge criterion
    guess = guess(keep_peak,:);

end

function guess = drop_peak_overlap(guess, proxThresh)
%       Checks whether to drop gaussians based on amount of overlap.
%
%       Parameters
%       ----------
%       guess : mx3 array, where m = No. of peaks.
%           Guess parameters for peak fits.
%       proxThresh: double
%           Proximity threshold (in st. dev. or gamma) between two peaks.
%
%       Returns
%       -------
%       guess : qx3 where q <= m No. of peaks.
%           Guess parameters for peak fits.
%
%       Note
%       -----
%       For any gaussians with an overlap that crosses the threshold,
%       the lowest height guess guassian is dropped.

    % Sort the peak guesses, so can check overlap of adjacent peaks
    guess = sortrows(guess);

    % Calculate standard deviation bounds for checking amount of overlap

    bounds = [guess(:,1) - guess(:,3) * proxThresh, ...
        guess(:,1), guess(:,1) + guess(:,3) * proxThresh];

    % Loop through peak bounds, comparing current bound to that of next peak
    drop_inds =  [];

    for ind = 1:size(bounds,1)-1

        b_0 = bounds(ind,:);
        b_1 = bounds(ind + 1,:);

        % Check if bound of current peak extends into next peak
        if b_0(2) > b_1(1)
            % If so, get the index of the gaussian with the lowest height (to drop)
            drop_inds = [drop_inds (ind - 1 + find(guess(ind:ind+1,2) == ...
                min(guess(ind,2),guess(ind+1,2))))];
        end
    end
    % Drop any peaks guesses that overlap too much, based on threshold.
    guess(drop_inds,:) = [];
end

function peak_params = fit_peak_guess(guess, freqs, flat_spec, peak_type, guess_weight, std_limits, hOT)
%     Fits a group of peak guesses with a fit function.
%
%     Parameters
%     ----------
%       guess : mx3 array, where m = No. of peaks.
%           Guess parameters for peak fits.
%       freqs : 1xn array
%           Frequency values for the power spectrum, in linear scale.
%       flat_iter : 1xn array
%           Flattened (aperiodic removed) power spectrum.
%       peakType : {'gaussian', 'cauchy', 'best'}
%           Which types of peaks are being fitted.
%       guess_weight : 'none', 'weak', 'strong'
%           Parameter to weigh initial estimates during optimization.
%       std_limits: 1x2 array
%           Minimum and maximum standard deviations for distribution.
%       hOT : 0 or 1
%           Defines whether to use constrained optimization, fmincon, or
%           basic simplex, fminsearch.
%
%       Returns
%       -------
%       peak_params : mx3, where m =  No. of peaks.
%           Peak parameters post-optimization.

    if hOT % Use OptimToolbox for fmincon
        options = optimset('Display', 'off', 'TolX', 1e-3, 'TolFun', 1e-5, ...
        'MaxFunEvals', 3000, 'MaxIter', 3000); % Tuned options
        lb = [guess(:,1)-guess(:,3)*2,zeros(size(guess(:,2))),ones(size(guess(:,3)))*std_limits(1)];
        ub = [guess(:,1)+guess(:,3)*2,inf(size(guess(:,2))),ones(size(guess(:,3)))*std_limits(2)];
        peak_params = fmincon(@error_model_constr,guess,[],[],[],[], ...
            lb,ub,[],options,freqs,flat_spec, peak_type);
    else % Use basic simplex approach, fminsearch, with guess_weight
        options = optimset('Display', 'off', 'TolX', 1e-4, 'TolFun', 1e-5, ...
        'MaxFunEvals', 5000, 'MaxIter', 5000);
        peak_params = fminsearch(@error_model,...
            guess, options, freqs, flat_spec, peak_type, guess, guess_weight);
    end
end

%% ===== ERROR FUNCTIONS =====
function err = error_expo_nk_function(params,xs,ys)
    ym = -log10(xs.^params(2)) + params(1);
    err = sum((ys - ym).^2);
end

function err = error_expo_function(params,xs,ys)
    ym = expo_function(xs,params);
    err = sum((ys - ym).^2);
end

function err = error_model(params, xVals, yVals, peak_type, guess, guess_weight)
    fitted_vals = 0;
    weak = 1E2;
    strong = 1E7;
    for set = 1:size(params,1)
        switch (peak_type)
            case 1 % Gaussian
                fitted_vals = fitted_vals + gaussian(xVals, params(set,1), params(set,2), params(set,3));
            case 2 % Cauchy
                fitted_vals = fitted_vals + cauchy(xVals, params(set,1), params(set,2), params(set,3));
        end
    end
    switch guess_weight
        case 'none'
            err = sum((yVals - fitted_vals).^2);
        case 'weak' % Add small weight to deviations from guess m and amp
            err = sum((yVals - fitted_vals).^2) + ...
                 weak*sum((params(:,1)-guess(:,1)).^2) + ...
                 weak*sum((params(:,2)-guess(:,2)).^2);
        case 'strong' % Add large weight to deviations from guess m and amp
            err = sum((yVals - fitted_vals).^2) + ...
                 strong*sum((params(:,1)-guess(:,1)).^2) + ...
                 strong*sum((params(:,2)-guess(:,2)).^2);
    end
end

function err = error_model_constr(params, xVals, yVals, peak_type)
    fitted_vals = 0;
    for set = 1:size(params,1)
        switch (peak_type)
            case 1 % Gaussian
                fitted_vals = fitted_vals + gaussian(xVals, params(set,1), params(set,2), params(set,3));
            case 2 % Cauchy
                fitted_vals = fitted_vals + cauchy(xVals, params(set,1), params(set,2), params(set,3));
        end
    end
    err = sum((yVals - fitted_vals).^2);
end

%% ===================================================================================
%  ===== FOOOF STATS =================================================================
%  ===================================================================================
function [ePeaks, eAper, eStats] = FOOOF_analysis(FOOOF_data, ChanNames, TF, max_peaks, sort_type, sort_param, sort_bands)
    % ===== EXTRACT PEAKS =====
    % Organize/extract peak components from FOOOF models
    nChan = numel(ChanNames);
    maxEnt = nChan * max_peaks;
    switch sort_type
        case 'param'
            % Initialize output struct
            ePeaks = struct('channel', [], 'center_frequency', [],...
                'amplitude', [], 'std_dev', []);
            % Collect data from all peaks
            i = 0;
            for chan = 1:nChan
                if ~isempty(FOOOF_data(chan).peak_params)
                    for p = 1:size(FOOOF_data(chan).peak_params,1)
                        i = i +1;
                        ePeaks(i).channel = ChanNames(chan);
                        ePeaks(i).center_frequency = FOOOF_data(chan).peak_params(p,1);
                        ePeaks(i).amplitude = FOOOF_data(chan).peak_params(p,2);
                        ePeaks(i).std_dev = FOOOF_data(chan).peak_params(p,3);
                    end
                end
            end
            % Apply specified sort
            switch sort_param
                case 'frequency'
                    [tmp,iSort] = sort([ePeaks.center_frequency]);
                    ePeaks = ePeaks(iSort);
                case 'amplitude'
                    [tmp,iSort] = sort([ePeaks.amplitude]);
                    ePeaks = ePeaks(iSort(end:-1:1));
                case 'std'
                    [tmp,iSort] = sort([ePeaks.std_dev]);
                    ePeaks = ePeaks(iSort);
            end
        case 'band'
            % Initialize output struct
            ePeaks = struct('channel', [], 'center_frequency', [],...
                'amplitude', [], 'std_dev', [], 'band', []);
            % Generate bands from input
            bands = process_tf_bands('Eval', sort_bands);
            % Collect data from all peaks
            i = 0;
            for chan = 1:nChan
                if ~isempty(FOOOF_data(chan).peak_params)
                    for p = 1:size(FOOOF_data(chan).peak_params,1)
                        i = i +1;
                        ePeaks(i).channel = ChanNames(chan);
                        ePeaks(i).center_frequency = FOOOF_data(chan).peak_params(p,1);
                        ePeaks(i).amplitude = FOOOF_data(chan).peak_params(p,2);
                        ePeaks(i).std_dev = FOOOF_data(chan).peak_params(p,3);
                        % Find name of frequency band from user definitions
                        bandRanges = cell2mat(bands(:,2));
                        iBand = find(ePeaks(i).center_frequency >= bandRanges(:,1) & ePeaks(i).center_frequency <= bandRanges(:,2));
                        if ~isempty(iBand)
                            ePeaks(i).band = bands{iBand,1};
                        else
                            ePeaks(i).band = 'None';
                        end
                    end
                end
            end
    end

    % ===== EXTRACT APERIODIC =====
    % Organize/extract aperiodic components from FOOOF models
    hasKnee = length(FOOOF_data(1).aperiodic_params) - 2;
    % Initialize output struct
    eAper = struct('channel', [], 'offset', [], 'exponent', []);
    for chan = 1:nChan
            eAper(chan).channel = ChanNames(chan);
            eAper(chan).offset = FOOOF_data(chan).aperiodic_params(1);
        if hasKnee % Legacy FOOOF alters order of parameters
            eAper(chan).exponent = FOOOF_data(chan).aperiodic_params(3);
            eAper(chan).knee_frequency = FOOOF_data(chan).aperiodic_params(2);
        else
            eAper(chan).exponent = FOOOF_data(chan).aperiodic_params(2);
        end
    end

    % ===== EXTRACT STAT =====
    % Organize/extract stats from FOOOF models
    % Initialize output struct
    eStats = struct('channel', ChanNames);
    for chan = 1:nChan
        eStats(chan).MSE = FOOOF_data(chan).error;
        eStats(chan).r_squared = FOOOF_data(chan).r_squared;
        spec = squeeze(log10(TF(chan,:,:)));
        fspec = squeeze(log10(FOOOF_data(chan).fooofed_spectrum))';
        eStats(chan).frequency_wise_error = abs(spec-fspec);
    end
end


% Snake AI: Dynamic Hamiltonian Cycle Repair (with some strategic stuff)
% by Brian Haidet for AlphaPhoenix
% published 4/11/2020
% CC non-commercial, attribution
% See readme for details

%there are two methods that this file uses to reconstruct the hamiltonian
%cycle. first, it checks to see if the turn the snake WANTS to take splits
%the hamcycle into two closed loops. if so, it looks for a location where
%both loops have a flat edge to each other and splices the loops together.

%if that fails, there's a special method for a "loop of two" where a single
%line of two nodes was isolated and needs to be re-fused to tha path.

%if that fails, then there's a more expensive algorythm that attempts to
%draw a brand new hamiltionian cycle by expanding the current path into
%unclaimed regions. It's not vry smart, but it does ocasionally succeed
%(but normally leaves little one-off missing nodes everywhere and doesn't
%complete.)

%if none of these methods solve for a new hamiltionian path, the snake
%follows the old hamiltionian path.

%I'm very sure that somewhere in this file there are some missing minus
%signs or some flipped indeces to matrices or vectors because i sometimes
%see the snake fail to repair the path in a way that I believe it should
%know how, but as you can see, there are an awful lot of IFs and (-)s in
%this file, and I lacked the patience to find every glitch when working on
%this initially. Apparently whatever sometimes fails simply causes this
%file to return no hamiltonian path instead of a wrong or illegal
%hamiltonian path, and that was good enough at the time. if you find
%anything, god you must be bored, but if you do, let me know!!  :)

hypsnake=[nextstepinds,snake];
newcycle=[]; % populated if sucessful

% map cycle #1 (the one with the snake)
oldcutoffpointer_pathind=find(hampathinds==nextstepinds);
cycle1=hampathinds(oldcutoffpointer_pathind:end);
%map cycle #2 (without the snake)
cycle2=hampathinds(1:oldcutoffpointer_pathind-1);
if mod(length(cycle2),2)==1
    disp("it's broken")
end
%does cycle2 close? (AND isn't a 2-node line!)
if (sum(abs(coords(cycle2(1))-coords(cycle2(end))))==1) && length(cycle2)>2 %yes it closes
    cycle1fill=zeros(l);
    cycle1fill(cycle1)=1;
    hypfield=zeros(l);
    hypfield(hypsnake)=1;
    hashperim=10*(cycle1fill-hypfield)+~cycle1fill;
    [hitx,hity]=find(conv2(hashperim,[1,1;1,1])==22);
    for c=1:length(hitx)
        %find out if the adjacent parts of the paths are "flat" ie. connectable
        %for each (manually)
        %find candidate coords
        candidatesite=[hitx(c),hity(c)]+[-1,-1;0,-1;-1,0;0,0];
        candidatesiteinds=inds(candidatesite);
        candidatechains=[0,0,0,0];%candidates belong to cycle1 (0) or cycle2 (1)
        %where does each point?
        candidatepointers=[0,0,0,0];
        cyloc=[0;0;0;0];
        for p=1:4
            if cycle1fill(inds(candidatesite(p,:)))==1
                cyloc(p)=find(cycle1==candidatesiteinds(p,:));
                candidatepointers(p)=cycle1(mod(cyloc(p),length(cycle1))+1);
            else
                candidatechains(p)=1;
                cyloc(p)=find(cycle2==candidatesiteinds(p,:));
                candidatepointers(p)=cycle2(mod(cyloc(p),length(cycle2))+1);
            end
        end
        matchers=intersect(candidatesiteinds,candidatepointers);
        if length(matchers)==2%they match!

            if sum((matchers(1)==candidatesiteinds)&candidatechains') % if matchers(2) belongs to cycle 1
                cy2ind=sum((matchers(1)==candidatesiteinds).*cyloc);
                cy1ind=sum((matchers(2)==candidatesiteinds).*cyloc);
            else
                cy1ind=sum((matchers(1)==candidatesiteinds).*cyloc);
                cy2ind=sum((matchers(2)==candidatesiteinds).*cyloc);
            end
            newcycle=[cycle1(mod((1:cy1ind-1)-1,length(cycle1))+1);cycle2(mod(cy2ind-1:cy2ind+length(cycle2)-2,length(cycle2))+1);cycle1(cy1ind:end)];
            break;
        end
    end
else%no it does not close
    cycle1save=cycle1;
    cycle1=[repmat([-1],l,1);cycle1(1:end-length(snake));repmat([-1],l,1)];
    if length(cycle2) <=900 %try to fix the thing
        makingprogress=false;
        madeprogress=false;
        spliced=false;
        pi=1;%segment being spliced
        while true
            spliced=false;
            if abs(cycle2(pi)-cycle2(pi+1))==1 %splicing vertical segment
                pc1=find(cycle1==cycle2(pi)-l);%look to the left of pi
                if length(pc1)==1% if that spot IS on chain 1
                    if cycle1(pc1+1)==cycle2(pi+1)-l
                        newcycle=[cycle1(1:pc1) ; cycle2(pi) ; cycle2(pi+1) ; cycle1(pc1+1:end)];
                        madeprogress=true;
                        spliced=true;
                    elseif cycle1(pc1-1)==cycle2(pi+1)-l
                        newcycle=[cycle1(1:pc1-1) ; cycle2(pi+1) ; cycle2(pi) ; cycle1(pc1:end)];
                        madeprogress=true;
                        spliced=true;
                    end
                end
                if spliced==false %look to the right if the left isnt on chain 1
                    pc1=find(cycle1==cycle2(pi)+l);%look to the right of pi
                    if length(pc1)==1 % if that spot IS on chain 1
                        if cycle1(pc1+1)==cycle2(pi+1)+l
                            newcycle=[cycle1(1:pc1) ; cycle2(pi) ; cycle2(pi+1) ; cycle1(pc1+1:end)];
                            madeprogress=true;
                            spliced=true;
                        elseif cycle1(pc1-1)==cycle2(pi+1)+l
                            newcycle=[cycle1(1:pc1-1) ; cycle2(pi+1) ; cycle2(pi) ; cycle1(pc1:end)];
                            madeprogress=true;
                            spliced=true;
                        end
                    end
                end
            else %splicing horizontal segment
                if ~(mod(cycle2(pi),l)==1) %if on top row, don't look up
                    pc1=find(cycle1==cycle2(pi)-1);%look to the up of pi
                    if length(pc1)==1 % if that spot IS on chain 1
                        if cycle1(pc1+1)==cycle2(pi+1)-1
                            newcycle=[cycle1(1:pc1) ; cycle2(pi) ; cycle2(pi+1) ; cycle1(pc1+1:end)];
                            madeprogress=true;
                            spliced=true;
                        elseif cycle1(pc1-1)==cycle2(pi+1)-1
                            newcycle=[cycle1(1:pc1-1) ; cycle2(pi+1) ; cycle2(pi) ; cycle1(pc1:end)];
                            madeprogress=true;
                            spliced=true;
                        end
                    end
                end
                if spliced==false%look to the down if the up isnt on chain 1
                    if ~(mod(cycle2(pi),l)==0) %if on bottom row, don't look down
                        pc1=find(cycle1==cycle2(pi)+1);%look to the down of pi
                        if length(pc1)==1 % if that spot IS on chain 1
                            if cycle1(pc1+1)==cycle2(pi+1)+1
                                newcycle=[cycle1(1:pc1) ; cycle2(pi) ; cycle2(pi+1) ; cycle1(pc1+1:end)];
                                madeprogress=true;
                                spliced=true;
                            elseif cycle1(pc1-1)==cycle2(pi+1)+1
                                newcycle=[cycle1(1:pc1-1) ; cycle2(pi+1) ; cycle2(pi) ; cycle1(pc1:end)];
                                madeprogress=true;
                                spliced=true;
                            end
                        end
                    end
                end
            end
            if spliced==true
                cycle2(pi:pi+1)=[];
                cycle1=newcycle;
            else
                pi=pi+2;
            end
            if isempty(cycle2)
%                 disp('stage1')
                break;
            end
            if pi>length(cycle2)
                if madeprogress
                    pi=1;
                    madeprogress=false;
%                     %%plotting for debug
%                     figure(3)
%                     plot(snakecoords(:,2),snakecoords(:,1),'Color',[0 1 0],'LineWidth',15)
%                     xlim([0.5,l+.5]);
%                     ylim([0.5,l+.5]);
%                     axis square
%                     set(gca,'Ydir','reverse')
%                     set(gca,'Color','k')
%                     hold on
%                     plot(optimalPath(:,2),optimalPath(:,1),'Color',[1 0 0],'LineWidth',5)
%                     ncc2=coords(cycle1);
%                     plot(ncc2(:,2),ncc2(:,1),'Color',[0 .3 0],'LineWidth',3)
%                     hold off
                else
%                     disp('stage1 - FAIL')
                    break;
                end
            end

        end
        if ~isempty(cycle2)
            makingprogress=false;
            madeprogress=false;
            spliced=false;
            pi=l+1;%segment being spliced
            while true
                spliced=false;
                if abs(cycle1(pi)-cycle1(pi+1))==1 %splicing vertical segment
                    pc2a=find(cycle2==cycle1(pi)-l);%look to the left of pi
                    pc2b=find(cycle2==cycle1(pi+1)-l);%look to the left of pi+1
                    if (length(pc2a)==1)&&(length(pc2b)==1)% if those spot ARE on chain 2
                        newcycle=[cycle1(1:pi) ; cycle2(pc2a) ; cycle2(pc2b) ; cycle1(pi+1:end)];
                        madeprogress=true;
                        spliced=true;
                    end
                    if spliced==false %look to the right if the left isnt on chain 1
                        pc2a=find(cycle2==cycle1(pi)+l);%look to the right of pi
                        pc2b=find(cycle2==cycle1(pi+1)+l);%look to the right of pi+1
                        if (length(pc2a)==1)&&(length(pc2b)==1)% if those spot ARE on chain 2
                            newcycle=[cycle1(1:pi) ; cycle2(pc2a) ; cycle2(pc2b) ; cycle1(pi+1:end)];
                            madeprogress=true;
                            spliced=true;
                        end
                    end
                else %splicing horizontal segment
                    if ~(mod(cycle1(pi),l)==1) %if on top row, don't look up
                        pc2a=find(cycle2==cycle1(pi)-1);%look to the up of pi
                        pc2b=find(cycle2==cycle1(pi+1)-1);%look to the up of pi+1
                        if (length(pc2a)==1)&&(length(pc2b)==1)% if those spot ARE on chain 2
                            newcycle=[cycle1(1:pi) ; cycle2(pc2a) ; cycle2(pc2b) ; cycle1(pi+1:end)];
                            madeprogress=true;
                            spliced=true;
                        end
                    end
                    if spliced==false%look to the down if the up isnt on chain 1
                        if ~(mod(cycle1(pi),l)==0) %if on bottom row, don't look down
                            pc2a=find(cycle2==cycle1(pi)+1);%look to the down of pi
                            pc2b=find(cycle2==cycle1(pi+1)+1);%look to the down of pi+1
                            if (length(pc2a)==1)&&(length(pc2b)==1)% if those spot ARE on chain 2
                                newcycle=[cycle1(1:pi) ; cycle2(pc2a) ; cycle2(pc2b) ; cycle1(pi+1:end)];
                                madeprogress=true;
                                spliced=true;
                            end
                        end
                    end
                end
                if spliced==true
                    cycle2([pc2a,pc2b])=[];
                    cycle1=newcycle;
                else
                    pi=pi+1;
                end
                if isempty(cycle2)
%                     disp('stage2')
                    break;
                end
                if pi>length(cycle1)-2*l
                    if madeprogress
                        pi=l+1;
                        madeprogress=false;
%                         %%plotting for debug
%                         figure(3)
%                         plot(snakecoords(:,2),snakecoords(:,1),'Color',[0 1 0],'LineWidth',15)
%                         xlim([0.5,l+.5]);
%                         ylim([0.5,l+.5]);
%                         axis square
%                         set(gca,'Ydir','reverse')
%                         set(gca,'Color','k')
%                         hold on
%                         plot(optimalPath(:,2),optimalPath(:,1),'Color',[1 0 0],'LineWidth',5)
%                         ncc2=coords(cycle1);
%                         plot(ncc2(:,2),ncc2(:,1),'Color',[0 .3 .7],'LineWidth',3)
%                         hold off
                    else
%                         disp('stage2 - FAIL')
                        break;
                    end
                end

            end
        end
        if madeprogress
            newcycle=[newcycle(l+1:end-l); cycle1save(end-length(snake)+1:end)];
        else
            newcycle=[];
        end
    end
end


function [locate_voi_min_x,locate_voi_max_x,locate_voi_min_y,locate_voi_max_y,locate_voi_min_z,locate_voi_max_z] = dicomrt_voiboundaries(dose_xmesh,dose_ymesh,dose_zmesh,VOI,voi2use,PatientPosition)
% dicomrt_voiboundaries(dose_xmesh,dose_ymesh,dose_zmesh,VOI,voi2use,PatientPosition)
%
% Returns the voxel index of the smallest and largest coordinate of the selected VOI in the specified Z slice.
%
% slice is the slice number where the mins and maxs will be calculated.
%      Slice number is relative to voiref.
% voiref is the voi of reference
% dose_xmesh,dose_ymesh, are coordinates of the center of the dose-pixel
% VOI is a cell array which contain the patients VOIs as read by dicomrt_loadvoi
% voi2use is a vector pointing to the number of VOIs to be used
%
% Example:
%
% [xmin,xmax,ymin,ymax]=dicomrt_voiboundariesZ(12,1,xmesh,xmesh,VOI,5)
%
% returns in xmin,xmax,ymin,ymax the pixel numbers which correspond to the edges of
% structure "5" (e.g. target) in slice number 12 with reference to structure number "1".
%
% See also dicomrt_voiboundaries, dicomrt_surfdose
%
% Copyright (C) 2002 Emiliano Spezi (emiliano.spezi@physics.org)

% Check input
[VOI_temp]=dicomrt_checkinput(VOI);
VOI=dicomrt_varfilter(VOI_temp);

nslices=size(VOI{voi2use,2},1);

voiref=1;

for kk=1:nslices
    [locate_voi_min_1,locate_voi_max_1,locate_voi_min_2,locate_voi_max_2]=dicomrt_voiboundaries_single(kk,dose_xmesh,dose_ymesh,VOI_temp,voi2use,PatientPosition);
    if kk==1
        locate_voi_min_x=locate_voi_min_1;
        locate_voi_max_x=locate_voi_max_1;
        locate_voi_min_y=locate_voi_min_2;
        locate_voi_max_y=locate_voi_max_2;
    else
        if locate_voi_min_1 < locate_voi_min_x
            locate_voi_min_x=locate_voi_min_1;
        end
        if locate_voi_max_1 > locate_voi_max_x
            locate_voi_max_x=locate_voi_max_1;
        end
        if locate_voi_min_2 < locate_voi_min_y
            locate_voi_min_y=locate_voi_min_2;
        end
        if locate_voi_max_2 > locate_voi_max_y
            locate_voi_max_y=locate_voi_max_2;
        end
    end
end

% z mesh is always sorted
min_z_voi=VOI_temp{2,1}{voi2use,2}{1}(1,3);
max_z_voi=VOI_temp{2,1}{voi2use,2}{end}(1,3);

locate_voi_min_z=dicomrt_findpointVECT(dose_zmesh,min_z_voi);
locate_voi_max_z=dicomrt_findpointVECT(dose_zmesh,max_z_voi);


% ==============================================================================
%
%    Software License Agreement (BSD License)
%    Copyright (c) 2019
%    (www.aimlab.wpi.edu)
%
%    All rights reserved.
%
%    Redistribution and use in source and binary forms, with or without
%    modification, are permitted provided that the following conditions
%    are met:
%
%    * Redistributions of source code must retain the above copyright
%    notice, this list of conditions and the following disclaimer.
%
%    * Redistributions in binary form must reproduce the above
%    copyright notice, this list of conditions and the following
%    disclaimer in the documentation and/or other materials provided
%    with the distribution.
%
%    * Neither the name of authors nor the names of its contributors may
%    be used to endorse or promote products derived from this software
%    without specific prior written permission.
%
%    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
%    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
%    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
%    FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
%    COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
%    INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
%    BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
%    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
%    CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
%    LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
%    ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
%    POSSIBILITY OF SUCH DAMAGE.
%
%    \author:    <http://www.aimlab.wpi.edu>
%    \author:    <amunawar@wpi.edu>
%    \author:    Adnan Munawar
%    \version:   0.1$
% ==============================================================================

function [nIterations,sizePath,run_time] =  RRTconnect3D(dim,segmentLength,random_world,show_output)
% dim = 2;
% segmentLength = 5;
% random_world = 0;
% standard length of path segments
if dim ==2
    start_cord = [5,5];
    goal_cord = [95,95];

else

    start_cord = [5,5,5];
    goal_cord = [95,95,95];
end

% create random world
Size = 100;
NumObstacles = 100;

if random_world ==1
    world = createWorld(NumObstacles,ones(1,dim)*Size,zeros(1,dim),dim);
else
    [world NumObstacles] = createKnownWorld(ones(1,dim)*Size,[0;0;0],dim);
end
% randomly select start and end nodes
%start_node = generateRandomNode(world,dim)
%end_node   = generateRandomNode(world,dim)
start_node = [start_cord,0,0,0];
end_node = [goal_cord,0,0,0];
% establish tree starting with the start node
tree = start_node;

a = clock;

% check to see if start_node connects directly to end_node
if ( (norm(start_node(1:dim)-end_node(1:dim))<segmentLength )...
        &&(collision(start_node,end_node,world,dim)==0) )
    path = [start_node; end_node];
else
    nIterations = 0;
    numPaths = 0;
    flag = 0;
    while numPaths<1,
        [tree,flag] = extendTree(tree,end_node,segmentLength,world,dim);
        numPaths = numPaths + flag;
        nIterations = nIterations+1;
    end
end

path = findMinimumPath(tree,end_node,dim);
sizePath = size(path,1);

b = clock;

run_time = 3600*(b(4)-a(4)) + 60 * (b(5)-a(5)) + (b(6) - a(6));

if show_output == 1
    % find path with minimum cost to end_node
    figure;
    plotExpandedTree(world,tree,dim);
    plotWorld(world,path,dim);
    % figure(2);
    % plotWorld(world,path,dim);
    %plotExpandedTree(world,tree,dim);
end
end

function world = createWorld(NumObstacles, endcorner, origincorner,dim)

if dim == 2

    % check to make sure that the region is nonempty
    if (endcorner(1) <= origincorner(1)) | (endcorner(2) <= origincorner(2))
        disp('Not valid corner specifications!')
        world=[];

        % create world data structure
    else
        world.NumObstacles = NumObstacles;
        world.endcorner = endcorner;
        world.origincorner = origincorner;

        % create NumObstacles
        maxRadius = min(endcorner(1)- origincorner(1), endcorner(2)-origincorner(2));
        maxRadius = 5*maxRadius/NumObstacles/2;
        for i=1:NumObstacles,
            % randomly pick radius
            world.radius(i) = maxRadius*rand;
            % randomly pick center of obstacles
            cx = origincorner(1) + world.radius(i)...
                + (endcorner(1)-origincorner(1)-2*world.radius(i))*rand;
            cy = origincorner(2) + world.radius(i)...
                + (endcorner(2)-origincorner(2)-2*world.radius(i))*rand;
            world.cx(i) = cx;
            world.cy(i) = cy;
        end
    end

elseif dim ==3;
    % check to make sure that the region is nonempty
    if (endcorner(1) <= origincorner(1)) || (endcorner(2) <= origincorner(2)) || (endcorner(3) <= origincorner(3))
        disp('Not valid corner specifications!')
        world=[];

        % create world data structure
    else
        world.NumObstacles = NumObstacles;
        world.endcorner = endcorner;
        world.origincorner = origincorner;

        % create NumObstacles
        bounds = [endcorner(1)- origincorner(1), endcorner(2)-origincorner(2), endcorner(3)-origincorner(3)];
        maxRadius = min(bounds);
        maxRadius = 5*maxRadius/NumObstacles;
        for i=1:NumObstacles,
            % randomly pick radius
            world.radius(i) = maxRadius*rand;
            % randomly pick center of obstacles
            cx = origincorner(1) + world.radius(i)...
                + (endcorner(1)-origincorner(1)-2*world.radius(i))*rand;
            cy = origincorner(2) + world.radius(i)...
                + (endcorner(2)-origincorner(2)-2*world.radius(i))*rand;
            cz = origincorner(2) + world.radius(i)...
                + (endcorner(2)-origincorner(2)-2*world.radius(i))*rand;
            world.cx(i) = cx;
            world.cy(i) = cy;
            world.cz(i) = cz;
        end
    end
end
end

function [world NumObstacles] = createKnownWorld(endcorner, origincorner,dim)
NumObstacles = 5;
if dim == 2
    % check to make sure that the region is nonempty
    if (endcorner(1) <= origincorner(1)) | (endcorner(2) <= origincorner(2)),
        disp('Not valid corner specifications!')
        world=[];
        % create world data structure
    else
        world.NumObstacles = NumObstacles;
        world.endcorner = endcorner;
        world.origincorner = origincorner;

        % create NumObstacles
        maxRadius = 10;

        world.radius(1) = maxRadius;
        cx = 50;
        cy = 50;
        world.cx(1) = cx;
        world.cy(1) = cy;

        world.radius(2) = maxRadius;
        cx = 75;
        cy = 25;
        world.cx(2) = cx;
        world.cy(2) = cy;

        world.radius(3) = maxRadius;
        cx = 25;
        cy = 75;
        world.cx(3) = cx;
        world.cy(3) = cy;

        world.radius(4) = maxRadius;
        cx = 25;
        cy = 25;
        world.cx(4) = cx;
        world.cy(4) = cy;

        world.radius(5) = maxRadius;
        cx = 75;
        cy = 75;
        world.cx(5) = cx;
        world.cy(5) = cy;
    end

elseif dim == 3

    NumObstacles = 9;
    % check to make sure that the region is nonempty
    if (endcorner(1) <= origincorner(1)) | (endcorner(2) <= origincorner(2)) | (endcorner(3) <= origincorner(3)),
        disp('Not valid corner specifications!')
        world=[];

        % create world data structure
    else
        world.NumObstacles = NumObstacles;
        world.endcorner = endcorner;
        world.origincorner = origincorner;

        % create NumObstacles
        maxRadius = 10;

        world.radius(1) = maxRadius;
        cx = 50;
        cy = 50;
        cz = 50;
        world.cx(1) = cx;
        world.cy(1) = cy;
        world.cz(1) = cz;

        world.radius(2) = maxRadius;
        cx = 25;
        cy = 25;
        cz = 25;
        world.cx(2) = cx;
        world.cy(2) = cy;
        world.cz(2) = cz;

        world.radius(3) = maxRadius;
        cx = 75;
        cy = 75;
        cz = 75;
        world.cx(3) = cx;
        world.cy(3) = cy;
        world.cz(3) = cz;

        world.radius(4) = maxRadius;
        cx = 25;
        cy = 25;
        cz = 75;
        world.cx(4) = cx;
        world.cy(4) = cy;
        world.cz(4) = cz;

        world.radius(5) = maxRadius;
        cx = 75;
        cy = 75;
        cz = 25;
        world.cx(5) = cx;
        world.cy(5) = cy;
        world.cz(5) = cz;

        world.radius(6) = maxRadius;
        cx = 25;
        cy = 75;
        cz = 25;
        world.cx(6) = cx;
        world.cy(6) = cy;
        world.cz(6) = cz;

        world.radius(7) = maxRadius;
        cx = 75;
        cy = 25;
        cz = 25;
        world.cx(7) = cx;
        world.cy(7) = cy;
        world.cz(7) = cz;

        world.radius(8) = maxRadius;
        cx = 75;
        cy = 25;
        cz = 75;
        world.cx(8) = cx;
        world.cy(8) = cy;
        world.cz(8) = cz;

        world.radius(9) = maxRadius;
        cx = 25;
        cy = 75;
        cz = 75;
        world.cx(9) = cx;
        world.cy(9) = cy;
        world.cz(9) = cz;
    end
end
end

function node=generateRandomNode(world,dim)

if dim ==2;
    % randomly pick configuration
    px       = (world.endcorner(1)-world.origincorner(1))*rand;
    py       = (world.endcorner(2)-world.origincorner(2))*rand;

    chi      = 0;
    cost     = 0;
    node     = [px, py, chi, cost, 0];

    % check collision with obstacle
    while collision(node, node, world,dim),
        px       = (world.endcorner(1)-world.origincorner(1))*rand;
        py       = (world.endcorner(2)-world.origincorner(2))*rand;

        chi      = 0;
        cost     = 0;
        node     = [px, py, chi, cost, 0];
    end

elseif dim ==3;
    % randomly pick configuration
    px       = (world.endcorner(1)-world.origincorner(1))*rand;
    py       = (world.endcorner(2)-world.origincorner(2))*rand;
    pz       = (world.endcorner(3)-world.origincorner(3))*rand;

    chi      = 0;
    cost     = 0;
    node     = [px, py, pz, chi, cost, 0];

    % check collision with obstacle
    while collision(node, node, world,dim),
        px       = (world.endcorner(1)-world.origincorner(1))*rand;
        py       = (world.endcorner(2)-world.origincorner(2))*rand;
        pz       = (world.endcorner(3)-world.origincorner(3))*rand;

        chi      = 0;
        cost     = 0;
        node     = [px, py, pz, chi, cost, 0];
    end

end

end

function collision_flag = collision(node, parent, world,dim)

collision_flag = 0;

for i=1:dim
    if (node(i)>world.endcorner(i))|(node(i)<world.origincorner(i))
        collision_flag = 1;
    end
end

if collision_flag == 0 && dim ==2
    for sigma = 0:.2:1,
        p = sigma*node(1:dim) + (1-sigma)*parent(1:dim);
        % check each obstacle
        for i=1:world.NumObstacles,
            if (norm([p(1);p(2)]-[world.cx(i); world.cy(i)])<=1*world.radius(i)),
                collision_flag = 1;
                break;
            end
        end
    end

elseif collision_flag == 0 && dim ==3
    for sigma = 0:.2:1,
        p = sigma*node(1:dim) + (1-sigma)*parent(1:dim);
        % check each obstacle
        for i=1:world.NumObstacles,
            if (norm([p(1);p(2);p(3)]-[world.cx(i); world.cy(i); world.cz(i)])<=1*world.radius(i)),
                collision_flag = 1;
                break;
            end
        end
    end
end
end

function collision_flag = is_point_valid(point, world,dim)

collision_flag = 0;

for i=1:dim
    if (point(i)>world.endcorner(i))||(point(i)<world.origincorner(i))
        collision_flag = 1;
    end
end

if collision_flag == 0 && dim ==2

    p = point(1:dim);
    % check each obstacle
    for i=1:world.NumObstacles,
        if (norm([p(1);p(2)]-[world.cx(i); world.cy(i)])<=1*world.radius(i)),
            collision_flag = 1;
            break;
        end
    end

elseif collision_flag == 0 && dim ==3
    p = point(1:dim);
    % check each obstacle
    for i=1:world.NumObstacles,
        if (norm([p(1);p(2);p(3)]-[world.cx(i); world.cy(i); world.cz(i)])<=1*world.radius(i)),
            collision_flag = 1;
            break;
        end
    end
end
end

function flag = canEndConnectToTree(tree,end_node,minDist,world,dim)
flag = 0;
% check only last node added to tree since others have been checked
if ( (norm(tree(end,1:dim)-end_node(1:dim))<minDist)...
        & (collision(tree(end,1:dim), end_node(1:dim), world,dim)==0) ),
    flag = 1;
end

end

function [new_tree,flag] = extendTree(tree,end_node,segmentLength,world,dim)

flag = 0;
% select a random point
randomPoint = zeros(1,dim);
for i=1:dim
    randomPoint(1,i) = (world.endcorner(i)-world.origincorner(i))*rand;
end

% find leaf on node that is closest to randomPoint
tmp = tree(:,1:dim)-ones(size(tree,1),1)*randomPoint;
sqrd_dist = sqr_eucl_dist(tmp,dim);
[min_dist,idx] = min(sqrd_dist);

min_parent_idx = idx;

new_point = tree(idx,1:dim);
new_node = tree(idx,:);

pflag = 0;

while norm(new_point-randomPoint)>0 && pflag==0

    if norm(new_point-randomPoint)<segmentLength

        pflag = collision(randomPoint,tree(min_parent_idx,:),world,dim);
        if pflag == 0

            new_point = randomPoint;
            min_cost = cost_np(tree(min_parent_idx,:),new_point,dim);
            new_node = [new_point,0,min_cost,min_parent_idx];
            tree = [tree;new_node];
            pflag = 1;
            goal_flag = is_goal(new_node,end_node,segmentLength,world,dim);

            if goal_flag == 1;
                tree(end,dim+1)=1;
                flag = 1;
            end
        end

    else

        new_point = (randomPoint-tree(min_parent_idx,1:dim));
        new_point = tree(min_parent_idx,1:dim)+(new_point/norm(new_point))*segmentLength;

        min_cost  = cost_np(tree(min_parent_idx,:),new_point,dim);
        new_node  = [new_point, 0, min_cost, min_parent_idx];

        pflag = collision(new_node,tree(min_parent_idx,:),world,dim);

        if pflag == 0
            tree = [tree ; new_node];
            min_parent_idx = size(tree,1);

            goal_flag = is_goal(new_node,end_node,segmentLength,world,dim);

            if goal_flag == 1;
                tree(end,dim+1)=1;  % mark node as connecting to end.
                pflag = 1;
                flag = 1;
            end
        end
    end
end

new_tree = tree;

end

function goal_flag = is_goal(node,end_node,segmentLength,world,dim)

goal_flag = 0;
if (norm(node(1:dim)-end_node(1:dim))<segmentLength )...
        && (collision(node,end_node,world,dim)==0)
    goal_flag = 1;
end

end

function e_dist = sqr_eucl_dist(array,dim)

sqr_e_dist = zeros(size(array,1),dim);
for i=1:dim

    sqr_e_dist(:,i) = array(:,i).*array(:,i);

end
e_dist = zeros(size(array,1),1);
for i=1:dim

    e_dist = e_dist+sqr_e_dist(:,i);

end

end

%calculate the cost from a node to a point
function [cost] = cost_np(from_node,to_point,dim)

diff = from_node(:,1:dim) - to_point;
eucl_dist = norm(diff);
cost = from_node(:,dim+2) + eucl_dist;

end

%calculate the cost from a node to a node
function [cost] = cost_nn(from_node,to_node,dim)

diff = from_node(:,1:dim) - to_node(:,1:dim);
eucl_dist = norm(diff);
cost = from_node(:,dim+2) + eucl_dist;

end

function [cost] = line_cost(from_node,to_point,dim)
diff = from_node(:,1:dim) - to_point;
cost = norm(diff);
end

function path = findMinimumPath(tree,end_node,dim)

% find nodes that connect to end_node
connectingNodes = [];
for i=1:size(tree,1),
    if tree(i,dim+1)==1,
        connectingNodes = [connectingNodes ; tree(i,:)];
    end
end

% find minimum cost last node
[tmp,idx] = min(connectingNodes(:,dim+2));

% construct lowest cost path
path = [connectingNodes(idx,:); end_node];
parent_node = connectingNodes(idx,dim+3);
while parent_node>1,
    parent_node = tree(parent_node,dim+3);
    path = [tree(parent_node,:); path];
end

end

function plotExpandedTree(world,tree,dim)
ind = size(tree,1);
while ind>0
    size(tree);
    branch = [];
    node = tree(ind,:);
    branch = [ branch ; node ];
    parent_node = node(dim+3);
    while parent_node > 1
        cur_parent = parent_node;
        branch = [branch; tree(parent_node,:)];
        parent_node = tree(parent_node,dim+3);
    end
    ind = ind - 1;

    if dim == 2
        X = branch(:,1);
        Y = branch(:,2);

        p = plot(X,Y);
        set(p,'Color','r','LineWidth',0.5,'Marker','.','MarkerEdgeColor','g');
        hold on;

    elseif dim == 3
        X = branch(:,1);
        Y = branch(:,2);
        Z = branch(:,3);

        p = plot3(X,Y,Z);
        set(p,'Color','r','LineWidth',0.5,'Marker','.','MarkerEdgeColor','g');
        hold on;
    end
end
end

function plotWorld(world,path,dim)
% the first element is the north coordinate
% the second element is the south coordinate
if dim ==2

    N = 10;
    th = 0:2*pi/N:2*pi;
    axis([world.origincorner(1),world.endcorner(1),...
        world.origincorner(2), world.endcorner(2)]);
    hold on

    for i=1:world.NumObstacles,
        X = world.radius(i)*sin(th) + world.cx(i);
        Y = world.radius(i)*cos(th) + world.cy(i);
        fill(X,Y,'blue');
    end

    X = path(:,1);
    Y = path(:,2);
    p = plot(X,Y);

elseif dim ==3
    axis([world.origincorner(1),world.endcorner(1),...
        world.origincorner(2), world.endcorner(2),...
        world.origincorner(3), world.endcorner(3)]);
    hold on

    for i=1:world.NumObstacles,
        [X Y Z] = sphere(10);
        X = (X*world.radius(i));
        Y = (Y*world.radius(i));
        Z = (Z*world.radius(i));
        surf(X+world.cx(i),Y+world.cy(i),Z+world.cz(i));
        colormap([0.5 0.2 0.3]);
    end

    X = path(:,1);
    Y = path(:,2);
    Z = path(:,3);
    p = plot3(X,Y,Z);
end
set(p,'Color','black','LineWidth',3)
xlabel('X axis');
ylabel('Y axis');
zlabel('Z axis');
title('RRT Connect Algorithm');
end


%%**********************************************************************
%% blktrace: compute <X1,Z1> + ... + <Xp,Zp>
%%
%% SDPT3: version 3.1
%% Copyright (c) 1997 by
%% K.C. Toh, M.J. Todd, R.H. Tutuncu
%% Last Modified: 16 Sep 2004
%%**********************************************************************

  function trXZ = blktrace(blk,X,Z,parbarrier);

  if (nargin == 3)
     trXZ = 0;
     for p = 1:size(blk,1)
        pblk = blk(p,:);
        if strcmp(pblk{1},'s')
           if (length(pblk{2}) == 1)
              trXZ = trXZ + sum(sum(X{p}.*Z{p}));
           else
              xx = mexsvec(pblk,X{p},0);
              zz = mexsvec(pblk,Z{p});
              trXZ = trXZ + xx'*zz;
           end
        else
           trXZ = trXZ + sum(X{p}.*Z{p});
        end
     end
  elseif (nargin == 4)
     trXZ = 0;
     for p = 1:size(blk,1)
        pblk = blk(p,:);
        if (norm(parbarrier{p}) == 0)
           if strcmp(pblk{1},'s')
              if (length(pblk{2}) == 1)
                 trXZ = trXZ + sum(sum(X{p}.*Z{p}));
              else
                 xx = mexsvec(pblk,X{p},0);
                 zz = mexsvec(pblk,Z{p});
                 trXZ = trXZ + xx'*zz;
              end
           else
              trXZ = trXZ + sum(X{p}.*Z{p});
           end
        else
           idx = find(parbarrier{p} == 0);
           if ~isempty(idx)
              if strcmp(pblk{1},'s')
                 sumXZ = sum(X{p}.*Z{p});
                 ss = [0,cumsum(pblk{2})];
                 for k = 1:length(idx)
                    idxtmp = [ss(idx(k))+1:ss(idx(k)+1)];
                    trXZ = trXZ + sum(sumXZ(idxtmp));
                 end
              elseif strcmp(pblk{1},'q')
	         tmp = qops(pblk,X{p},Z{p},1);
                 trXZ = trXZ + sum(tmp(idx));
              elseif strcmp(pblk{1},'l')
                 trXZ = trXZ + sum(X{p}(idx).*Z{p}(idx));
              end
           end
        end
     end
  end
%%**********************************************************************