function [source] = ft_dipolefitting(cfg, data)

% FT_DIPOLEFITTING perform grid search and non-linear fit with one or multiple
% dipoles and try to find the location where the dipole model is best able
% to explain the measured EEG or MEG topography.
%
% This function will initially scan the whole brain with a single dipole on
% a regular coarse grid, and subsequently start at the most optimal location
% with a non-linear search. Alternatively you can specify the initial
% location of the dipole(s) and the non-linear search will start from there.
%
% Use as
%   [source] = ft_dipolefitting(cfg, data)
%
% The configuration has the following general fields
%   cfg.numdipoles  = number, default is 1
%   cfg.symmetry    = 'x', 'y' or 'z' symmetry for two dipoles, can be empty (default = [])
%   cfg.channel     = Nx1 cell-array with selection of channels (default = 'all'),
%                     see FT_CHANNELSELECTION for details
%   cfg.gridsearch  = 'yes' or 'no', perform global search for initial
%                     guess for the dipole parameters (default = 'yes')
%   cfg.nonlinear   = 'yes' or 'no', perform nonlinear search for optimal
%                     dipole parameters (default = 'yes')
%
% If you start with a grid search, the complete grid with dipole positions is
% constructed using FT_PREPARE_SOURCEMODEL. It can be specified as as a regular 3-D
% grid that is aligned with the axes of the head coordinate system using
%   cfg.xgrid               = vector (e.g. -20:1:20) or 'auto' (default = 'auto')
%   cfg.ygrid               = vector (e.g. -20:1:20) or 'auto' (default = 'auto')
%   cfg.zgrid               = vector (e.g.   0:1:20) or 'auto' (default = 'auto')
%   cfg.resolution          = number (e.g. 1 cm) for automatic grid generation
% If the source model destribes a triangulated cortical sheet, it is described as
%   cfg.sourcemodel.pos     = N*3 matrix with the vertex positions of the cortical sheet
%   cfg.sourcemodel.tri     = M*3 matrix that describes the triangles connecting the vertices
% Alternatively the position of a few dipoles at locations of interest can be
% user-specified, for example obtained from an anatomical or functional MRI
%   cfg.sourcemodel.pos     = N*3 matrix with position of each source
%   cfg.sourcemodel.inside  = N*1 vector with boolean value whether grid point is inside brain (optional)
%   cfg.sourcemodel.dim     = [Nx Ny Nz] vector with dimensions in case of 3-D grid (optional)
%
% If you do not start with a grid search, you have to give a starting location
% for the nonlinear search
%   cfg.dip.pos     = initial dipole position, matrix of Ndipoles x 3
%
% The conventional approach is to fit dipoles to event-related averages, which
% within FieldTrip can be obtained from the FT_TIMELOCKANALYSIS or from
% the FT_TIMELOCKGRANDAVERAGE function. This has the additional options
%   cfg.latency     = [begin end] in seconds or 'all' (default = 'all')
%   cfg.model       = 'moving' or 'regional'
% A moving dipole model has a different position (and orientation) for each
% timepoint, or for each component. A regional dipole model has the same
% position for each timepoint or component, and a different orientation.
%
% You can also fit dipoles to the spatial topographies of an independent
% component analysis, obtained from the FT_COMPONENTANALYSIS function.
% This has the additional options
%   cfg.component   = array with numbers (can be empty -> all)
%
% You can also fit dipoles to the spatial topographies that are present
% in the data in the frequency domain, which can be obtained using the
% FT_FREQANALYSIS function. This has the additional options
%   cfg.frequency   = single number (in Hz)
%
% Low level details of the fitting can be specified in the cfg.dipfit structure
%   cfg.dipfit.display      = level of display, can be 'off', 'iter', 'notify' or 'final' (default = 'iter')
%   cfg.dipfit.optimfun     = function to use, can be 'fminsearch' or 'fminunc' (default is determined automatic)
%   cfg.dipfit.maxiter      = maximum number of function evaluations allowed (default depends on the optimfun)
%   cfg.dipfit.checkinside  = boolean, check that the dipole remains in the source compartment (default = false)
%
% Optionally, you can modify the leadfields by reducing the rank, i.e. remove the weakest orientation
%   cfg.reducerank    = 'no', or number (default = 3 for EEG, 2 for MEG)
%   cfg.backproject   = 'yes' or 'no',  determines when reducerank is applied whether the
%                       lower rank leadfield is projected back onto the original linear
%                       subspace, or not (default = 'yes')
%
% The volume conduction model of the head should be specified as
%   cfg.headmodel     = structure with volume conduction model, see FT_PREPARE_HEADMODEL
%
% The EEG or MEG sensor positions can be present in the data or can be specified as
%   cfg.elec          = structure with electrode positions or filename, see FT_READ_SENS
%   cfg.grad          = structure with gradiometer definition or filename, see FT_READ_SENS
%
% To facilitate data-handling and distributed computing you can use
%   cfg.inputfile   =  ...
%   cfg.outputfile  =  ...
% If you specify one of these (or both) the input data will be read from a *.mat
% file on disk and/or the output data will be written to a *.mat file. These mat
% files should contain only a single variable, corresponding with the
% input/output structure.
%
% See also FT_SOURCEANALYSIS, FT_PREPARE_LEADFIELD, FT_PREPARE_HEADMODEL

% TODO change the output format, more suitable would be something like:
% dip.label
% dip.time
% dip.avg (instead of Vdata)
% dip.dip.pos
% dip.dip.mom
% dip.dip.model, or dip.dip.avg
% dip.dimord

% Undocumented local options:
%   cfg.dipfit.constr   = Source model constraints, depends on cfg.symmetry
% Optionally, you can include a noise covariance structure to sphere the data (is useful when using both
% magnetometers and gradiometers to fit your dipole)
%   cfg.dipfit.noisecov       = noise covariance matrix, see e.g. FT_TIMELOCK_ANALYSIS

% Copyright (C) 2004-2013, Robert Oostenveld
%
% This file is part of FieldTrip, see http://www.fieldtriptoolbox.org
% for the documentation and details.
%
%    FieldTrip is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    FieldTrip is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with FieldTrip. If not, see <http://www.gnu.org/licenses/>.
%
% $Id$

% these are used by the ft_preamble/ft_postamble function and scripts
ft_revision = '$Id$';
ft_nargin   = nargin;
ft_nargout  = nargout;

% do the general setup of the function
ft_defaults
ft_preamble init
ft_preamble debug
ft_preamble loadvar data
ft_preamble provenance data

% the ft_abort variable is set to true or false in ft_preamble_init
if ft_abort
  return
end

% check if the input data is valid for this function
data = ft_checkdata(data, 'datatype', {'comp', 'timelock', 'freq'}, 'feedback', 'yes');

% check if the input cfg is valid for this function
cfg = ft_checkconfig(cfg, 'forbidden',  {'channels'}); % prevent accidental typos, see issue 1729
cfg = ft_checkconfig(cfg, 'renamed',    {'elecfile', 'elec'});
cfg = ft_checkconfig(cfg, 'renamed',    {'gradfile', 'grad'});
cfg = ft_checkconfig(cfg, 'renamed',    {'optofile', 'opto'});
cfg = ft_checkconfig(cfg, 'renamed',    {'hdmfile', 'headmodel'});
cfg = ft_checkconfig(cfg, 'renamed',    {'vol',     'headmodel'});
cfg = ft_checkconfig(cfg, 'renamed',    {'grid',    'sourcemodel'});

% get the defaults
cfg.channel         = ft_getopt(cfg, 'channel', 'all');
cfg.component       = ft_getopt(cfg, 'component', 'all');   % for comp input
cfg.frequency       = ft_getopt(cfg, 'frequency');          % for freq input
cfg.latency         = ft_getopt(cfg, 'latency', 'all');     % for timelock input
cfg.feedback        = ft_getopt(cfg, 'feedback', 'text');
cfg.gridsearch      = ft_getopt(cfg, 'gridsearch', 'yes');
cfg.nonlinear       = ft_getopt(cfg, 'nonlinear', 'yes');
cfg.symmetry        = ft_getopt(cfg, 'symmetry');
cfg.dipfit          = ft_getopt(cfg, 'dipfit', []);     % the default for this is handled below

cfg = ft_checkconfig(cfg, 'renamed',    {'tightgrid', 'tight'});  % this is moved to cfg.sourcemodel.tight by the subsequent createsubcfg
cfg = ft_checkconfig(cfg, 'renamed',    {'sourceunits', 'unit'}); % this is moved to cfg.sourcemodel.unit  by the subsequent createsubcfg

% put the low-level options pertaining to the sourcemodel in their own field
cfg = ft_checkconfig(cfg, 'createsubcfg', {'sourcemodel'});
% move some fields from cfg.sourcemodel back to the top-level configuration
cfg = ft_checkconfig(cfg, 'createtopcfg', {'sourcemodel'});

% determine data type
iscomp = ft_datatype(data, 'comp');           % it can also be raw+comp, timelock+comp or freq+comp
isfreq = ft_datatype(data, 'freq');           % it might also be freq+comp, in that case it should be treated as component data
istimelock = ft_datatype(data, 'timelock');   % it might also be timelock+comp, in that case it should be treated as component data

% the default for this depends on the data type
if ~isfield(cfg, 'model')
  if iscomp
    % each component is fitted independently
    cfg.model = 'moving';
  elseif isfreq
    % fit the data with a dipole at one location
    cfg.model = 'regional';
  elseif istimelock
    % fit the data with a dipole at one location
    cfg.model = 'regional';
  end
end

if ~isfield(cfg, 'numdipoles')
  if isfield(cfg, 'dip')
    cfg.numdipoles = size(cfg.dip(1).pos,1);
  else
    cfg.numdipoles = 1;
  end
end

% set up the symmetry constraints
if ~isempty(cfg.symmetry)
  if cfg.numdipoles~=2
    ft_error('symmetry constraints are only supported for two-dipole models');
  elseif strcmp(cfg.symmetry, 'x')
    % this structure is passed onto the low-level FT_INVERSE_DIPOLEFIT function
    cfg.dipfit.constr.reduce = [1 2 3];         % select the parameters [x1 y1 z1]
    cfg.dipfit.constr.expand = [1 2 3 1 2 3];   % repeat them as [x1 y1 z1 x1 y1 z1]
    cfg.dipfit.constr.mirror = [1 1 1 -1 1 1];  % multiply each of them with 1 or -1, resulting in [x1 y1 z1 -x1 y1 z1]
  elseif strcmp(cfg.symmetry, 'y')
    % this structure is passed onto the low-level FT_INVERSE_DIPOLEFIT function
    cfg.dipfit.constr.reduce = [1 2 3];         % select the parameters [x1 y1 z1]
    cfg.dipfit.constr.expand = [1 2 3 1 2 3];   % repeat them as [x1 y1 z1 x1 y1 z1]
    cfg.dipfit.constr.mirror = [1 1 1 1 -1 1];  % multiply each of them with 1 or -1, resulting in [x1 y1 z1 x1 -y1 z1]
  elseif strcmp(cfg.symmetry, 'z')
    % this structure is passed onto the low-level FT_INVERSE_DIPOLEFIT function
    cfg.dipfit.constr.reduce = [1 2 3];         % select the parameters [x1 y1 z1]
    cfg.dipfit.constr.expand = [1 2 3 1 2 3];   % repeat them as [x1 y1 z1 x1 y1 z1]
    cfg.dipfit.constr.mirror = [1 1 1 1 1 -1];  % multiply each of them with 1 or -1, resulting in [x1 y1 z1 x1 y1 -z1]
  else
    ft_error('unrecognized symmetry constraint');
  end
elseif ~isfield(cfg, 'dipfit') || ~isfield(cfg.dipfit, 'constr')
  % no symmetry constraints have been specified
  cfg.dipfit.constr = [];
end

if ft_getopt(cfg.dipfit.constr, 'sequential', false) && strcmp(cfg.model, 'moving')
  ft_error('the moving dipole model does not combine with the sequential constraint')
  % see http://bugzilla.fieldtriptoolbox.org/show_bug.cgi?id=3119
end

if iscomp
  % transform the data into a representation on which the timelocked dipole fit can perform its trick
  data = comp2timelock(cfg, data);

  % default component selection is all components
  if ischar(cfg.component) && strcmp(cfg.component, 'all')
    cfg.component = (1:size(data.avg, 2));
  end
elseif isfreq
  % transform the data into a representation on which the timelocked dipole fit can perform its trick
  data = freq2timelock(cfg, data);
elseif istimelock
  % no transformation is needed
end

% collect and preprocess the electrodes/gradiometer and head model
% this will also update cfg.channel to match the electrodes/gradiometers
[headmodel, sens, cfg] = prepare_headmodel(cfg, data);

% construct the low-level options for the leadfield computation as key-value pairs, these are passed to FT_COMPUTE_LEADFIELD and FT_INVERSE_DIPOLEFIT
leadfieldopt = {};
leadfieldopt = ft_setopt(leadfieldopt, 'reducerank',     ft_getopt(cfg, 'reducerank'));
leadfieldopt = ft_setopt(leadfieldopt, 'backproject',    ft_getopt(cfg, 'backproject'));
leadfieldopt = ft_setopt(leadfieldopt, 'normalize',      ft_getopt(cfg, 'normalize'));
leadfieldopt = ft_setopt(leadfieldopt, 'normalizeparam', ft_getopt(cfg, 'normalizeparam'));
leadfieldopt = ft_setopt(leadfieldopt, 'weight',         ft_getopt(cfg, 'weight'));

% construct the low-level options for the dipole fitting as key-value pairs, these are passed to FT_INVERSE_DIPOLEFIT
dipfitopt = ft_cfg2keyval(cfg.dipfit);

% select the desired channels, ordered according to the sensor structure or configuration
[selcfg, seldata] = match_str(cfg.channel, data.label);
% take the selected channels from the data structure
Vdata = data.avg(seldata, :);

% sphere the date using the noise covariance matrix supplied, if any
% this affects both the gridsearch and the nonlinear optimization
noisecov = ft_getopt(cfg.dipfit, 'noisecov');
if ~isempty(noisecov)
  [u, s] = svd(noisecov);
  tol = max(size(noisecov)) * eps(norm(s, inf));
  s = diag(s);
  r1 = sum(s > tol) + 1;
  s(1:(r1 - 1)) = 1 ./ sqrt(s(1:(r1 - 1)));
  s(r1:end)     = 0;
  sphere = diag(s) * u';
  % apply the sphering to the data
  Vdata = sphere * Vdata;
  % apply the sphering as a pre-multiplication to the sensor definition
  montage = [];
  montage.labelold = cfg.channel;
  montage.labelnew = cfg.channel;
  montage.tra = sphere;
  sens = ft_apply_montage(sens, montage, 'balancename', 'sphering');
end

if iscomp
  % select the desired component topographies
  Vdata = Vdata(:, cfg.component);
elseif isfreq
  % the desired frequencies have already been selected
  Vdata = Vdata(:, :);
elseif istimelock
  % select the desired latencies
  if ischar(cfg.latency) && strcmp(cfg.latency, 'all')
    cfg.latency = data.time([1 end]);
  end
  tbeg = nearest(data.time, cfg.latency(1));
  tend = nearest(data.time, cfg.latency(end));
  cfg.latency = [data.time(tbeg) data.time(tend)];
  Vdata = Vdata(:, tbeg:tend);
end

nchans = size(Vdata,1);
ntime  = size(Vdata,2);
Vmodel = zeros(nchans, ntime);
ft_info('selected %d channels\n', nchans);
ft_info('selected %d topographies\n', ntime);

if nchans<cfg.numdipoles*3
  ft_warning('not enough channels to perform a dipole fit');
end

if ntime<1
  ft_error('no spatial topography selected');
end

% check whether EEG is average referenced
if ft_senstype(sens, 'eeg')
  if any(rv(Vdata, avgref(Vdata))>0.001)
    ft_warning('the EEG data is not average referenced, correcting this');
  end
  Vdata = avgref(Vdata);
end

% set to zeros if no initial dipole was specified
if ~isfield(cfg, 'dip')
  cfg.dip.pos = zeros(cfg.numdipoles, 3);
  cfg.dip.mom = zeros(3*cfg.numdipoles, 1);
end

% set to zeros if no initial dipole position was specified
if ~isfield(cfg.dip, 'pos')
  cfg.dip.pos = zeros(cfg.numdipoles, 3);
end

% set to zeros if no initial dipole moment was specified
if ~isfield(cfg.dip, 'mom')
  cfg.dip.mom = zeros(3*cfg.numdipoles, 1);
end

% check the specified dipole model
if numel(cfg.dip.pos)~=cfg.numdipoles*3 || numel(cfg.dip.mom)~=cfg.numdipoles*3
  ft_error('inconsistent number of dipoles in configuration')
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% perform the dipole scan, this is usefull for generating an initial guess
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if strcmp(cfg.gridsearch, 'yes')
  % test whether we have a valid configuration for dipole scanning
  if cfg.numdipoles==1
    % this is ok
  elseif cfg.numdipoles==2 && ~isempty(cfg.dipfit.constr)
    % this is also ok
  elseif isfield(cfg.sourcemodel, 'pos') && size(cfg.sourcemodel.pos,2)==cfg.numdipoles*3
    % this is also ok
  else
    ft_error('dipole scanning is only possible for a single dipole or a symmetric dipole pair');
  end

  if isfield(cfg.sourcemodel, 'leadfield')
    ft_notice('using precomputed leadfields for the gridsearch');

    sourcemodel = keepfields(cfg.sourcemodel, {'pos', 'tri', 'dim', 'inside', 'leadfield', 'leadfielddimord', 'label'});

    % select the channels corresponding to the data and the user configuration
    tmpcfg = keepfields(cfg, 'channel');
    sourcemodel = ft_selectdata(tmpcfg, sourcemodel);

    % sort the channels to be consistent with the data
    [dum, chansel] = match_str(data.label, sourcemodel.label);
    sourcemodel.label = sourcemodel.label(chansel);
    for i=1:numel(sourcemodel.leadfield)
      if ~isempty(sourcemodel.leadfield{i})
        sourcemodel.leadfield{i} = sourcemodel.leadfield{i}(chansel, :);
      end
    end

    % ensure that the channels are consistent with the data
    assert(isequal(sourcemodel.label, cfg.channel), 'cannot match the channels in the sourcemodel to those in the data')

  else
    ft_notice('computing the leadfields for the gridsearch on the fly');

    % construct the dipole positions on which the source reconstruction will be done
    tmpcfg           = keepfields(cfg, {'sourcemodel', 'mri', 'headshape', 'symmetry', 'smooth', 'threshold', 'spheremesh', 'inwardshift', 'xgrid' 'ygrid', 'zgrid', 'resolution', 'tight', 'warpmni', 'template', 'showcallinfo', 'trackcallinfo', 'trackusage', 'trackdatainfo', 'trackmeminfo', 'tracktimeinfo', 'checksize'});
    tmpcfg.headmodel = headmodel;
    if ft_senstype(sens, 'eeg')
      tmpcfg.elec = sens;
    elseif ft_senstype(sens, 'meg')
      tmpcfg.grad = sens;
    end
    sourcemodel = ft_prepare_sourcemodel(tmpcfg);

  end % if precomputed leadfield or not

  ngrid = size(sourcemodel.pos,1);

  switch cfg.model
    case 'regional'
      sourcemodel.error = nan(ngrid, 1);
    case 'moving'
      sourcemodel.error = nan(ngrid, ntime);
    otherwise
      ft_error('unsupported cfg.model');
  end

  insideindx = find(sourcemodel.inside);
  ft_progress('init', cfg.feedback, 'scanning grid');
  for i=1:length(insideindx)
    ft_progress(i/length(insideindx), 'scanning grid location %d/%d\n', i, length(insideindx));
    thisindx = insideindx(i);
    if isfield(sourcemodel, 'leadfield')
      % reuse the previously computed leadfield
      lf = sourcemodel.leadfield{thisindx};
    else
      lf = ft_compute_leadfield(sourcemodel.pos(thisindx,:), sens, headmodel, leadfieldopt{:});
    end
    % the model is V=lf*mom+noise, therefore mom=pinv(lf)*V estimates the
    % dipole moment this makes the model potential U=lf*pinv(lf)*V and the
    % model error is norm(V-U) = norm(V-lf*pinv(lf)*V) = norm((eye-lf*pinv(lf))*V)
    if any(isnan(lf(:)))
      % this might happen if one of the dipole locations of the grid is
      % outside the brain compartment
      lf(:) = 0;
    end
    switch cfg.model
      case 'regional'
        % sum the error over all latencies
        sourcemodel.error(thisindx,1) = sum(sum(((eye(nchans)-lf*pinv(lf))*Vdata).^2));
      case 'moving'
        % remember the error for each latency independently
        sourcemodel.error(thisindx,:) = sum(((eye(nchans)-lf*pinv(lf))*Vdata).^2);
      otherwise
        ft_error('unsupported cfg.model');
    end % switch model
  end % looping over the grid
  ft_progress('close');

  switch cfg.model
    case 'regional'
      % find the source position with the minimum error
      [err, indx] = min(sourcemodel.error);
      dip.pos = sourcemodel.pos(indx,:);                % note that for a symmetric dipole pair this results in a vector
      dip.pos = reshape(dip.pos,3,cfg.numdipoles)';     % convert to a Nx3 array
      dip.mom = zeros(cfg.numdipoles*3,1);              % set the dipole moment to zero
      if cfg.numdipoles==1
        ft_info('found minimum after scanning on grid point [%g %g %g]\n', dip.pos(1), dip.pos(2), dip.pos(3));
      elseif cfg.numdipoles==2
        ft_info('found minimum after scanning on grid point [%g %g %g; %g %g %g]\n', dip.pos(1,1), dip.pos(1,2), dip.pos(1,3), dip.pos(2,1), dip.pos(2,2), dip.pos(2,3));
      end

    case 'moving'
      for t=1:ntime
        % find the source position with the minimum error
        [err, indx] = min(sourcemodel.error(:,t));
        dip(t).pos = sourcemodel.pos(indx,:);                 % note that for a symmetric dipole pair this results in a vector
        dip(t).pos = reshape(dip(t).pos,3,cfg.numdipoles)';   % convert to a Nx3 array
        dip(t).mom = zeros(cfg.numdipoles*3,1);               % set the dipole moment to zero
        if cfg.numdipoles==1
          ft_info('found minimum after scanning for topography %d on grid point [%g %g %g]\n', t, dip(t).pos(1), dip(t).pos(2), dip(t).pos(3));
        elseif cfg.numdipoles==2
          ft_info('found minimum after scanning for topography %d on grid point [%g %g %g; %g %g %g]\n', t, dip(t).pos(1,1), dip(t).pos(1,2), dip(t).pos(1,3), dip(t).pos(2,1), dip(t).pos(2,2), dip(t).pos(2,3));
        end
      end

    otherwise
      ft_error('unsupported cfg.model');
  end % switch model

elseif strcmp(cfg.gridsearch, 'no')
  % there is no grid needed for dipole scanning
  sourcemodel = [];
  % use the initial guess supplied in the configuration for the remainder
  switch cfg.model
    case 'regional'
      dip = cfg.dip;
    case 'moving'
      for t=1:ntime
        dip(t) = cfg.dip;
      end
    otherwise
      ft_error('unsupported cfg.model');
  end % switch model

end % if gridsearch yes/no

% multiple dipoles can be represented either as a 1x(N*3) vector or as a Nx3 matrix,
% i.e. [x1 y1 z1 x2 y2 z2] or [x1 y1 z1; x2 y2 z2]
switch cfg.model
  case 'regional'
    dip = fixdipole(dip);
  case 'moving'
    for t=1:ntime
      dip(t) = fixdipole(dip(t));
    end
  otherwise
    ft_error('unsupported cfg.model');
end % switch model

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% perform the non-linear fit
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if strcmp(cfg.nonlinear, 'yes')
  switch cfg.model
    case 'regional'
      % perform the non-linear dipole fit for all latencies together
      % catch errors due to non-convergence
      try
        dip = ft_inverse_dipolefit(dip, sens, headmodel, Vdata, dipfitopt{:}, leadfieldopt{:});
        success = 1;
        if cfg.numdipoles==1
          ft_info('found minimum after non-linear optimization on [%g %g %g]\n', dip.pos(1), dip.pos(2), dip.pos(3));
        elseif cfg.numdipoles==2
          ft_info('found minimum after non-linear optimization on [%g %g %g; %g %g %g]\n', dip.pos(1,1), dip.pos(1,2), dip.pos(1,3), dip.pos(2,1), dip.pos(2,2), dip.pos(2,3));
        end
      catch
        success = 0;
        disp(lasterr);
      end

    case 'moving'
      % perform the non-linear dipole fit for each latency independently
      % instead of using dip(t) = ft_inverse_dipolefit(dip(t),...), I am using temporary variables dipin and dipout
      % to prevent errors like "Subscripted assignment between dissimilar structures"
      dipin = dip;
      for t=1:ntime
        % catch errors due to non-convergence
        try
          dipout(t) = ft_inverse_dipolefit(dipin(t), sens, headmodel, Vdata(:,t), dipfitopt{:}, leadfieldopt{:});
          success(t) = 1;
          if cfg.numdipoles==1
            ft_info('found minimum after non-linear optimization for topography %d on [%g %g %g]\n', t, dipout(t).pos(1), dipout(t).pos(2), dipout(t).pos(3));
          elseif cfg.numdipoles==2
            ft_info('found minimum after non-linear optimization for topography %d on [%g %g %g; %g %g %g]\n', t, dipout(t).pos(1,1), dipout(t).pos(1,2), dipout(t).pos(1,3), dipout(t).pos(2,1), dipout(t).pos(2,2), dipout(t).pos(2,3));
          end
        catch
          % keep the position and moment according to the initial guess
          dipout(t).pos = dipin(t).pos;
          dipout(t).mom = dipin(t).mom;
          success(t) = 0;
          disp(lasterr);
        end
      end
      dip = dipout;
      clear dipin dipout
    otherwise
      ft_error('unsupported cfg.model');
  end % switch model
end % if nonlinear

if strcmp(cfg.nonlinear, 'no')
  % the optimal dipole positions are either obtained from scanning
  % or from the initial configured specified by the user
  switch cfg.model
    case 'regional'
      success = 1;
    case 'moving'
      success = ones(1,ntime);
    otherwise
      ft_error('unsupported cfg.model');

  end % switch model
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% compute the model potential distribution and the residual variance
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
switch cfg.model
  case 'regional'
    if success
      % re-compute the leadfield in order to compute the model potential and dipole moment
      lf = ft_compute_leadfield(dip.pos, sens, headmodel, leadfieldopt{:});
      if isfield(dip, 'mom') && isfield(dip, 'ampl')
        % the orientation and amplitude have already been estimated, this applies to the case of a fixed dipole orientation
        dip.pot = (lf * dip.mom) * dip.ampl;
      else
        % compute all details of the final dipole model using linear estimation
        dip.mom = pinv(lf)*Vdata;
        dip.pot = lf*dip.mom;
      end
      dip.rv  = rv(Vdata, dip.pot);
      Vmodel  = dip.pot;
    end
  case 'moving'
    for t=1:ntime
      if success(t)
        % re-compute the leadfield in order to compute the model potential and dipole moment
        lf = ft_compute_leadfield(dip(t).pos, sens, headmodel, leadfieldopt{:});
        % compute all details of the final dipole model
        dip(t).mom = pinv(lf)*Vdata(:,t);
        dip(t).pot = lf*dip(t).mom;
        dip(t).rv  = rv(Vdata(:,t), dip(t).pot);
        Vmodel(:,t) = dip(t).pot;
      end
    end
  otherwise
    ft_error('unsupported cfg.model');
end % switch model

switch cfg.model
  case 'regional'
    if isfreq
      % the matrix with the dipole moment is encrypted and cannot be interpreted straight away
      % reconstruct the frequency representation of the data at the source level
      if isfield(dip, 'mom') && isfield(dip, 'ampl')
        % this applies to the case of a fixed dipole orientation
        [dip.pow, dip.csd, dip.fourier] = timelock2freq(dip.mom * dip.ampl);
      else
        [dip.pow, dip.csd, dip.fourier] = timelock2freq(dip.mom);
      end
    end
  case 'moving'
    if isfreq
      % although this is technically possible so far, it does not make any sense
      ft_warning('a moving dipole model in the frequency domain is not supported');
    end
  otherwise
    ft_error('unsupported cfg.model');
end % switch model

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% collect the results
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
source.label  = cfg.channel; % these channels were used in fitting
source.dip    = dip;
source.Vdata  = Vdata;  % FIXME this should be renamed (if possible w.r.t. EEGLAB)
source.Vmodel = Vmodel; % FIXME this should be renamed (if possible w.r.t. EEGLAB)

% the units of the fitted source are the same as the units of the headmodel and the sensor array
for i=1:length(source.dip)
  if isfield(headmodel, 'unit')
    source.dip(i).unit = headmodel.unit;
  elseif isfield(sourcemodel, 'unit')
    source.dip(i).unit = sourcemodel.unit;
  end
end

% assign a latency, frequeny or component axis to the output
if iscomp
  source.component = cfg.component;
  % FIXME assign Vdata to an output variable, idem for the model potential
elseif isfreq
  source.freq   = cfg.frequency;
  source.dimord = 'chan_freq';
  % FIXME assign Vdata to an output variable, idem for the model potential
elseif istimelock
  tbeg = nearest(data.time, cfg.latency(1));
  tend = nearest(data.time, cfg.latency(end));
  source.time   = data.time(tbeg:tend);
  source.dimord = 'chan_time';
end

% do the general cleanup and bookkeeping at the end of the function
ft_postamble debug
ft_postamble previous   data
ft_postamble provenance source
ft_postamble history    source
ft_postamble savevar    source


% Astar
% A* navigation class
%
% A concrete subclass of the Navigation class that implements the A*
% navigation algorithm. Methods included are for the standard case,
% multiobjective optimization (MOO) -- i.e. optimizes over several
% objectives/criteria -- and the A*-PO algorithms for MOO that utilizes
% Pareto optimality.
%
% Methods:
% 	plan            Compute the cost map given a goal and map
% 	path            Compute a path to the goal
% 	visualize       Display the obstacle map (deprecated)
% 	plot            Display the obstacle map
% 	costmap_modify 	Modify the costmap
% 	costmap_get     Return the current costmap
% 	costmap_set     Set the current costmap
% 	display         Print the parameters in human readable form
% 	char            Convert to string
%
% Properties:
% TBD
%
% Example 1::
%        load map1              % load map
%        goal = [50;30];
%        start=[20;10];
%        as = Astar(map);       % create Navigation object
%        as.plan(goal,2,3,0);   % setup costmap for specified goal;
%                               % standard D* algorithm w/ 2 objectives
%                               % and 3 costmap layers
%        as.path(start);        % plan solution path start-to-goal, animate
%        P = as.path(start);    % plan solution path start-to-goal, return
%                               % path
% Example 2::
%        goal = [100;100];
%        start = [1;1];
%        as = Astar(0);          % create Navigation object with pseudo-
%                                % random occupancy grid
%        ds.addCost(terrain);    % terrain is a 100x100 matrix of
%                                % elevations [0,1]
% 	     ds.plan(goal,3,4,0);    % setup costmap for specified goal
%                                % (3 and 4 include the added terrain cost)
%        as.path(start);         % plan solution path start-goal, animate
%        P = as.path(start);     % plan solution path start-goal, return
%                                % path
%
% Notes
% - Obstacles are represented by Inf in the costmap.
%
% References
% - A Pareto Optimal D* Search Algorithm for Multiobjective Path Planning,
%   A. Lavin.
% - A Pareto Front-Based Multiobjective Path Planning Algorithm, A. Lavin.
% - Robotics, Vision & Control, Sec 5.2.2, Peter Corke, Springer, 2011.
%
% See Also Navigation, Dstar

% Copyright (C) 1993-2017, by Peter I. Corke
%
% This file is part of The Robotics Toolbox for MATLAB (RTB).
%
% RTB is free software: you can redistribute it and/or modify
% it under the terms of the GNU Lesser General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% RTB is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU Lesser General Public License for more details.
%
% You should have received a copy of the GNU Leser General Public License
% along with RTB.  If not, see <http://www.gnu.org/licenses/>.
%
% http://www.petercorke.com

% Implementation notes:
%
% X is an index into the array of states.
% State pointers are kept as matlab array index rather than row,col format.

classdef Astar < Navigation

    properties (SetAccess=private, GetAccess=private)

        % essential world info
        costmap   % cost layers (1st is map)
        G         % index of goal point
        N         % number of objectives
        L         % number of cost layers

        % info kept per cell (state):
        b         % backpointer (0 means not set)
        t         % tag: NEW/OPEN/CLOSED

        algorithm % A*, A*-MOO, or A*-PO
        tie
        openlist  % priority queue with states and their costs
        niter
        changed
        openlist_maxlen
        quiet     % specifies verbosity

        % tag state values
        NEW = 0;
        OPEN = 1;
        CLOSED = 2;
    end

    methods  % start of public methods

        function as = Astar(world, varargin)
            %Astar.Astar A* constructor
            %
            % AS = Astar(MAP, OPTIONS) is a A* navigation object, and MAP
            % is an occupancy grid, a representation of a planar world as
            % a matrix whose elements are 0 (free space) or 1 (occupied).
            % The occupancy grid is coverted to a costmap with a unit cost
            % for traversing a cell.
            %
            % Options::
            % 'world' = 0   will call for a pseudo-random occupancy grid
            % 'goal',G      Specify the goal point (2x1)
            % 'metric',M    Specify the distance metric as 'Euclidean'
            %               (default) or 'cityblock'
            % 'inflate',K   Inflate all obstacles by K cells
            % 'quiet'       Don't display the progress spinner
            %
            % Other options are supported by the Navigation superclass.
            %
            % See also Navigation.Navigation.

            % Invoke the superclass constructor
            as = as@Navigation(world, varargin{:});  % includes the occgrid

            % Options
            opt.quiet = false;
            opt = tb_optparse(opt, varargin);
            as.quiet = opt.quiet;

            as.occgrid2costmap(as.occgrid);

            % Initialize A* state variables
            as.reset();
            if ~isempty(as.goal)
                as.goal_change();
            end
            as.changed = false;
        end

        function reset(as)
            %Astar.reset Reset the planner
            %
            % AS.reset() resets the A* planner.  The next instantiation
            % of AS.plan() will perform a global replan.

            % Build the matrices required to hold the state of each cell
            as.b = zeros(size(as.occgrid), 'uint32');     % backpointers
            as.t = zeros(size(as.occgrid), 'uint8');      % tags, all NEW=0
            as.costmap(:,:,2) = zeros(size(as.occgrid));  % path cost g
            as.costmap(:,:,3) = zeros(size(as.occgrid));  % path cost h

            % Priority queue has col for each open state, one row for the
            % state location, and rows for each cost layer
            as.openlist = zeros(as.L+1,0);
            as.openlist_maxlen = -Inf;
        end

        function goal_change(as)
            %Astar.goal_change Changes the costlayers due to new goal
            %position
            if isempty(as.b)
                return;
            end
            goal = as.goal;

            % Keep goal in index rather than row,col format
            as.G = sub2ind(size(as.occgrid), goal(2), goal(1));
            as.INSERT(as.G, as.projectCost(as.G), 'goalset');
            as.costmap(goal(2),goal(1),2) = 0;

            % If new goal modifies costs for a layer, recalculate here
            as.calcHeuristic(as.occgrid, as.goal);
        end

        function s = char(as)
            %Astar.char Convert Navigation object to string
            %
            % AS.char() is a string representing the state of the Astar
            % object in human-readable form.
            %
            % See also Astar.display, Navigation.char.

            % Work is done by the superclass
            s = char@Navigation(as);
        end

        function plot(as, varargin)
            %Astar.plot Visualize navigation environment
            %
            % AS.plot() displays the occupancy grid and the goal distance
            % in a new figure.  The goal distance is shown by intensity
            % which increases with distance from the goal.  Obstacles are
            % overlaid and shown in red.
            %
            % AS.plot(P) as above but also overlays a path given by the set
            % of points P (Mx2).
            %
            % See also Navigation.plot.

            plot@Navigation(as, 'distance', as.costmap(:,:,3), varargin{:});
        end

        function n = next(as, current)
            % Invoked by Navigation.step
            % Backpropagate from goal to start
            % Return [col;row] of previous step
            if as.changed
                error('Cost map has changed, replan');
            end
            X = sub2ind(size(as.occgrid), current(2), current(1));
            % Set X as the backpointer of X
            X = as.b(X);
            if X == 0
                % Goal (no further backpointer)
                n = [];
            else
                [r,c] = ind2sub(size(as.occgrid), X);
                n = [c;r];
            end
        end

        function plan(as, goal, N, layers, algorithm)
            %Astar.plan Prep the grid for planning.
            %
            % AS.plan() updates AS with a costmap of distance to the
            % goal from every non-obstacle point in the map.  The goal is
            % as specified to the constructor.
            %
            % Inputs:
            %   goal: goal state coordinates
            %   N: number of optimization objectives; standard A* is 2
            %   (i.e. distance and heuristic)
            %   layers: number of cost layers in costmap
            %   algorithm: specify standard A*(0), A*-MOO (1), A*-PO (2)

            % Setup parameters
            if nargin < 3
                N = 2;
            end
            if nargin < 4
                layers = 3;
            end
            if nargin < 5
                algorithm = 0;
            end
            as.N = N;
            as.L = layers;
            as.algorithm = algorithm;
            as.openlist = zeros(as.L+1,0);

            % Initialize cost layers
            for a = 2:as.L
                as.costmap(:,:,a) = zeros(size(as.occgrid));
            end

            % Cost priority/tiebreaker: layer 2 (distance to node)
            as.tie = 2;

            % Set goal
            if nargin > 1
                as.goal = goal;  % invokes superclass method set.goal()
            end
            if isempty(as.goal)
                error('must specify a goal point');
            end

            % Populate heuristic cost layer
            as.calcHeuristic(as.occgrid, as.goal);
        end

        function P = path(as, start)
        %Astar.path Find a path between two points
        %
        % AS.path(START) finds and displays a path from START to GOAL
        % which is overlaid on the occupancy grid.
        %
        % P = AS.path(START) returns the path (2xM) from START to GOAL.
            if nargin < 1
                error('must specify start state');
            end

            % Invoke the superclass path function, which iterates on our
            % next method
%             start = [start; 1];  % specifies backpropogation for NAV.path()
%             temp = start;
%             start = as.goal;
%             as.goal = temp;

            if nargout == 0
                path@Navigation(as, start);
            else
                P = path@Navigation(as, start);
            end
        end

        % Handler invoked by Navigation.path() to start the navigation
        % process -- calculate the solution path.
        % Line comments Ln reference A* pseudocode in Lavin's "A Pareto
        % Front-Based Multiobjective Path Planning Algorithm" where n is
        % the line number.
        function navigate_init(as, start)
            as.openlist = zeros(as.L+1,0);  % openlist must be empty
            % Begin search with the start node
            start = sub2ind(size(as.occgrid), start(2), start(1));
            as.openlist(1,1) = start;
            as.t(start) = as.OPEN;

            % Plan the A* path
            as.niter = 0; flag = 0;
            while ~isempty(as.openlist)                              % L4
                % Normalize costs on the open list, choose expansion state
                queue = normc(as.openlist(2:size(as.openlist,1),:)');
                if as.algorithm == 2
                    % Get Pareto optimal point off the open list
                    front = as.openlist(:,paretofront(queue));
                    [~,col] = min(front(as.tie+1,:));
                    X = front(1,col);                                % L5
                else
                    [~,ind]=min(sum(queue,2));
                    X = as.openlist(1,ind);                          % L5
                end
                    as.DELETE(X);                                    % L6

                as.niter = as.niter + 1;
                if ~as.quiet && mod(as.niter, 20) == 0
                    as.spinner();
                end

                % Populate the openlist
                for Y=as.neighbors(X)                                % L7,8
                    if(Y==as.G)                                      % L9
                        as.b(Y) = X;
                        as.updateCosts(Y,X,as.N)
                        flag = 1;  % flag for goal
                        break;
                    end
                    if as.t(Y)==as.NEW && as.costmap(Y)~=Inf
                        as.b(Y) = X;
                        as.updateCosts(Y,X,as.N);
                        % Project node's costs into objective space:
                        objspace = as.projectCost(Y,X);
                        as.INSERT(Y, objspace, '');
                    end
                end
                if as.verbose
                    disp(' ')
                end
                if flag==1 % goal found
                    break;
                end
            end
            if ~as.quiet
                fprintf('\r');
            end
            as.changed = false;
        end

        function layer = cost_get(as, layer)
        %Astar.cost_get Get the specified cost layer
            layer = as.costmap(:,:,layer);
        end

        function c = heurstic_get(as)
        %Astar.heuristice_get Get the current heuristic map
        %
        % C = AS.heuristice_get() is the current heuristic layer. It is
        % computed in Astar.plan.
        %
        % See also Astar.plan.
            c = as.costmap(:,:,3);
        end

        function c = costmap_get(as)
        %Astar.costmap_get Get the current costmap
        %
        % C = AS.costmap_get() is the current costmap.
        % The value of each element represents the cost of traversing the
        % cell.  It is autogenerated by the class constructor from the
        % occupancy grid such that:
        % - free cell (occupancy 0) has a cost of 1
        % - occupied cell (occupancy >0) has a cost of Inf
        %
        % See also Astar.costmap_set, Astar.costmap_modify.
            c = as.costmap;
        end

        function costmap_set(as, costmap)
        %Astar.costmap_set Set the current costmap
        %
        % AS.costmap_set(C) sets the current costmap.
        % This method accepts the full costmap -- i.e. all layers.
        %
        % Notes:
        % - After the cost map is changed the path should be replanned by
        %   calling AS.plan().
        %
        % See also Astar.costmap_get, Astar.costmap_modify.
            [i,j,k] = size(costmap);
            if ~all([i,j] == size(as.occgrid))
                error('costmap must be same size as occupancy grid');
            end
            as.L = k;  % set the number of cost layers
            as.costmap = costmap;
            as.changed = true;
        end

        function costmap_modify(as, point, newcost)
        %Astar.costmap_modify Modify cost map
        %
        % AS.costmap_modify(P, NEW) modifies the cost map at P=[X,Y] to
        % have the value NEW.  If P (2xM) and NEW (1xM) then the cost of
        % the points defined by the columns of P are set to the corresponding
        % elements of NEW.
        %
        % Notes::
        % - After one or more point costs have been updated the path
        %   should be replanned by calling AS.plan().
        %
        % See also Astar.costmap_set, Astar.costmap_get.
            if (newcost < 0) || (1 < newcost)
                error('new cost value must be normlaized [0,1]')
            end

            [i,j,k] = size(as.costmap);
            if (point(1) < 0) || (point(1) > i)
                error('1st dimension of point is out of bounds')
            end
            if (point(2) < 0) || (point(2) > j)
                error('2nd dimension of point is out of bounds')
            end
            if (point(2) < 0) || (point(3) > k)
                error('3rd dimension of point is out of bounds')
            end

            as.costmap(point) = newcost;
        end

        function addCost(as, values)
        %Astar.addCost Add an additional cost layer
        %
        % AS.addCost(values) adds the matrix specified by values as a
        % cost layer.
        % Inputs
        %   values: normalized matrix the size of the environment
            [i,j,k] = size(as.costmap);

            if [i,j]~=size(as.occgrid)
                error('layer size does not match the environment')
            end
            if max(max(values))~=1 || min(min(values))~=0
                error('layer values are not normalized [0,1]')
            end

            as.costmap(:,:,k+1) = values;
        end

        function flag = backProp(as)
            flag = 1;
        end

    end  % end of public methods

    methods (Access=protected)  % start of private methods

        function occgrid2costmap(as, og, cost)
            if nargin < 3
                cost = 1;
            end
            og(og==1) = Inf;  % occupied cells -> infinite path cost
            og(og==0) = cost;  % unoccupied cells -> path cost
            as.costmap(:,:,1) = og;
        end

        function calcHeuristic(as, grid, goal)
            as.costmap(:,:,3) = zeros(size(grid));
            for ii=1:size(grid,1)
                for jj=1:size(grid,2)
                    as.costmap(ii,jj,3) = sqrt((ii-goal(1))^2+(jj-goal(2))^2);
                end
            end
        end

        function k_new = updateCosts(as, a, b, obj)
            % NOTE: Only for costs that accumulate (i.e. sum) over the
            % path, and for dynamic costs.
            % E.g. the heuristic parameter only needs updating when the
            % goal state changes; its values are stored for each cell.
            %
            % Location moving from state b to a.
            %
            % The costs are coded to be (1) distance, (2) heuristic, (3)
            % elevation, (4) solar deviation, and (5) risk. If deviating
            % from these costs (in this order) you MUST EDIT THIS METHOD.
            [i,j,~] = size(as.costmap);

            if nargout > 0
                k_new = as.costmap(i*j+b) + as.dc(b,a);
                return
            end
            if obj == 0
                % Return what the new priority cost would be (k_new)
                return
            end
            if obj > 1
                % Standard A* search
                as.costmap(i*j+a) = as.costmap(i*j+b) + as.dc(b,a);
                % (no heuristic update needed)
            end
            if obj > 2
                % W/ elevation costs
                % (no elevation update needed)
            end
            if obj > 3
                % W/ solar costs
                % Rotate the solar vector 1rad per 100 steps
                sV = [cos(as.niter/100);sin(as.niter/100)];
                as.costmap(4*i*j+a) = dot(sV,as.vc(b,a));
            end
            if obj > 4
                % W/ risk costs
                % (no risk update needed)
            end
        end

        function pt = projectCost(as, a, b)
            % Returns the projection of state a into objective space. If
            % specified, location is moving from b to a (case 3).
            [i,j,k] = size(as.costmap);
            pt(1) = as.costmap(a);
            switch nargin
                case 2
                    pt(2) = as.costmap(i*j+a);
                case 3
                    pt(2) = as.costmap(i*j+b) + as.dc(a,b);
                otherwise
                    return
            end
            for n=3:k
                pt(n) = as.costmap((n-1)*i*j+a);
            end
        end

        function INSERT(as, X, pt, where)
            % Add state X to the openlist with objective space values
            % specified by pt.

            if nargin>2
                as.message('insert (%s) %d = %f\n', where, X, pt);
            end

            i = find(as.openlist(1,:) == X);
            if length(i) > 1
                error('A*:INSERT: state in open list %d times', X);
            end

            [i,j,~] = size(as.costmap);
            if (as.t(X) == as.OPEN || as.CLOSED) && ...
               (pt(as.tie) > as.costmap((as.tie-1)*i*j+X))
                % L13/14: If a node with same position as successor is in
                % the OPEN/CLOSED list & has a lower f than successor,
                % then skip this successor.
            else
                % Add a new column to the open list for this node
                as.openlist = [as.openlist [X; pt(:)]];
            end

            % Keep track of the max length of the openlist
            if numcols(as.openlist) > as.openlist_maxlen
                as.openlist_maxlen = numcols(as.openlist);
            end

            % Tag state X as open
            as.t(X) = as.OPEN;
        end

        function DELETE(as, X)
            as.message('delete %d\n', X);
            i = find(as.openlist(1,:) == X);
            if length(i) ~= 1
                error('A*:DELETE: state %d does not exist', X);
            end

            % Remove the column, close the state
            as.openlist(:,i) = [];
            as.t(X) = as.CLOSED;
        end

        function cost = dc(as, X, Y)
            % Return the distance cost of moving from state X to state Y
            [r,c] = ind2sub(size(as.occgrid), [X; Y]);
            dist = sqrt(sum(diff([r c]).^2));
            dcost = (as.costmap(X) + as.costmap(Y))/2;

            cost = dist * dcost;
        end

        function vector = vc(as, X, Y)
            % Return the robot unit vector -- direction of moving from
            % state X to state Y
            [Xi,Xj] = ind2sub(size(as.occgrid),X);
            [Yi,Yj] = ind2sub(size(as.occgrid),Y);
            vector = [Yi-Xi;Yj-Xj];
            vector = vector/norm(vector);
        end

        function Y = neighbors(as, X)
            % Return indices of neighbor states (max 8) as a row vector
            dims = size(as.occgrid);
            [r,c] = ind2sub(dims, X);

            % Of 8-way neighbors, only use those w/in grid bounds
            Y = [r-1 r-1 r-1 r r  r+1 r+1 r+1; c-1 c c+1 c-1 c+1 c-1 c c+1];
            k = (min(Y)>0) & (Y(1,:)<=dims(1)) & (Y(2,:)<=dims(2));
            Y = Y(:,k);
            Y = sub2ind(dims, Y(1,:)', Y(2,:)')';
        end

    end  % end of private methods

end


function retPhases = atlasEstimatePhases(vw, areaCorners, scanNum, fieldName)
% Define retinotopy phases for an angle and wedge map
%
%   retPhases = atlasEstimatePhases(vw, areaCorners, scanNum, fieldName)
%
% Lets the user interactively define the data phases corresponding to the
% fovea, periphery, upper and lower vertical meridians for the atlas and
% data fits.
%
% scanNum is the scan number of the wedge and ring in this view (and for
% the view's selected data type). fieldName is a 1 x 2 cell array
% specifying the field containing the polar angle and eccentricity data,
% respectively. (For PRF analyses, these are not necessarily loaded into
% the phase slot.)
%
if notDefined('fieldName'),		fieldName = {'ph' 'ph'};			end

nAreas = length(areaCorners);
for ii=1:nAreas
    corners = areaCorners{ii};
    retPhases(ii,1) = atlasEstimateBoundaryPhases(vw, corners(1,:), ...
							corners(2,:), scanNum(1), fieldName{1});
    retPhases(ii,2) = atlasEstimateBoundaryPhases(vw, corners(3,:), ...
							corners(4,:), scanNum(1), fieldName{1});
    retPhases(ii,3) = atlasEstimateBoundaryPhases(vw, corners(1,:), ...
							corners(4,:), scanNum(2), fieldName{2});
    retPhases(ii,4) = atlasEstimateBoundaryPhases(vw, corners(2,:), ...
							corners(3,:), scanNum(2), fieldName{2});
end

if nAreas > 1, retPhases = meanPhase(retPhases); end

%% this part of the code seems to attempt to adjust the phase difference
%% between the stored data in the flat view, and real world units. A couple
%% comments / points of confusion here:
%% (1) It seems that things would be much simplified if the downstream code
%% to this function dealt exclusively with data already in real-world
%% units.
%% (2) for traveling-wave data, the functions polarAngle and eccentricity
%% map between phase and real-world units. Perhaps this should be used?
%% (3) for pRF data, the loaded data are already in real world units.
%% (4) what exactly do the retPhases do?
%% ras, 04/09.
prompt = {'Foveal', 'Peripheral', 'Lower Phase (angle map)'};
def = {num2str(retPhases(1)), num2str(retPhases(2)), num2str(retPhases(3))};
dlgTitle = 'Adjust retinal phase estimates';
lineNo = 1;
answer = inputdlg(prompt, dlgTitle, lineNo, def);

% adjust the retPhases based on the user response
angleShift = retPhases(3) - str2double(answer{3});
for ii=1:3, retPhases(ii) = str2double(answer{ii}); end

% We shift the fourth (UVM) by the same amount as the LVM
retPhases(4) = retPhases(4) - angleShift;

return;

%----------------------------------------------------
function meanPh = atlasEstimateBoundaryPhases(vw, p1, p2, scanNum, fieldName)
%
%   meanPh = atlasEstimateBoundaryPhases(vw, p1, p2, scanNum, fieldName)
%
% Author:  Wandell
% Purpose:
%     Estimate the average phase along a line between two points.  This
%     code is used to provide first estimates of the foveal, peripheral,
%     UVM and LVM phase in building atlases.
%
% Example:
%
%   scanNum = 1
%   retPhase(1)= atlasEstimateBoundaryPhases(vw,corners(1,:),corners(2,:),scanNum);
%
%
% ras 04/2009:  allows the field to be passed as a parameter -- not only
% confined to phase field, and doesn't require multiple different scans for
% polar angle / eccentricity. This is critical for use with pRF models.

curSlice = viewGet(vw, 'Current Slice');
[x, y] = findLinePoints([p1(1) p1(2)], [p2(1) p2(2)]);

newCoords = zeros(3,length(x));
newCoords(1,:) = y;
newCoords(2,:) = x;
newCoords(3,:) = curSlice*ones(1,length(x));

% Convert coords to canonical frame of reference
newCoords = curOri2CanOri(vw, newCoords);
ph = getCurDataROI(vw, fieldName, scanNum, newCoords);
if ~isequal( lower(fieldName), 'ph' )
	% put the data into the range [0 2*pi].
	% there may be a chance that it is not correct that the ph data span
	% this range -- for instance, if someone provides polar angle data in a
	% different field. I should deal with this contingency, but it might be
	% even better to have the downstream code deal with real-world-unit
	% data, and leave the conversion to these units up to the user before
	% entering this function (ras 04/09)...
	ph = normalize(ph, 0, 2*pi);
end
cxph = exp(sqrt(-1)*ph);

% We add pi to the output so that the variables run from [0,2pi] instead of
% from [-pi,pi].  This is consistent with mrLoadRet encoding of phase.
meanPh = angle(mean(cxph));

if meanPh < 0, meanPh = meanPh + 2*pi; end
if meanPh > 2*pi, meanPh = meanPh - 2*pi; end

return;


function pesq_mos= pesq_psychoacoustic_model (ref_data, ref_Nsamples, deg_data, ...
    deg_Nsamples )

global CALIBRATE Nfmax Nb Sl Sp
global nr_of_hz_bands_per_bark_band centre_of_band_bark
global width_of_band_hz centre_of_band_hz width_of_band_bark
global pow_dens_correction_factor abs_thresh_power
global Downsample SEARCHBUFFER DATAPADDING_MSECS Fs Nutterances
global Utt_Start Utt_End Utt_Delay NUMBER_OF_PSQM_FRAMES_PER_SYLLABE
global Fs Plot_Frame

% Plot_Frame= 75; % this is the frame whose spectrum will be plotted

FALSE= 0;
TRUE= 1;
NUMBER_OF_PSQM_FRAMES_PER_SYLLABE= 20;

maxNsamples = max (ref_Nsamples, deg_Nsamples);
Nf = Downsample * 8;
MAX_NUMBER_OF_BAD_INTERVALS = 1000;

start_frame_of_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
stop_frame_of_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
start_sample_of_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
stop_sample_of_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
number_of_samples_in_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
delay_in_samples_in_bad_interval= zeros( 1, MAX_NUMBER_OF_BAD_INTERVALS);
number_of_bad_intervals= 0;
there_is_a_bad_frame= FALSE;

Whanning= hann( Nf, 'periodic');
Whanning= Whanning';

D_POW_F = 2;
D_POW_S = 6;
D_POW_T = 2;
A_POW_F = 1;
A_POW_S = 6;
A_POW_T = 2;
D_WEIGHT= 0.1;
A_WEIGHT= 0.0309;

CRITERIUM_FOR_SILENCE_OF_5_SAMPLES = 500;
samples_to_skip_at_start = 0;
sum_of_5_samples= 0;
while ((sum_of_5_samples< CRITERIUM_FOR_SILENCE_OF_5_SAMPLES) ...
        && (samples_to_skip_at_start < maxNsamples / 2))
    sum_of_5_samples= sum( abs( ref_data( samples_to_skip_at_start...
        + SEARCHBUFFER * Downsample + 1: samples_to_skip_at_start...
        + SEARCHBUFFER * Downsample + 5)));

    if (sum_of_5_samples< CRITERIUM_FOR_SILENCE_OF_5_SAMPLES)
        samples_to_skip_at_start = samples_to_skip_at_start+ 1;
    end
end
% fprintf( 'samples_to_skip_at_start is %d\n', samples_to_skip_at_start);

samples_to_skip_at_end = 0;
sum_of_5_samples= 0;
while ((sum_of_5_samples< CRITERIUM_FOR_SILENCE_OF_5_SAMPLES) ...
        && (samples_to_skip_at_end < maxNsamples / 2))
    sum_of_5_samples= sum( abs( ref_data( maxNsamples - ...
        SEARCHBUFFER* Downsample + DATAPADDING_MSECS* (Fs/ 1000) ...
        - samples_to_skip_at_end - 4: maxNsamples - ...
        SEARCHBUFFER* Downsample + DATAPADDING_MSECS* (Fs/ 1000) ...
        - samples_to_skip_at_end)));
    if (sum_of_5_samples< CRITERIUM_FOR_SILENCE_OF_5_SAMPLES)
        samples_to_skip_at_end = samples_to_skip_at_end+ 1;
    end
end
% fprintf( 'samples_to_skip_at_end is %d\n', samples_to_skip_at_end);

start_frame = floor( samples_to_skip_at_start/ (Nf/ 2));
stop_frame = floor( (maxNsamples- 2* SEARCHBUFFER* Downsample ...
    + DATAPADDING_MSECS* (Fs/ 1000)- samples_to_skip_at_end) ...
    / (Nf/ 2))- 1;
% number of frames in speech data plus DATAPADDING_MSECS
% fprintf( 'start/end frame is %d/%d\n', start_frame, stop_frame);

D_disturbance= zeros( stop_frame+ 1, Nb);
DA_disturbance= zeros( stop_frame+ 1, Nb);

power_ref = pow_of (ref_data, SEARCHBUFFER* Downsample, ...
    maxNsamples- SEARCHBUFFER* Downsample+ DATAPADDING_MSECS* (Fs/ 1000),...
    maxNsamples- 2* SEARCHBUFFER* Downsample+ DATAPADDING_MSECS* (Fs/ 1000));
power_deg = pow_of (deg_data, SEARCHBUFFER * Downsample, ...
    maxNsamples- SEARCHBUFFER* Downsample+ DATAPADDING_MSECS* (Fs/ 1000),...
    maxNsamples- 2* SEARCHBUFFER* Downsample+ DATAPADDING_MSECS* (Fs/ 1000));
% fprintf( 'ref/deg power is %f/%f\n', power_ref, power_deg);

hz_spectrum_ref             = zeros( 1, Nf/ 2);
hz_spectrum_deg             = zeros( 1, Nf/ 2);
frame_is_bad                = zeros( 1, stop_frame + 1);
smeared_frame_is_bad        = zeros( 1, stop_frame + 1);
silent                      = zeros( 1, stop_frame + 1);

pitch_pow_dens_ref          = zeros( stop_frame + 1, Nb);
pitch_pow_dens_deg          = zeros( stop_frame + 1, Nb);

frame_was_skipped           = zeros( 1, stop_frame + 1);
frame_disturbance           = zeros( 1, stop_frame + 1);
frame_disturbance_asym_add  = zeros( 1, stop_frame + 1);

avg_pitch_pow_dens_ref      = zeros( 1, Nb);
avg_pitch_pow_dens_deg      = zeros( 1, Nb);
loudness_dens_ref           = zeros( 1, Nb);
loudness_dens_deg           = zeros( 1, Nb);
deadzone                    = zeros( 1, Nb);
disturbance_dens            = zeros( 1, Nb);
disturbance_dens_asym_add   = zeros( 1, Nb);

time_weight                 = zeros( 1, stop_frame + 1);
total_power_ref             = zeros( 1, stop_frame + 1);

% fid= fopen( 'tmp_mat.txt', 'wt');

for frame = 0: stop_frame
    start_sample_ref = 1+ SEARCHBUFFER * Downsample + frame* (Nf/ 2);
    hz_spectrum_ref= short_term_fft (Nf, ref_data, Whanning, ...
        start_sample_ref);

    utt = Nutterances;
    while ((utt >= 1) && ((Utt_Start(utt)- 1)* Downsample+ 1 ...
            > start_sample_ref))
        utt= utt - 1;
    end

    if (utt >= 1)
        delay = Utt_Delay(utt);
    else
        delay = Utt_Delay(1);
    end

    start_sample_deg = start_sample_ref + delay;

    if ((start_sample_deg > 0) && (start_sample_deg + Nf- 1 < ...
            maxNsamples+ DATAPADDING_MSECS* (Fs/ 1000)))
        hz_spectrum_deg= short_term_fft (Nf, deg_data, Whanning, ...
            start_sample_deg);
    else
        hz_spectrum_deg( 1: Nf/ 2)= 0;
    end

    pitch_pow_dens_ref( frame+ 1, :)= freq_warping (...
        hz_spectrum_ref, Nb, frame);
    %peak = maximum_of (pitch_pow_dens_ref, 0, Nb);
    pitch_pow_dens_deg( frame+ 1, :)= freq_warping (...
        hz_spectrum_deg, Nb, frame);

    total_audible_pow_ref = total_audible (frame, pitch_pow_dens_ref, 1E2);
    total_audible_pow_deg = total_audible (frame, pitch_pow_dens_deg, 1E2);
    silent(frame+ 1) = (total_audible_pow_ref < 1E7);

end
% fclose( fid);

avg_pitch_pow_dens_ref= time_avg_audible_of (stop_frame + 1, ...
    silent, pitch_pow_dens_ref, floor((maxNsamples- 2* SEARCHBUFFER* ...
    Downsample+ DATAPADDING_MSECS* (Fs/ 1000))/ (Nf / 2))- 1);
avg_pitch_pow_dens_deg= time_avg_audible_of (stop_frame + 1, ...
    silent, pitch_pow_dens_deg, floor((maxNsamples- 2* SEARCHBUFFER* ...
    Downsample+ DATAPADDING_MSECS* (Fs/ 1000))/ (Nf/ 2))- 1);

% fid= fopen( 'tmp_mat.txt', 'wt');
% fprintf( fid, '%f\n', avg_pitch_pow_dens_deg);
% fclose( fid);

if (CALIBRATE== 0)
    pitch_pow_dens_ref= freq_resp_compensation (stop_frame + 1, ...
        pitch_pow_dens_ref, avg_pitch_pow_dens_ref, ...
        avg_pitch_pow_dens_deg, 1000);
    if (Plot_Frame>= 0) % plot pitch_pow_dens_ref
        figure;
        subplot( 1, 2, 1);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            pitch_pow_dens_ref( Plot_Frame+ 1, :)));
        axis( [0 Fs/2 0 95]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'reference signal bark spectrum with frequency compensation');
        subplot( 1, 2, 2);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            pitch_pow_dens_deg( Plot_Frame+ 1, :)));
        axis( [0 Fs/2 0 95]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'degraded signal bark spectrum');
    end

end
% tmp1= pitch_pow_dens_ref';

MAX_SCALE = 5.0;
MIN_SCALE = 3e-4;
oldScale = 1;
THRESHOLD_BAD_FRAMES = 30;
for frame = 0: stop_frame

    total_audible_pow_ref = total_audible (frame, pitch_pow_dens_ref, 1);
    total_audible_pow_deg = total_audible (frame, pitch_pow_dens_deg, 1);
    total_power_ref (1+ frame) = total_audible_pow_ref;

    scale = (total_audible_pow_ref + 5e3)/ (total_audible_pow_deg + 5e3);
    if (frame > 0)
        scale = 0.2 * oldScale + 0.8 * scale;
    end
    oldScale = scale;

    if (scale > MAX_SCALE)
        scale = MAX_SCALE;
    elseif (scale < MIN_SCALE)
        scale = MIN_SCALE;
    end

    pitch_pow_dens_deg( 1+ frame, :) = ...
        pitch_pow_dens_deg( 1+ frame, :) * scale;

    if (frame== Plot_Frame)
        figure;
        subplot( 1, 2, 1);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            pitch_pow_dens_ref( Plot_Frame+ 1, :)));
        axis( [0 Fs/2 0 95]); %xlabel( 'Hz'); ylabel( 'Db');
        subplot( 1, 2, 2);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            pitch_pow_dens_deg( Plot_Frame+ 1, :)));
        axis( [0 Fs/2 0 95]); %xlabel( 'Hz'); ylabel( 'Db');
    end

    loudness_dens_ref = intensity_warping_of (frame, pitch_pow_dens_ref);
    loudness_dens_deg = intensity_warping_of (frame, pitch_pow_dens_deg);
    disturbance_dens = loudness_dens_deg - loudness_dens_ref;

    if (frame== Plot_Frame)
        figure;
        subplot( 1, 2, 1);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            loudness_dens_ref));
        axis( [0 Fs/2 0 15]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'reference signal loudness density');
        subplot( 1, 2, 2);
        plot( centre_of_band_hz, 10* log10( eps+ ...
            loudness_dens_deg));
        axis( [0 Fs/2 0 15]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'degraded signal loudness density');
    end

    for band =1: Nb
        deadzone (band) = 0.25* min (loudness_dens_deg (band), ...
            loudness_dens_ref (band));
    end

    for band = 1: Nb
        d = disturbance_dens (band);
        m = deadzone (band);

        if (d > m)
            disturbance_dens (band) = disturbance_dens (band)- m;
%             disturbance_dens (band) = d- m;
        else
            if (d < -m)
                disturbance_dens (band) = disturbance_dens (band)+ m;
%                 disturbance_dens (band) = d+ m;
            else
                disturbance_dens (band) = 0;
            end
        end
    end

    if (frame== Plot_Frame)
        figure;
        subplot( 1, 2, 1);
        plot( centre_of_band_hz, disturbance_dens);
        axis( [0 Fs/2 -1 50]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'disturbance');
    end
    D_disturbance( frame+ 1, :)= disturbance_dens;

    frame_disturbance (1+ frame) = pseudo_Lp (disturbance_dens, D_POW_F);
    if (frame_disturbance (1+ frame) > THRESHOLD_BAD_FRAMES)
        there_is_a_bad_frame = TRUE;
    end

    disturbance_dens= multiply_with_asymmetry_factor (...
        disturbance_dens, frame, pitch_pow_dens_ref, pitch_pow_dens_deg);

    if (frame== Plot_Frame)
        subplot( 1, 2, 2);
        plot( centre_of_band_hz, disturbance_dens);
        axis( [0 Fs/2 -1 50]); %xlabel( 'Hz'); ylabel( 'Db');
        title( 'disturbance after asymmetry processing');
    end
    DA_disturbance( frame+ 1, :)= disturbance_dens;

    frame_disturbance_asym_add (1+ frame) = ...
        pseudo_Lp (disturbance_dens, A_POW_F);
end
% fid= fopen( 'tmp_mat.txt', 'wt');
% fprintf( fid, '%f\n', frame_disturbance);
% fclose( fid);

frame_was_skipped (1: 1+ stop_frame) = FALSE;

for utt = 2: Nutterances
    frame1 = floor (((Utt_Start(utt)- 1- SEARCHBUFFER )* Downsample+ 1+ ...
        Utt_Delay(utt))/ (Nf/ 2));
    j = floor( floor(((Utt_End(utt-1)- 1- SEARCHBUFFER)* Downsample+ 1+ ...
        Utt_Delay(utt-1)))/(Nf/ 2));
    delay_jump = Utt_Delay(utt) - Utt_Delay(utt-1);
    if (frame1 > j)
        frame1 = j;
    elseif (frame1 < 0)
        frame1 = 0;
    end
%     fprintf( 'frame1, j, delay_jump is %d, %d, %d\n', frame1, ...
%         j, delay_jump);

    if (delay_jump < -(Nf/ 2))
        frame2 = floor (((Utt_Start(utt)- 1- SEARCHBUFFER)* Downsample+ 1 ...
            + max (0, abs (delay_jump)))/ (Nf/ 2)) + 1;

        for frame = frame1: frame2
            if (frame < stop_frame)
                frame_was_skipped (1+ frame) = TRUE;
                frame_disturbance (1+ frame) = 0;
                frame_disturbance_asym_add (1+ frame) = 0;
            end
        end
    end
end

nn = DATAPADDING_MSECS* (Fs/ 1000) + maxNsamples;
tweaked_deg = zeros( 1, nn);
% fprintf( 'nn is %d\n', nn);

for i= SEARCHBUFFER* Downsample+ 1: nn- SEARCHBUFFER* Downsample
    utt = Nutterances;

    while ((utt >= 1) && ((Utt_Start (utt)- 1)* Downsample> i))
        utt = utt- 1;
    end
    if (utt >= 1)
        delay = Utt_Delay (utt);
    else
        delay = Utt_Delay (1);
    end

    j = i + delay;
    if (j < SEARCHBUFFER * Downsample+ 1)
        j = SEARCHBUFFER * Downsample+ 1;
    end
    if (j > nn - SEARCHBUFFER * Downsample)
        j = nn - SEARCHBUFFER * Downsample;
    end
    tweaked_deg (i) = deg_data (j);
end

if (there_is_a_bad_frame)

    for frame = 0: stop_frame
        frame_is_bad (1+ frame) = (frame_disturbance (1+ frame)...
            > THRESHOLD_BAD_FRAMES);
        smeared_frame_is_bad (1+ frame) = FALSE;
    end
    frame_is_bad (1) = FALSE;
    SMEAR_RANGE = 2;

    for frame = SMEAR_RANGE: stop_frame- 1- SMEAR_RANGE
        max_itself_and_left = frame_is_bad (1+ frame);
        max_itself_and_right = frame_is_bad (1+ frame);

        for i = -SMEAR_RANGE: 0
            if (max_itself_and_left < frame_is_bad (1+ frame+ i))
                max_itself_and_left = frame_is_bad (1+ frame+ i);
            end
        end

        for i = 0: SMEAR_RANGE
            if (max_itself_and_right < frame_is_bad (1+ frame + i))
                max_itself_and_right = frame_is_bad (1+ frame + i);
            end
        end

        mini = max_itself_and_left;
        if (mini > max_itself_and_right)
            mini = max_itself_and_right;
        end

        smeared_frame_is_bad (1+ frame) = mini;
    end

    MINIMUM_NUMBER_OF_BAD_FRAMES_IN_BAD_INTERVAL = 5;
    number_of_bad_intervals = 0;
    frame = 0;
    while (frame <= stop_frame)
        while ((frame <= stop_frame) && (~smeared_frame_is_bad (1+ frame)))
            frame= frame+ 1;
        end

        if (frame <= stop_frame)
            start_frame_of_bad_interval(1+ number_of_bad_intervals)= ...
                1+ frame;

            while ((frame <= stop_frame) && (...
                    smeared_frame_is_bad (1+ frame)))
                frame= frame+ 1;
            end

            if (frame <= stop_frame)
                stop_frame_of_bad_interval(1+ number_of_bad_intervals)= ...
                    1+ frame;
                if (stop_frame_of_bad_interval(1+ number_of_bad_intervals)- ...
                        start_frame_of_bad_interval(1+ number_of_bad_intervals)...
                        >= MINIMUM_NUMBER_OF_BAD_FRAMES_IN_BAD_INTERVAL)
                    number_of_bad_intervals= number_of_bad_intervals+ 1;
                end
            end
        end
    end

    for bad_interval = 0: number_of_bad_intervals - 1
        start_sample_of_bad_interval(1+ bad_interval) = ...
            (start_frame_of_bad_interval(1+ bad_interval)- 1) * (Nf/ 2) ...
            + SEARCHBUFFER * Downsample+ 1;
        stop_sample_of_bad_interval(1+ bad_interval) = ...
            (stop_frame_of_bad_interval(1+ bad_interval)- 1) * (Nf/ 2) ...
            + Nf + SEARCHBUFFER* Downsample;
        if (stop_frame_of_bad_interval(1+ bad_interval) > stop_frame+ 1)
            stop_frame_of_bad_interval(1+ bad_interval) = stop_frame+ 1;
        end

        number_of_samples_in_bad_interval(1+ bad_interval) = ...
            stop_sample_of_bad_interval(1+ bad_interval) - ...
            start_sample_of_bad_interval(1+ bad_interval)+ 1;
    end
%     fprintf( 'number of bad intervals %d\n', number_of_bad_intervals);
%     fprintf( '%d %d\n', number_of_samples_in_bad_interval(1), ...
%         number_of_samples_in_bad_interval(2));
%     fprintf( '%d %d\n', start_sample_of_bad_interval(1), ...
%         start_sample_of_bad_interval(2));

    SEARCH_RANGE_IN_TRANSFORM_LENGTH = 4;
    search_range_in_samples= SEARCH_RANGE_IN_TRANSFORM_LENGTH * Nf;

    for bad_interval= 0: number_of_bad_intervals- 1
        ref = zeros (1, 2 * search_range_in_samples + ...
            number_of_samples_in_bad_interval (1+ bad_interval));
        deg = zeros (1, 2 * search_range_in_samples + ...
            number_of_samples_in_bad_interval (1+ bad_interval));

        ref(1: search_range_in_samples) = 0;

        ref (search_range_in_samples+ 1: search_range_in_samples+ ...
                number_of_samples_in_bad_interval (1+ bad_interval)) = ...
                ref_data (start_sample_of_bad_interval( 1+ bad_interval) + 1: ...
                start_sample_of_bad_interval( 1+ bad_interval) + ...
                number_of_samples_in_bad_interval (1+ bad_interval));

        ref (search_range_in_samples + ...
                number_of_samples_in_bad_interval (1+ bad_interval) + 1: ...
                search_range_in_samples + ...
                number_of_samples_in_bad_interval (1+ bad_interval) + ...
                search_range_in_samples) = 0;

        for i = 0: 2 * search_range_in_samples + ...
                number_of_samples_in_bad_interval (1+ bad_interval) - 1
            j = start_sample_of_bad_interval (1+ bad_interval) - ...
                search_range_in_samples + i;
            nn = maxNsamples - SEARCHBUFFER * Downsample + ...
                DATAPADDING_MSECS  * (Fs / 1000);
            if (j <= SEARCHBUFFER * Downsample)
                j = SEARCHBUFFER * Downsample+ 1;
            end
            if (j > nn)
                j = nn;
            end
            deg (1+ i) = tweaked_deg (j);
        end

        [delay_in_samples, best_correlation]= compute_delay ...
            (1, 2 * search_range_in_samples + ...
            number_of_samples_in_bad_interval (1+ bad_interval), ...
            search_range_in_samples, ref, deg);
        delay_in_samples_in_bad_interval (1+ bad_interval) =  ...
            delay_in_samples;
%         fprintf( 'delay_in_samples, best_correlation is \n\t%d, %f\n', ...
%             delay_in_samples, best_correlation);
%
        if (best_correlation < 0.5)
            delay_in_samples_in_bad_interval  (1+ bad_interval) = 0;
        end
    end

    if (number_of_bad_intervals > 0)
        doubly_tweaked_deg = tweaked_deg( 1: maxNsamples + ...
            DATAPADDING_MSECS  * (Fs / 1000));
        for bad_interval= 0: number_of_bad_intervals- 1
            delay = delay_in_samples_in_bad_interval (1+ bad_interval);

            for i = start_sample_of_bad_interval (1+ bad_interval): ...
                    stop_sample_of_bad_interval (1+ bad_interval)
                j = i + delay;
                if (j < 1)
                    j = 1;
                end
                if (j > maxNsamples)
                    j = maxNsamples;
                end
                h = tweaked_deg (j);
                doubly_tweaked_deg (i) = h;
            end
        end

        untweaked_deg = deg_data;
        deg_data = doubly_tweaked_deg;

        for bad_interval= 0: number_of_bad_intervals- 1
            for frame = start_frame_of_bad_interval (1+ bad_interval): ...
                    stop_frame_of_bad_interval (1+ bad_interval)- 1
                frame= frame- 1;
                start_sample_ref = SEARCHBUFFER * Downsample + ...
                    frame * Nf / 2+ 1;
                start_sample_deg = start_sample_ref;
                hz_spectrum_deg= short_term_fft (Nf, deg_data, ...
                    Whanning, start_sample_deg);
                pitch_pow_dens_deg( 1+ frame, :)= freq_warping (...
                    hz_spectrum_deg, Nb, frame);
            end

            oldScale = 1;
            for frame = start_frame_of_bad_interval (1+ bad_interval): ...
                    stop_frame_of_bad_interval (1+ bad_interval)- 1
                frame= frame- 1;
                % see implementation for detail why 1 needed to be
                % subtracted
                total_audible_pow_ref = total_audible (frame, ...
                    pitch_pow_dens_ref, 1);
                total_audible_pow_deg = total_audible (frame, ...
                    pitch_pow_dens_deg, 1);
                scale = (total_audible_pow_ref + 5e3) / ...
                    (total_audible_pow_deg + 5e3);
                if (frame > 0)
                    scale = 0.2 * oldScale + 0.8*scale;
                end
                oldScale = scale;
                if (scale > MAX_SCALE)
                    scale = MAX_SCALE;
                end
                if (scale < MIN_SCALE)
                    scale = MIN_SCALE;
                end

                pitch_pow_dens_deg (1+ frame, :) = ...
                    pitch_pow_dens_deg (1+ frame, :)* scale;
                loudness_dens_ref= intensity_warping_of (frame, ...
                    pitch_pow_dens_ref);
                loudness_dens_deg= intensity_warping_of (frame, ...
                    pitch_pow_dens_deg);
                disturbance_dens = loudness_dens_deg - loudness_dens_ref;

                for band = 1: Nb
                    deadzone(band) = min (loudness_dens_deg(band), ...
                        loudness_dens_ref(band));
                    deadzone(band) = deadzone(band)* 0.25;
                end

                for band = 1: Nb
                    d = disturbance_dens (band);
                    m = deadzone (band);

                    if (d > m)
                        disturbance_dens (band) = ...
                            disturbance_dens (band)- m;
                    else
                        if (d < -m)
                            disturbance_dens (band) = ...
                                disturbance_dens (band)+ m;
                        else
                            disturbance_dens (band) = 0;
                        end
                    end
                end

                frame_disturbance( 1+ frame) = min (...
                    frame_disturbance( 1+ frame), pseudo_Lp(...
                    disturbance_dens, D_POW_F));
                disturbance_dens= multiply_with_asymmetry_factor ...
                    (disturbance_dens, frame, pitch_pow_dens_ref, ...
                    pitch_pow_dens_deg);
                frame_disturbance_asym_add(1+ frame) = min (...
                    frame_disturbance_asym_add(1+ frame), ...
                    pseudo_Lp (disturbance_dens, A_POW_F));
            end
        end
        deg_data = untweaked_deg;
    end
end

for frame = 0: stop_frame
    h = 1;
    if (stop_frame + 1 > 1000)
        n = floor( (maxNsamples - 2 * SEARCHBUFFER * Downsample)...
            / (Nf / 2)) - 1;
        timeWeightFactor = (n - 1000) / 5500;
        if (timeWeightFactor > 0.5)
            timeWeightFactor = 0.5;
        end
        h = (1.0 - timeWeightFactor) + timeWeightFactor * frame / n;
    end

    time_weight (1 +frame) = h;
end

% fid= fopen( 'tmp_mat1.txt', 'at');
% fprintf( '\n');
for frame = 0: stop_frame
    h = ((total_power_ref (1+ frame) + 1e5) / 1e7)^ 0.04;
%     if (frame== 118)
%         fprintf( '%f\n', h);
%         fprintf( '%f\n', frame_disturbance( 1+ frame));
%     end
    frame_disturbance( 1+ frame) = frame_disturbance( 1+ frame)/ h;

%     if (frame== 118)
%         fprintf( '%f\n', frame_disturbance( 1+ frame));
%     end
%
    frame_disturbance_asym_add( 1+ frame) = ...
        frame_disturbance_asym_add( 1+ frame)/ h;
    if (frame_disturbance( 1+ frame) > 45)
        frame_disturbance( 1+ frame) = 45;
    end
    if (frame_disturbance_asym_add( 1+ frame)> 45)
        frame_disturbance_asym_add( 1+ frame) = 45;
    end
end
% fclose ( fid);

d_indicator = Lpq_weight (start_frame, stop_frame, ...
    D_POW_S, D_POW_T, frame_disturbance, time_weight);
a_indicator = Lpq_weight (start_frame, stop_frame, ...
    A_POW_S, A_POW_T, frame_disturbance_asym_add, time_weight);

pesq_mos = 4.5 - D_WEIGHT * d_indicator - A_WEIGHT * a_indicator;

if (Plot_Frame> 0)
    figure;
    subplot( 1, 2, 1);
    mesh( 0: stop_frame, centre_of_band_hz, D_disturbance');
    title( 'disturbance');
    subplot( 1, 2, 2);
    mesh( 0: stop_frame, centre_of_band_hz, DA_disturbance');
    title( 'disturbance after asymmetry processing');
end

% fid= fopen( 'tmp_mat.txt', 'wt');
% fprintf( fid, 'time_weight\n');
% fprintf( fid, '%f\n', time_weight);
% fprintf( fid, 'frame_disturbance:\n');
% fprintf( fid, '%f\n', frame_disturbance);
% fprintf( fid, 'frame_disturbance_asym_add\n');
% fprintf( fid, '%f\n', frame_disturbance_asym_add);
% fclose( fid);

function result_time= Lpq_weight(start_frame, stop_frame, ...
        power_syllable, power_time, frame_disturbance, time_weight)

global NUMBER_OF_PSQM_FRAMES_PER_SYLLABE

% fid= fopen( 'tmp_mat1.txt', 'at');
% fprintf( 'result_time:\n');

result_time= 0;
total_time_weight_time = 0;
% fprintf( 'start/end frame: %d/%d\n', start_frame, stop_frame);
for start_frame_of_syllable = start_frame: ...
        NUMBER_OF_PSQM_FRAMES_PER_SYLLABE/2: stop_frame
    result_syllable = 0;
    count_syllable = 0;

    for frame = start_frame_of_syllable: ...
            start_frame_of_syllable + NUMBER_OF_PSQM_FRAMES_PER_SYLLABE- 1
        if (frame <= stop_frame)
            h = frame_disturbance(1+ frame);
%             if (start_frame_of_syllable== 101)
%                 fprintf( fid, '%f\n', h);
%             end
            result_syllable = result_syllable+ (h^ power_syllable);
        end
        count_syllable = count_syllable+ 1;
    end

    result_syllable = result_syllable/ count_syllable;
    result_syllable = result_syllable^ (1/power_syllable);

    result_time= result_time+ (time_weight (...
        1+ start_frame_of_syllable - start_frame) * ...
        result_syllable)^ power_time;
    total_time_weight_time = total_time_weight_time+ ...
        time_weight (1+ start_frame_of_syllable - start_frame)^ power_time;

%     fprintf( fid, '%f\n', result_time);
end
% fclose (fid);

% fprintf( 'total_time_weight_time is %f\n', total_time_weight_time);
result_time = result_time/ total_time_weight_time;
result_time= result_time^ (1/ power_time);
% fprintf( 'result_time is %f\n\n', result_time);

function [best_delay, max_correlation] = compute_delay (...
    start_sample, stop_sample, search_range, ...
    time_series1, time_series2)

n = stop_sample - start_sample+ 1;
power_of_2 = 2^ (ceil( log2( 2 * n)));

power1 = pow_of (time_series1, start_sample, stop_sample, n)* ...
    n/ power_of_2;
power2 = pow_of (time_series2, start_sample, stop_sample, n)* ...
    n/ power_of_2;
normalization = sqrt (power1 * power2);
% fprintf( 'normalization is %f\n', normalization);

if ((power1 <= 1e-6) || (power2 <= 1e-6))
    max_correlation = 0;
    best_delay= 0;
end

x1( 1: power_of_2)= 0;
x2( 1: power_of_2)= 0;
y( 1: power_of_2)= 0;

x1( 1: n)= abs( time_series1( start_sample: ...
    stop_sample));
x2( 1: n)= abs( time_series2( start_sample: ...
    stop_sample));

x1_fft= fft( x1, power_of_2)/ power_of_2;
x2_fft= fft( x2, power_of_2);
x1_fft_conj= conj( x1_fft);
y= ifft( x1_fft_conj.* x2_fft, power_of_2);

best_delay = 0;
max_correlation = 0;

% these loop can be rewritten
for i = -search_range: -1
    h = abs (y (1+ i + power_of_2)) / normalization;
    if (h > max_correlation)
        max_correlation = h;
        best_delay= i;
    end
end
for i = 0: search_range- 1
    h = abs (y (1+i)) / normalization;
    if (h > max_correlation)
        max_correlation = h;
        best_delay= i;
    end
end
best_delay= best_delay- 1;

function mod_disturbance_dens= multiply_with_asymmetry_factor (...
    disturbance_dens, frame, pitch_pow_dens_ref, pitch_pow_dens_deg)

global Nb
for i = 1: Nb
    ratio = (pitch_pow_dens_deg(1+ frame, i) + 50)...
        / (pitch_pow_dens_ref (1+ frame, i) + 50);
    h = ratio^ 1.2;
    if (h > 12)
        h = 12;
    elseif (h < 3)
        h = 0.0;
    end
    mod_disturbance_dens (i) = disturbance_dens (i) * h;
end

function loudness_dens = intensity_warping_of (...
    frame, pitch_pow_dens)

global abs_thresh_power Sl Nb centre_of_band_bark
ZWICKER_POWER= 0.23;
for band = 1: Nb
    threshold = abs_thresh_power (band);
    input = pitch_pow_dens (1+ frame, band);

    if (centre_of_band_bark (band) < 4)
        h =  6 / (centre_of_band_bark (band) + 2);
    else
        h = 1;
    end

    if (h > 2)
        h = 2;
    end
    h = h^ 0.15;
    modified_zwicker_power = ZWICKER_POWER * h;
    if (input > threshold)
        loudness_dens (band) = ((threshold / 0.5)^ modified_zwicker_power)...
            * ((0.5 + 0.5 * input / threshold)^ modified_zwicker_power- 1);
    else
        loudness_dens (band) = 0;
    end

    loudness_dens (band) = loudness_dens (band)* Sl;
end

function result= pseudo_Lp (x, p)

global Nb width_of_band_bark
totalWeight = 0;
result = 0;
for band = 2: Nb
    h = abs (x (band));
    w = width_of_band_bark (band);
    prod = h * w;

    result = result+ prod^ p;
    totalWeight = totalWeight+ w;
end
result = (result/ totalWeight)^ (1/p);
result = result* totalWeight;

function mod_pitch_pow_dens_ref= freq_resp_compensation (number_of_frames, ...
    pitch_pow_dens_ref, avg_pitch_pow_dens_ref, ...
    avg_pitch_pow_dens_deg, constant)

global Nb

for band = 1: Nb
    x = (avg_pitch_pow_dens_deg (band) + constant) / ...
        (avg_pitch_pow_dens_ref (band) + constant);
    if (x > 100.0)
        x = 100.0;
    elseif (x < 0.01)
        x = 0.01;
    end

    for frame = 1: number_of_frames
        mod_pitch_pow_dens_ref(frame, band) = ...
            pitch_pow_dens_ref(frame, band) * x;
    end
end

function avg_pitch_pow_dens= time_avg_audible_of(number_of_frames, ...
    silent, pitch_pow_dens, total_number_of_frames)

global Nb abs_thresh_power

for band = 1: Nb
    result = 0;
    for frame = 1: number_of_frames
        if (~silent (frame))
            h = pitch_pow_dens (frame, band);
            if (h > 100 * abs_thresh_power (band))
                result = result + h;
            end
        end

        avg_pitch_pow_dens (band) = result/ total_number_of_frames;
    end
end

function hz_spectrum= short_term_fft (Nf, data, Whanning, start_sample)

x1= data( start_sample: start_sample+ Nf-1).* Whanning;
x1_fft= fft( x1);
hz_spectrum= abs( x1_fft( 1: Nf/ 2)).^ 2;
hz_spectrum( 1)= 0;

function pitch_pow_dens= freq_warping( hz_spectrum, Nb, frame)

global nr_of_hz_bands_per_bark_band pow_dens_correction_factor
global Sp

hz_band = 1;
for bark_band = 1: Nb
    n = nr_of_hz_bands_per_bark_band (bark_band);
    sum = 0;
    for i = 1: n
        sum = sum+ hz_spectrum( hz_band);
        hz_band= hz_band+ 1;
    end
    sum = sum* pow_dens_correction_factor (bark_band);
    sum = sum* Sp;
    pitch_pow_dens (bark_band) = sum;

end

function total_audible_pow = total_audible (frame, ...
    pitch_pow_dens, factor)

global Nb abs_thresh_power

total_audible_pow = 0;
for band= 2: Nb
    h = pitch_pow_dens (frame+ 1,band);
    threshold = factor * abs_thresh_power (band);
    if (h > threshold)
        total_audible_pow = total_audible_pow+ h;
    end
end


function [caNodeIndices, vResolution_] = ex_CreateIndexCatalog3D(mCatalog, mPolygon, bMap, nGriddingMode, fSmpValue, fSmpBnd, fSizeRectHorizontal, fSizeRectDepth)
    % Creates a cell-array with subcatalogs for every grid node defined by mPolygon.
    %
    % [caNodeIndices] = ex_CreateIndexCatalog(mCatalog, mPolygon, bMap, nGriddingMode,
    %                                                  nNumberEvents, fRadius, fSizeRectHorizontal, fSizeRectDepth)
    % -------------------------------------------------------------------------------------------------------------
    % Creates a cell-array with subcatalogs for every grid node defined by mPolygon. These subcatalogs
    %   contain only indices to the earthquake "rows" in mCatalog.
    %
    % Input parameters:
    %   mCatalog              Earthquake catalog
    %   mPolygon              Polygon (defined by ex_selectgrid)
    %   bMap                  Calculate cell-array for a map (true) or a cross-section (false)
    %   nGriddingMode         Mode of creating grid node subcatalogs
    %                         0: Constant number of events
    %                         1: Constant radius
    %                         2: Rectangular grid node samples
    %                         3: Spherical grid node samples with constant no
    %                            of events
    %                         4: Spherical grid node samples with constant radius
    %   nNumberEvents         Number of events per grid node (nGriddingMode == 0)
    %   fRadius               Radius of grid node sample (nGriddingMode == 1)
    %   fSizeRectHorizontal   Latitude/horizontal size of rectangle (nGriddingMode == 2)
    %   fSizeRectDepth        Longitude/depth size of rectangle (nGriddingMode == 2)
    %
    % Output parameters:
    %   caNodeIndices         Cell-array with index-catalogs per grid node of mPolygon
    %
    % Danijel Schorlemmer
    % June 17, 2002

    report_this_filefun();

    % Create the catalogs for each node with pointers to the overall catalog
    nNumberNodes_ = length(mPolygon(:,1));
    caNodeIndices = cell(nNumberNodes_, 1);

    % If cross-section calculate the length along cross-section
    if ~bMap
        nRow_ = mCatalog.Count;
        vXSecX_ = mCatalog(:,nColumn_);  % length along x-section
        vXSecY_ = (-1) * mCatalog.Depth;  % depth of hypocenters
    end

    % vResolution give the radius (for nGriddingMode = 4) and no. of events
    % (for nGriddingMode = 3)
    vResolution_(:,1)=nan(nNumberNodes_,1)
    if ((nGriddingMode == 3) || (nGriddingMode ==4)) clear vResolution_; end

    % Loop over all points of the polygon
    for nNode_ = 1:nNumberNodes_
        % Get the grid node coordinates
        fX_ = mPolygon(nNode_, 1);
        fY_ = mPolygon(nNode_, 2);
        if size(mPolygon,2)==3
            fZ_ = mPolygon(nNode_, 3);
        end

        if (nGriddingMode == 0) | (nGriddingMode == 1)  % Fixed radius or fixed number
            % Calculate distance from center point
            if bMap
                vDistances_ = sqrt(((mCatalog.Longitude-fX_)*cosd(fY_)*111).^2 + ((mCatalog.Latitude-fY_)*111).^2);
            else
                vDistances_ = sqrt(((vXSecX_ - fX_)).^2 + ((vXSecY_ - fY_)).^2);
            end
            if nGriddingMode == 0 % Fixed number
                if mCatalog.Count == 0
                    caNodeIndices{nNode_} = [];
                    % NaN for no events
                    vResolution_(nNode_) = nan;
                elseif nNumberEvents > mCatalog.Count
                    caNodeIndices{nNode_} = vIndices(1:mCatalog.Count);
                    % take the maximal distance for all eq. in the catalog
                    vResolution_(nNode_) = max(vdistances_);
                else
                    % Use first nNumberEvents events
                    [vTmp, vIndices] = sort(vDistances_);
                    caNodeIndices{nNode_} = vIndices(1:nNumberEvents);
                    % radius of the nNumberEvents-th event in the sorted vDistances_
                    vResolution_(nNode_) = vTmp(nNumberEvents);
                end
            else % Fixed radius
                % Use all events within fRadius
                caNodeIndices{nNode_} = find(vDistances_ <= fRadius);
                vResolution_(nNode_) = length(find(vDistances_ <= fRadius));
            end
        elseif nGriddingMode ==2  % Rectangular gridding (nGriddingMode == 2)
            if bMap
                vSel_ = ((mCatalog.Longitude >= (fX_ - fSizeRectHorizontal/2)) & (mCatalog.Longitude < (fX_ + fSizeRectHorizontal/2)) & ...
                    (mCatalog.Latitude >= (fY_ - fSizeRectDepth/2)) & (mCatalog.Latitude < (fY_ + fSizeRectDepth/2)));
                vResolution_(nNode_) = length(find(vSel_ > 0))
            else
                vSel_ = ((vXSecX_ >= (fX_ - fSizeRectHorizontal/2)) & (vXSecX_ < (fX_ + fSizeRectHorizontal/2)) & ...
                    (vXSecY_ >= (fY_ - fSizRectDepth/2)) & (vXSecY_ < (fY_ + fSizeRectDepth/2)));
                vResolution_(nNode_) = length(find(vSel_ > 0))
            end
            caNodeIndices{nNode_} = find(vSel_ == 1);
        elseif ((nGriddingMode == 3) || (nGriddingMode == 4))     % Spherical grid node samples with constant no of events
            vDistances_ = sqrt(((mCatalog.Longitude-fX_)*cosd(fY_)*111).^2 + ((mCatalog.Latitude-fY_)*111).^2 + (mCatalog.Depth-fZ_).^2);
            if nGriddingMode == 3  % Spherical grid node samples with constant no of events
                nNumberEvents=fSmpValue;
                fRadius=fSmpBnd;
                if mCatalog.Count == 0
                    caNodeIndices{nNode_} = [];
                    % NaN for no events
                    vResolution_{nNode_} = nan;
                    %               vResolution_(nNode_) = nan;
                elseif nNumberEvents > mCatalog.Count
                    caNodeIndices{nNode_} = vIndices(1:mCatalog.Count);
                    % take the maximal distance for all eq. in the catalog
                    vResolution_{nNode_} = vdistances_;
                    %               vResolution_(nNode_) = max(vdistances_);
                else
                    % Use first nNumberEvents events
                    [vTmp, vIndices] = sort(vDistances_);
                    caNodeIndices{nNode_} = vIndices(1:nNumberEvents);
                    % radius of the nNumberEvents-th event in the sorted vDistances_
                    vResolution_{nNode_} = vTmp(nNumberEvents);
                    %               vResolution_(nNode_) = vTmp(nNumberEvents);
                end
            elseif nGriddingMode == 4     % Spherical grid node samples with constant radius
                nNumberEvents=fSmpBnd;
                fRadius=fSmpValue;
                %  muss noch gemacht werden....

            end
        end
    end % of for nNode_
end


% This file is part of TREEQSM.
%
% TREEQSM is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% TREEQSM is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with TREEQSM.  If not, see <http://www.gnu.org/licenses/>.

function RS = relative_size(P,cover,segment)

% ---------------------------------------------------------------------
% RELATIVE_SIZE.M   Determines relative cover set size for points in new covers
%
% Version 2.00
% Latest update     16 Aug 2017
%
% Copyright (C) 2014-2017 Pasi Raumonen
% ---------------------------------------------------------------------
%
% Uses existing segmentation and its branching structure to determine
% relative size of the cover sets distributed over new covers. The idea is
% to decrease the relative size as the branch size decreases. This is
% realised so that the relative size at the base of a branch is
% proportional to the size of the stem's base, measured as number of
% cover sets in the first few layers. Also when we approach the
% tip of the branch, the relative size decreases to the minimum.
% Maximum relative size is 256 at the bottom of the
% stem and the minimum is 1 at the tip of every branch.
%
% Output:
% RS    Relative size (1-256), uint8-vector, (n_points x 1)

Bal = cover.ball;
Cen = cover.center;
Nei = cover.neighbor;
Segs = segment.segments;
SChi = segment.ChildSegment;
np = size(P,1);     % number of points
ns = size(Segs,1);  % number of segments

%% Use branching order and height as apriori info
% Determine the branch orders of the segments
Ord = zeros(ns,1);
C = SChi{1};
order = 0;
while ~isempty(C)
    order = order+order;
    Ord(C) = order;
    C = vertcat(SChi{C});
end
maxO = order+1; % maximum branching order (plus one)

% Determine tree height
Top = max(P(Cen,3));
Bot = min(P(Cen,3));
H = Top-Bot;

%% Determine "base size" compared to the stem base
% BaseSize is the relative size of the branch base compared to the stem
% base, measured as number of cover sets in the first layers of the cover
% sets. If it is larger than apriori upper limit based on branching order
% and branch height, then correct to the apriori limit
BaseSize = zeros(ns,1);
% Determine first the base size at the stem
S = Segs{1};
n = size(S,1);
if n >= 2
    m = min([6 n]);
    BaseSize(1) = mean(cellfun(@length,S(2:m)));
else
    BaseSize(1) = length(S{1});
end
% Then define base size for other segments
for i = 2:ns
    S = Segs{i};
    n = size(S,1);
    if n >= 2
        m = min([6 n]);
        BaseSize(i) = ceil(mean(cellfun(@length,S(2:m)))/BaseSize(1)*256);
    else
        BaseSize(i) = length(S{1})/BaseSize(1)*256;
    end
    bot = min(P(Cen(S{1}),3));
    h = bot-Bot; % height of the segment's base
    BS = ceil(256*(maxO-Ord(i))/maxO*(H-h)/H); % maximum apriori base size
    if BaseSize(i) > BS
        BaseSize(i) = BS;
    end
end
BaseSize(1) = 256;

%% Determine relative size for points
TS = 1;
RS = zeros(np,1,'uint8');
for i = 1:ns
    S = Segs{i};
    s = size(S,1);
    for j = 1:s
        Q = S{j};
        RS(vertcat(Bal{Q})) = BaseSize(i)-(BaseSize(i)-TS)*sqrt((j-1)/s);
    end
end

%% Adjust the relative size at the base of child segments
RS0 = RS;
for i = 1:ns
    C = SChi{i};
    n = length(C);
    if n > 0
        for j = 1:n
            S = Segs{C(j)};
            B = S{1};
            N = vertcat(Nei{B});
            if size(S,1) > 1
                N = setdiff(N,S{2});
            end
            N = union(N,B);
            N = vertcat(Bal{N});
            RS(N) = RS0(N)/2;
        end
    end
end