function [y_fit, coeffs, rms_error, num_under] = least_squares(x, y, ...
    x_fit, num_funcs, basis_funcs)
% Find the least squares best fit to observations for given basis functions.
%
% function [y_fit, coeffs, rms_error, num_under] = least_squares(x, y, ...
%          x_fit, num_funcs, basis_funcs)
%
%    INPUT:
%
% x: column vector of positions at which observations are made
% y: matrix of observations dimensioned (length(x), N) so that N sets of
%     observations can be fitted in the same call. Note that y may contain
%     NaNs, but this will slow things down.
% x_fit: if this is a vector then it is used to determine the x values for
%        which y_fit is calculated; if it is a scalar then x is used.
% num_funcs: the number of basis functions.
% basis_funcs: a string specifying the basis functions to be used. The
%              default if no string is passed is 'fourier', the only type
%              written at present.
%
%    OUTPUT:
%
% y_fit: a set of values on the best fit curve. It is returned as a matrix
%        dimensioned (length(x_fit), N) or (length(x), N) if x_fit is a
%        scalar. If a given column in y describes an underdetermined problem
%        then NaNs are returned.
% coeffs: the best fit coefficients of the basis functions, returned as a
%         matrix dimensioned (num_funcs, N). If a given column in y
%         describes an underdetermined problem then NaNs are returned.
% rms_error: a row vector giving the rms value (actually standard deviation
%            as calculated by std) for the goodness of fit to the observations.
% num_under: The number of columns in y where there are less than num_funcs
%            non-NaN values, i.e., where the problem is under-determined.

% $Id: least_squares.m,v 1.6 1997/09/22 03:59:34 mansbrid Exp $
% Copyright J. V. Mansbridge, CSIRO, Mon Sep 15 12:45:01 EST 1997

if nargin < 4
  help least_squares
  error('too few arguments')
end

% Check that the input arrays have the correct dimensions.

[len_x, n_test] = size(x);
if n_test ~= 1
  error('x must be a column vector')
end
[m_test, n_y] = size(y);
if m_test ~= len_x
  error('y must have the same number of rows as x')
end

% Check that the method is only applied for over-determined problems.

if len_x < num_funcs
  disp('System is under-determined')
  coeffs = repmat(NaN, num_funcs, n_y);
  if length(x_fit) == 1
    y_fit = repmat(NaN, len_x, n_y);
  else
    y_fit = repmat(NaN, length(x_fit), n_y);
  end
  num_under = n_y;
  return
else
  num_under = 0;
end

% Set the basis functions if necessary.

if nargin == 4
  basis_funcs = 'fourier';
else
  basis_funcs = lower(basis_funcs);
end

% Find the NaNs so that y can be divided up into those columns that involve
% NaNs and those that don't.

nan_pos = isnan(y);
nan_cols = sum(nan_pos, 1);
ff_nan = find(nan_cols ~= 0);
len_ff_nan = length(ff_nan);

if len_ff_nan == 0 % The original y does not contain any NaNs.
  [y_fit, coeffs, rms_error] = l_s(x, y, x_fit, num_funcs, basis_funcs);
else
  % initialisation
  
  ff_non_nan = find(nan_cols == 0);
  len_ff_non_nan = length(ff_non_nan);
  coeffs = repmat(NaN, num_funcs, n_y);
  rms_error = repmat(NaN, 1, n_y);
  if length(x_fit) == 1
    y_fit = repmat(NaN, len_x, n_y);
  else
    y_fit = repmat(NaN, length(x_fit), n_y);
  end
  
  % Do the inversion stuff for all of the columns that don't contain NaNs.
  
  if len_ff_non_nan > 0
    y_non_nan = y(:, ff_non_nan);
    [y_fit_non_nan, coeffs_non_nan, rms_error_non_nan] = ...
	l_s(x, y_non_nan, x_fit, num_funcs, basis_funcs);
    coeffs(:, ff_non_nan) = coeffs_non_nan;
    rms_error(ff_non_nan) = rms_error_non_nan;
    y_fit(:, ff_non_nan) = y_fit_non_nan;
  end
  
  % Do the inversion stuff, one column at a time, for the columns that
  % contain NaNs.
  
  for ii = 1:len_ff_nan
    col = ff_nan(ii);
    ff = find(nan_pos(:, col) == 0);
    if length(ff) > num_funcs % Only try a fit for over-determined problems
      x_nan = x(ff);
      y_nan = y(ff, col);
      [y_fit_nan, coeffs_nan, rms_error_nan] = l_s(x_nan, y_nan, x_fit, ...
	  num_funcs, basis_funcs);
      coeffs(:, col) = coeffs_nan;
      rms_error(col) = rms_error_nan;
      y_fit(:, col) = y_fit_nan;
    else
      num_under = num_under + 1;
    end
  end
end

%%%%%%%%%%  function l_s  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [y_fit, coeffs, rms_error] = l_s(x, y, x_fit, num_funcs, basis_funcs)

% Inversion function which does not need to do checks and assumes that y
% contains no NaNs. 

func = find_basis_func(x, num_funcs, basis_funcs);

% Set the parts of the inversion problem.

func_t = func';
b = func_t*y;
A = func_t*func;
coeffs = (A\b);
y_fit = func*coeffs;
error = y_fit - y;
% rms_error = sqrt(mean(error.*error, 1))
rms_error = std(error);
if length(x_fit) ~= 1 % recalculate func to use x_fit to calculate y_fit
  func = find_basis_func(x_fit, num_funcs, basis_funcs);
  y_fit = func*coeffs;
end

%%%%%%%%%%  function find_basis_func  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function func = find_basis_func(x, num_funcs, basis_funcs)

% Define the array of values of the basis functions

len_x = length(x);
switch basis_funcs
  case 'fourier'
    n_max = (num_funcs - 1)/2;
    if (n_max ~= round(n_max)) | (num_funcs < 1)
      error(['num_funcs should be a positive odd integer but is ' ...
	    num2str(num_funcs)])
    end
    two_pi = 2*pi;
    func = ones(len_x, 2*n_max + 1);
    for ii = 1:n_max
      func(:, 2*ii) = cos(ii*two_pi*x);
      func(:, 2*ii+1) = sin(ii*two_pi*x);
    end
  otherwise
    error(['unrecognised basis functions: ' basis_funcs])
end
