function B=inpaint_nans_bc(A,method,bcclass)
% INPAINT_NANS_BC: in-paints over nans in an array, with spherical or toroidal boundary conditions
% usage: B=inpaint_nsns_bc(A)          % default method
% usage: B=inpaint_nsns_bc(A,method)   % specify method used
% usage: B=inpaint_nsns_bc(A,method,bcclass)   % specify class of boundary conditions applied
%
% Solves approximation to one of several pdes to
% interpolate and extrapolate holes in an array.
% Depending upon the boundary conditions specified,
% the array will effectively be treated as if it lies
% on either the surface of a sphere or a toroid.
%
% arguments (input):
%  A - nxm array with some NaNs to be filled in
%
%  method - (OPTIONAL) scalar numeric flag - specifies
%     which approach (or physical metaphor to use
%     for the interpolation.) All methods are capable
%     of extrapolation, some are better than others.
%     There are also speed differences, as well as
%     accuracy differences for smooth surfaces.
%
%     The methods employed here are a subset of the
%     methods of the original inpaint_nans.
%
%     methods {0,1} use a simple plate metaphor.
%     method  4 uses a spring metaphor.
%
%     method == 0 --> (DEFAULT) see method 1, but
%       this method does not build as large of a
%       linear system in the case of only a few
%       NaNs in a large array.
%       Extrapolation behavior is linear.
%         
%     method == 1 --> simple approach, applies del^2
%       over the entire array, then drops those parts
%       of the array which do not have any contact with
%       NaNs. Uses a least squares approach, but it
%       does not modify known values.
%       In the case of small arrays, this method is
%       quite fast as it does very little extra work.
%       Extrapolation behavior is linear.
%         
%     method == 4 --> Uses a spring metaphor. Assumes
%       springs (with a nominal length of zero)
%       connect each node with every neighbor
%       (horizontally, vertically and diagonally)
%       Since each node tries to be like its neighbors,
%       extrapolation is as a constant function where
%       this is consistent with the neighboring nodes.
%
%     DEFAULT: 0
%
%  bcclass - (OPTIONAL) character flag, indicating how
%       the array boundaries will be treated in the
%       inpainting operation. bcclass may be either
%       'sphere' or 'toroid', or any simple contraction
%       of these words.
%
%     bcclass = 'sphere' --> The first and last rows
%       of the array will be treated as if they are
%       at the North and South poles of a sphere.
%       Adjacent to those rows will be singular
%       phantom nodes at each pole.
%
%     bcclass = 'toroid' --> The first and last rows
%       of the array will be treated as if they are
%       adjacent to ech other. As well, the first and
%       last columns will be adjacent to each other.
%
%     DEFAULT: 'sphere'
%
% arguments (output):
%   B - nxm array with NaNs replaced
%
%
% Example:
%  [x,y] = meshgrid(0:.01:1);
%  z0 = exp(x+y);
%  znan = z0;
%  znan(20:50,40:70) = NaN;
%  znan(30:90,5:10) = NaN;
%  znan(70:75,40:90) = NaN;
%
%  z = inpaint_nans(znan);
%
%
% See also: griddata, interp1
%
% Author: John D'Errico
% e-mail address: woodchips@rochester.rr.com
% Release: 2
% Release date: 4/15/06


% I always need to know which elements are NaN,
% and what size the array is for any method
[n,m]=size(A);
A=A(:);
nm=n*m;
k=isnan(A(:));

% list those nodes which are known, and which will
% be interpolated
nan_list=find(k);
known_list=find(~k);

% how many nans overall
nan_count=length(nan_list);

% convert NaN indices to (r,c) form
% nan_list==find(k) are the unrolled (linear) indices
% (row,column) form
[nr,nc]=ind2sub([n,m],nan_list);

% both forms of index in one array:
% column 1 == unrolled index
% column 2 == row index
% column 3 == column index
nan_list=[nan_list,nr,nc];

% supply default method
if (nargin<2) || isempty(method)
  method = 0;
elseif ~ismember(method,[0 1 4])
  error('INPAINT_NANS_BC:improperargument', ...
    'If supplied, method must be one of: {0,1,4}.')
end

% supply default value for bcclass
if (nargin < 3) || isempty(bcclass)
  bcclass = 'sphere';
elseif ~ischar(bcclass)
  error('INPAINT_NANS_BC:improperargument', ...
    'If supplied, bcclass must be ''sphere'' or ''toroid''')
else
  % it was a character string
  valid = {'sphere' 'toroid'};
  
  % check to see if it is valid
  [bcclass,errorclass] = validstring(arg,valid);
  
  if ~isempty(errorclass)
    error('INPAINT_NANS_BC:improperargument', ...
      'If supplied, bcclass must be ''sphere'' or ''toroid''')
  end
end

% choice of methods
switch method
 case 0
  % The same as method == 1, except only work on those
  % elements which are NaN, or at least touch a NaN.
  
  % horizontal and vertical neighbors only
  talks_to = [-1 0;0 -1;1 0;0 1];
  neighbors_list=identify_neighbors(n,m,nan_list,talks_to);
  
  % list of all nodes we have identified
  all_list=[nan_list;neighbors_list];
  
  % generate sparse array with second partials on row
  % variable for each element in either list, but only
  % for those nodes which have a row index > 1 or < n
  L = find((all_list(:,2) > 1) & (all_list(:,2) < n)); 
  nl=length(L);
  if nl>0
    fda=sparse(repmat(all_list(L,1),1,3), ...
      repmat(all_list(L,1),1,3)+repmat([-1 0 1],nl,1), ...
      repmat([1 -2 1],nl,1),nm,nm);
  else
    fda=spalloc(n*m,n*m,size(all_list,1)*5);
  end
  
  % 2nd partials on column index
  L = find((all_list(:,3) > 1) & (all_list(:,3) < m)); 
  nl=length(L);
  if nl>0
    fda=fda+sparse(repmat(all_list(L,1),1,3), ...
      repmat(all_list(L,1),1,3)+repmat([-n 0 n],nl,1), ...
      repmat([1 -2 1],nl,1),nm,nm);
  end
  
  % eliminate knowns
  rhs=-fda(:,known_list)*A(known_list);
  k=find(any(fda(:,nan_list(:,1)),2));
  
  % and solve...
  B=A;
  B(nan_list(:,1))=fda(k,nan_list(:,1))\rhs(k);
  
 case 1
  % least squares approach with del^2. Build system
  % for every array element as an unknown, and then
  % eliminate those which are knowns.

  % Build sparse matrix approximating del^2 for
  % every element in A.
  % Compute finite difference for second partials
  % on row variable first
  [i,j]=ndgrid(1:n,1:m);
  ind=i(:)+(j(:)-1)*n;
  np=n*m;
  switch bcclass
    case 'sphere'
      % we need to have two phantom nodes at the poles
      np = np + 2;
      
      
      
      
      
  end
  
  
  
  fda=sparse(repmat(ind,1,3),[ind-1,ind,ind+1], ...
      repmat([1 -2 1],np,1),n*m,n*m);
  
  % now second partials on column variable
  [i,j]=ndgrid(1:n,2:(m-1));
  ind=i(:)+(j(:)-1)*n;
  np=n*(m-2);
  fda=fda+sparse(repmat(ind,1,3),[ind-n,ind,ind+n], ...
      repmat([1 -2 1],np,1),nm,nm);
  
  % eliminate knowns
  rhs=-fda(:,known_list)*A(known_list);
  k=find(any(fda(:,nan_list),2));
  
  % and solve...
  B=A;
  B(nan_list(:,1))=fda(k,nan_list(:,1))\rhs(k);
  
 case 4
  % Spring analogy
  % interpolating operator.
  
  % list of all springs between a node and a horizontal
  % or vertical neighbor
  hv_list=[-1 -1 0;1 1 0;-n 0 -1;n 0 1];
  hv_springs=[];
  for i=1:4
    hvs=nan_list+repmat(hv_list(i,:),nan_count,1);
    k=(hvs(:,2)>=1) & (hvs(:,2)<=n) & (hvs(:,3)>=1) & (hvs(:,3)<=m);
    hv_springs=[hv_springs;[nan_list(k,1),hvs(k,1)]];
  end

  % delete replicate springs
  hv_springs=unique(sort(hv_springs,2),'rows');
  
  % build sparse matrix of connections, springs
  % connecting diagonal neighbors are weaker than
  % the horizontal and vertical springs
  nhv=size(hv_springs,1);
  springs=sparse(repmat((1:nhv)',1,2),hv_springs, ...
     repmat([1 -1],nhv,1),nhv,nm);
  
  % eliminate knowns
  rhs=-springs(:,known_list)*A(known_list);
  
  % and solve...
  B=A;
  B(nan_list(:,1))=springs(:,nan_list(:,1))\rhs;
  
end

% all done, make sure that B is the same shape as
% A was when we came in.
B=reshape(B,n,m);

end % mainline


% ====================================================
%      end of main function
% ====================================================
% ====================================================
%      begin subfunctions
% ====================================================
function neighbors_list=identify_neighbors(n,m,nan_list,talks_to)
% identify_neighbors: identifies all the neighbors of
%   those nodes in nan_list, not including the nans
%   themselves
%
% arguments (input):
%  n,m - scalar - [n,m]=size(A), where A is the
%      array to be interpolated
%  nan_list - array - list of every nan element in A
%      nan_list(i,1) == linear index of i'th nan element
%      nan_list(i,2) == row index of i'th nan element
%      nan_list(i,3) == column index of i'th nan element
%  talks_to - px2 array - defines which nodes communicate
%      with each other, i.e., which nodes are neighbors.
%
%      talks_to(i,1) - defines the offset in the row
%                      dimension of a neighbor
%      talks_to(i,2) - defines the offset in the column
%                      dimension of a neighbor
%      
%      For example, talks_to = [-1 0;0 -1;1 0;0 1]
%      means that each node talks only to its immediate
%      neighbors horizontally and vertically.
% 
% arguments(output):
%  neighbors_list - array - list of all neighbors of
%      all the nodes in nan_list

if ~isempty(nan_list)
  % use the definition of a neighbor in talks_to
  nan_count=size(nan_list,1);
  talk_count=size(talks_to,1);
  
  nn=zeros(nan_count*talk_count,2);
  j=[1,nan_count];
  for i=1:talk_count
    nn(j(1):j(2),:)=nan_list(:,2:3) + ...
        repmat(talks_to(i,:),nan_count,1);
    j=j+nan_count;
  end
  
  % form the same format 3 column array as nan_list
  neighbors_list=[sub2ind([n,m],nn(:,1),nn(:,2)),nn];
  
  % delete replicates in the neighbors list
  neighbors_list=unique(neighbors_list,'rows');
  
  % and delete those which are also in the list of NaNs.
  neighbors_list=setdiff(neighbors_list,nan_list,'rows');
  
else
  neighbors_list=[];
end

end % function identify_neighbors

function [str,errorclass] = validstring(arg,valid)
% validstring: compares a string against a set of valid options
% usage: [str,errorclass] = validstring(arg,valid)
%
% If a direct hit, or any unambiguous shortening is found, that
% string is returned. Capitalization is ignored.
%
% arguments: (input)
%  arg - character string, to be tested against a list
%        of valid choices. Capitalization is ignored.
%
%  valid - cellstring array of alternative choices
%
% Arguments: (output)
%  str - string - resulting choice resolved from the
%        list of valid arguments. If no unambiguous
%        choice can be resolved, then str will be empty.
%
%  errorclass - string - A string argument that explains
%        the error. It will be one of the following
%        possibilities:
%
%        ''  --> No error. An unambiguous match for arg
%                was found among the choices.
%
%        'No match found' --> No match was found among 
%                the choices provided in valid.
%
%        'Ambiguous argument' --> At least two ambiguous
%                matches were found among those provided
%                in valid.
%        
%
% Example:
%  valid = {'off' 'on' 'The sky is falling'}
%  
%
% See also: parse_pv_pairs, strmatch, strcmpi
%
% Author: John D'Errico
% e-mail: woodchips@rochester.rr.com
% Release: 1.0
% Release date: 3/25/2010

ind = strmatch(lower(arg),lower(valid));
if isempty(ind)
  % No hit found
  errorclass = 'No match found';
  str = '';
elseif (length(ind) > 1)
  % Ambiguous arg, hitting more than one of the valid options
  errorclass = 'Ambiguous argument';
  str = '';
  return
else
  errorclass = '';
  str = valid{ind};
end

end % function validstring