function [par,beta]=covest_ml(D,z,par0,AB,method,par_start)
% covest_ml estimates a Matérn covariance with the Maximum Likelihood method.
%
% par=covest_ml(D,z,par_fixed)
%  or
% [par,beta]=covest_ml(D,y,par_fixed,AB)
%  or
% [par,beta]=covest_ml(D,y,[],AB)
%  or
% [par,beta]=covest_ml(D,y,par_fixed,AB,method,par_start)
%
% par = [sigma2; kappa; nu; sigma2_epsilon]
% D is the distance matrix for the observations
% z are the residuals  z=y-mu,  for known expectation  mu
% y are the observations, if  mu  is unknown, but  AB  supplied
%
% Optional parameters:
% par_fixed has the same structure as  par.  (default=zeros(4,1))
%           Allows setting some parameters to known values.
%           A non-zero entry for a parameter fixes that parameter
%           to that value.
% AB is the product of the observation matrix A and the expectation
%    basis matrix B.
%    If AB is supplied, the optimisation also estimates
%    beta, for mu=B*beta.
% If AB is supplied,  beta  are the estimated expectation parameters.
% method if method==1 estimates the the quota sigma2/sigma2_epsilon first
%        snd separate into sigma2 and sigma2_epsilon afterwards. This can
%        NOT be used if either sigma or sigma2_epsilon are fixed.
% par_start a vector of initial values for the optimisation.
%
% Setting a parameter to [] gives reasonable defaults.

% $Id: covest_ml.m 4410 2011-09-08 14:08:01Z johanl $

if (nargin<3), par0 = []; end
if (nargin<4), AB = []; end
if (nargin<5), method = []; end
if (nargin<6), par_start = []; end
if isempty(par0), par0 = zeros(4,1); end

fixed = (par0>0);
if isempty(method)
  if (fixed(1) || fixed(4))
    method = 0;
  else
    method = 1;
  end
end

%find unique distances
[dunique,I,J] = unique(D(D>0));

if (method==0)
  fcn = @(par) ...
        covest_matern_ml_loss0(D,z,...
                               exp(par).*(~fixed)+...
                               par0.*fixed,AB,I,J);
  par = fminsearch(@(par) fcn(par)+(fixed'*(exp(par)-par0).^2),...
                   log([var(z)/2;par0(2:3)+1;var(z)/2])+0.001);
  par = exp(par).*(~fixed)+par0.*fixed;

  if (~isempty(AB)) % AB supplied, estimate beta
    sigma2 = par(1);
    kappa = par(2);
    nu = par(3);
    sigma2_eps = par(4);
    
    Sigma_yy = matern_covariance(D,sigma2,kappa,nu,I,J)+...
        sigma2_eps*eye(length(z));
    R = chol(Sigma_yy);
    
    tmp = ([z,AB]'/R)';
    beta = tmp(:,2)\tmp(:,1);
  end
else
  if (fixed(1) || fixed(4))
    error(['This method can only be used when both variance ',...
           'parameters are to be estimated.']);
  end
  if isempty(par_start)
    par_start = [0;log(par0(2:3)+1)+0.001];
  else
    par_start = [log(par_start(1)+0.001)-log(par_start(4)+0.001);...
                 log(par_start(2:3))];
  end
  fcn = @(par) ...
        covest_matern_ml_loss1(D,z,...
                               [1/(1+exp(-par(1)));...
                      exp(par(2:3)).*(~fixed(2:3))+...
                      par0(2:3).*fixed(2:3)],AB,I,J);
  par = fminsearch(@(par) fcn(par)+(fixed(2:3)'*...
                                    (exp(par(2:3))-par0(2:3)).^2),...
                   par_start);
  par(1) = 1/(1+exp(-par(1)));
  par(2:3) = exp(par(2:3)).*(~fixed(2:3))+par0(2:3).*fixed(2:3);
  
  if (isempty(AB))
    sigma2_ratio = par(1);
    kappa = par(2);
    nu = par(3);
    Sigma_yy = sigma2_ratio*matern_covariance(D,1,kappa,nu,I,J)+...
        (1-sigma2_ratio)*eye(length(z));
    R = chol(Sigma_yy);
    
    sigma2_sum = sum((z'/R).^2)/length(z);
    
    par(1) = sigma2_ratio*sigma2_sum;
    par(4) = (1-sigma2_ratio)*sigma2_sum;
  else % AB supplied, estimate beta
    sigma2_ratio = par(1);
    kappa = par(2);
    nu = par(3);
    
    Sigma_yy = sigma2_ratio*matern_covariance(D,1,kappa,nu,I,J)+...
        (1-sigma2_ratio)*eye(length(z));
    R = chol(Sigma_yy);
    
    tmp = ([z,AB]'/R)';
    beta = tmp(:,2:end)\tmp(:,1);

    sigma2_sum = sum((tmp(:,1)-tmp(:,2:end)*beta).^2)/length(z);
    
    par(1) = sigma2_ratio*sigma2_sum;
    par(4) = (1-sigma2_ratio)*sigma2_sum;
  end
end


function S=covest_matern_ml_loss0(D,z,par,AB,I,J)

sigma2 = par(1);
kappa = par(2);
nu = par(3);
sigma2_eps = par(4);

Sigma_yy = matern_covariance(D,sigma2,kappa,nu,I,J)+...
    sigma2_eps*eye(length(z));
R = chol(Sigma_yy);

if isempty(AB)
  S = sum(log(diag(R)))+0.5*sum((z'/R).^2);
else
  tmp = ([z,AB]'/R)';
  beta = tmp(:,2:end)\tmp(:,1);
  S = sum(log(diag(R)))+0.5*sum((tmp(:,1)-tmp(:,2:end)*beta).^2);
end



function S=covest_matern_ml_loss1(D,z,par,AB,I,J)

sigma2_ratio = par(1);
kappa = par(2);
nu = par(3);

Sigma_yy = sigma2_ratio*matern_covariance(D,1,kappa,nu,I,J)+...
    (1-sigma2_ratio)*eye(length(z));
[R,p] = chol(Sigma_yy);

if p~=0 || any(isnan(R(:)))
    S = realmax;
    return
end
if isempty(AB) 
  sigma2_sum = sum((z'/R).^2)/length(z);
  S = sum(log(diag(R)))+length(z)/2*log(sigma2_sum);
else
  tmp = ([z,AB]'/R)';
  beta = tmp(:,2:end)\tmp(:,1);
  sigma2_sum = sum((tmp(:,1)-tmp(:,2:end)*beta).^2)/length(z);
  S = sum(log(diag(R)))+length(z)/2*log(sigma2_sum);
end