function regstr = mlreg(x,t,tout,t0,s0,K,h,thest,maxit,rndst)
%function regstr = mlreg(x,t,tout,theta0,tau,K,h,thest,maxit,rndst)
%
%Maximum Likelihood Registration (as of 8/22/05)
%
%INPUT
%   x:      (n x m)     Data matrix (subjects in rows)
%   t:      (1 x m)     Input time grid
%   tout:   (1 x mout)  Output time grid
%   theta0: (1 x p)     Location parameter of truncated Normal for Theta
%   tau:    (1 x p)     Scale parameter of truncated Normal for Theta
%   K:      (real)      Monte Carlo size for Theta
%                       (suggested: K = 200)
%   h:      (real)      Bandwidth for density estimator of g(t,Theta)
%                       (if h<=0, the default 75th percentile of 
%                       oversmoothing bandwidths is used)
%   thest:  (char)      Method of estimation for individual Thetas:
%                           'mode': mode of conditional distribution
%                           'mean': mean of conditional distribution
%   maxit:  (real)      Maximum number of iterations
%   rndst:  (real)      Initial state for random number generator
%                       (if negative, current system state is used)
%
%OUTPUT
% regstr: struct with the following fields:
%   t:      (1 x m)     Same as input
%   tout:   (1 x mout)  Same as input 
%   theta0: (1 x p)     Same as input
%   tau:    (1 x p)     Same as input
%   K:      (real)      Same as input
%   h:      (real)      Same as input
%   mu:     (1 x mout)  Estimated structural mean
%   s:      (real)      Estimated error standard deviation
%   L:      (real)      Final value of log-likelihood function
%   theta:  (n x p)     Estimated Thetas
%   g:      (n x p)     Functions g(t,Theta)
%   g1:     (n x p)     Warping functions, inverses of g(t,Theta)
%   xreg:   (n x m)     Raw registered curves
%   xregsm: (n x mout)  Smoothed registered curves
%
%External functions used:
% NORMPDF (Statistics Toolbox) and FMINSEARCH (Optimization Toolbox).
% These may not be available in your particular Matlab installation.


% Checking input arguments
if nargin<10
    disp('Not enough input arguments')
    regstr = [];
    return
end
[n,m] = size(x);
if size(t,2)~=m | size(t,1)~=1
    disp('Dimensions of X and T do not match')
    regstr = [];
    return
end
if thest~='mean' | thest~='mode'
    disp('Wrong input THEST')
    regstr = [];
    return
end
if size(tout,1)>1,
    tout = tout';
end
mout = length(tout);

% Monte Carlo thetas and g(t,theta)
if rndst>=0
    randn('state',rndst)
end
p = length(t0);
a = min(t);
b = max(t);
thsamp = Inf*ones(K,p);
gthsamp = zeros(K,m);
for i = 1:K
    while thsamp(i,1)<a|thsamp(i,p)>b|any(thsamp(i,1:p-1)>=thsamp(i,2:p))
        thsamp(i,:) = t0 + s0.*randn(1,p);
    end
    gthsamp(i,:) = interp1([a thsamp(i,:) b], [a t0 b], t, 'cubic');
end
ker = zeros(m,K,mout);
if h<=0
    h = (40*sqrt(pi)/K)^(1/5) * mean(std(gthsamp));
end
for k = 1:mout
    ker(:,:,k) = (.75*(1-((gthsamp-tout(k))/h).^2).* ...
        (abs(gthsamp-tout(k))<=h)/h)';
end

% Initialization
mu = interp1(t,mean(x),tout,'cubic');
s = sqrt(mean(var(x)));
mug = zeros(K,m);
for j = 1:m
    mug(:,j) = interp1(tout,mu,gthsamp(:,j),'cubic');
end
fx = zeros(n,1);        % Marginal density f(x)
fxth = zeros(K,1);      % Conditional density f(x|theta)
ss = 0;
for i = 1:n
    fxth = prod(normpdf(repmat(x(i,:),K,1),mug,s),2);
    fx(i) = mean(fxth);
    ss = ss + mean(sum((repmat(x(i,:),K,1)-mug).^2,2).*fxth/fx(i));
end
s = sqrt(ss/(n*m));
L = mean(log(fx));
disp(['Initial Loglik: ' num2str(L) ', S: ' num2str(s)])

% Iterations
err = 1;
iter = 0;
Lopt = L;
w = zeros(n,m,mout);    % ML weights
while err>1e-3 & iter<maxit
    iter = iter+1;
    L0 = L;
    % Weight computation
    for i = 1:n
        fxth = prod(normpdf(repmat(x(i,:),K,1),mug,s),2);
        for k = 1:mout
            w(i,:,k) = (ker(:,:,k)*fxth/K)'/fx(i);
        end
    end
    % Mean update
    for k = 1:mout
        mu(k) = sum(sum(x.*w(:,:,k)))/sum(sum(w(:,:,k)));
    end
    % Log-likelihood and Sigma update
    for j = 1:m
        mug(:,j) = interp1(tout,mu,gthsamp(:,j),'cubic');
    end
    ss = 0;
    for i = 1:n
    	fxth = prod(normpdf(repmat(x(i,:),K,1),mug,s),2);
	    fx(i) = mean(fxth);
        ss = ss + mean(sum((repmat(x(i,:),K,1)-mug).^2,2).*fxth/fx(i));
    end
    s = sqrt(ss/(n*m));
    L = mean(log(fx));
    if L>Lopt
        Lopt = L;   % Update optimum L
    end
    % Convergence check-up
    err = abs(L-L0)/abs(L0);
    disp(['Iteration ' num2str(iter) ', Error ' num2str(err) ...
        ', Loglik ' num2str(L) ', S ' num2str(s) ...
        ', Current opt ' num2str(Lopt)])
end

% Theta estimation and curve registration
disp('Estimating Thetas ...')
theta = zeros(n,p);
g = zeros(n,m);
g1 = zeros(n,m);
xreg = zeros(n,m);
for i = 1:n
    if thest=='mode'
        theta(i,:) = fminsearch(@F,t0,[],t0,s0,t,tout,mu,s,x(i,:));
    else
        fxth = prod(normpdf(repmat(x(i,:),K,1),mug,s),2);
        fx(i) = mean(fxth);
        theta(i,:) = mean(thsamp.*repmat(fxth/fx(i),1,p));
    end
    g(i,:) = interp1([a theta(i,:) b], [a t0 b], t, 'cubic');
    g1(i,:) = interp1([a t0 b], [a theta(i,:) b], t, 'cubic');
    xreg(i,:) = interp1(t,x(i,:),g1(i,:),'cubic');
end
xregsm = zeros(n,mout);
for k = 1:mout
    xregsm(:,k) = sum(x.*w(:,:,k),2)./sum(w(:,:,k),2);
end

% Output
regstr = struct('t',t,'tout',tout,'theta0',t0,'tau',s0,'K',K,'h',h, ...
    'mu',mu,'s',s,'L',L,'theta',theta,'g',g,'g1',g1,'xreg',xreg, ...
    'xregsm',xregsm);


% --------------------------------------------------

% Auxiliary function
function y = F(theta,theta0,tau,t,tout,mu,s,xi)
m = length(t);
a = t(1);
b = t(m);
if any([a theta]>=[theta b])
    y = Inf;
    return
end
g = interp1([a theta b], [a theta0 b], t, 'cubic');
mug = interp1(tout,mu,g,'cubic');
y = sum((xi-mug).^2)/(2*s^2) - log(prod(normpdf(theta,theta0,tau)));