function [x, g, accept, L, X, xnew] = mhmcmc(fun, x, varargin)
% MHMCMC Hamiltonian Markov chain Monte Carlo
%
% Usage
%   [x, g, accept, L, X] = mhmcmc(fun, x, options)
%
% Input
%   fun          Function handle to objective function of the following form
%                   f = fun(x, '0');       % Returns objective 
%                   g = fun(x, '1');       % Returns gradient
%                   [f, g] = fun(x, '01'); % Returns objective and gradient
%   x            Initial value (N x 1)-vector
%   options
%      epsilon   Step size (default 10^(-5*rand-2))
%      tau       Number of leapfrog steps (default 20)
%      A,b       Linear inequality constraints, A*x>b
%      Aeq,beq   Linear equality constraints, Aeq*x=beq
%      verbose   Display information
%
% Output
%   x            Final value of x
%   g            Final gradient at x
%   accept       New step accepted (true/false)
%   L            Negative log likelihood of final x
%   X            Cell array with intermediate results
%   xnew         Final value of x (which might have been rejected)

% Constants etc.
[N,J] = size(x);

% Parse inputs
opts = mgetopt(varargin);
epsilon = mgetopt(opts, 'epsilon', 10.^(-5*rand-2));
tau = mgetopt(opts, 'tau', 20);
A = mgetopt(opts, 'A', zeros(0,N));
b = mgetopt(opts, 'b', zeros(0,1));
Aeq = mgetopt(opts, 'Aeq', []);
beq = mgetopt(opts, 'beq', []);
verbose = mgetopt(opts, 'verbose', false);

if nargout>=5, 
    X = cell(tau+1,J); 
    for j=1:J, X{1,j} = x(:,j); end
end

% Compute orthogonal basis of equality constraints + new origin
if isempty(Aeq) % No equality constraints: use standard basis
    r = 0;
    P = zeros(0,N);
    K = eye(N);
    cx = zeros(N,1);
else
    [U,S] = svd(Aeq');
    [m,n] = size(Aeq);
    if m>1, s=diag(S); elseif m==1, s=S(1); else s=0; end
    tol = max(m,n)*max(s)*eps;
    r = sum(s>tol);
    P = U(:,1:r)';
    K = U(:,r+1:end)';
    cx = Aeq\beq;
end


[L,g] = fun(x, '01');

% p = randcg(zeros(N,J), eye(N), [], [], Aeq, zeros(size(beq)), zeros(N,1), 10);
% p = (K'*K)*randn(N,J); % Is this right? no
% p = K'*randn(N-r,J); % Or is this? no

% p = randcg(zeros(N,J), eye(N), [], [], Aeq, zeros(size(beq)), K'*randn(N-r,J), 10);
% p = repmat(randcg(zeros(N,1), eye(N), [], [], Aeq, zeros(size(beq)), K'*randn(N-r,1), 10),1,J);

p = K'*randn(N-r,J);


H = sum(sum(p.^2,1))/2 + L;

xnew = x;
gnew = g;
if verbose, mprogress; end;
for t = 1:tau
    % Half-step in p
    p = p - epsilon*(K'*K)*gnew/2;     
     
    % Full step in x
    if isempty(A) % No inequality constraints
        xnew = xnew + epsilon*p;
        if nargout>=5, for j=1:J, X{t+1,j} = xnew(:,j); end; end
    else % Inequality constraints
        maxEpsK = bsxfun(@minus, b, A*xnew)./(A*p);
        maxEpsK(maxEpsK<=0 & A*p>=0) = nan;        
        stepEps = min(maxEpsK, [], 1);
        if isnan(stepEps), stepEps=inf; end;
        iFullStep = stepEps>epsilon;
        iReflect = stepEps<=epsilon;
        xnew(:,iFullStep) = xnew(:,iFullStep) + epsilon*p(:,iFullStep);
        if nargout>=5, for j=find(iFullStep), X{t+1,j} = xnew(:,j); end; end
        for j = find(iReflect)
            restEps = epsilon;
            iRefl = [];
            while restEps>0
                maxEpsK =(b-A*xnew(:,j))./(A*p(:,j));
                maxEpsK(maxEpsK<=0 & A*p(:,j)>=0) = nan; 
                maxEpsK(iRefl) = nan; % Don't consider variables that have just been reflected
                stepEps = min(maxEpsK, [], 1);
                if isnan(stepEps), stepEps=inf; end;
                stepEps = min(stepEps, restEps);
%                 stepEps = min(min(maxEpsK(maxEpsK>0 & A*p(:,j)<0), [], 1), restEps);
                iRefl = find(maxEpsK==stepEps);
                xnew(:,j) = xnew(:,j) + stepEps*p(:,j);
                if nargout>=5, X{t+1,j} = [X{t+1,j} xnew(:,j)]; end;
                restEps = restEps - stepEps;            
                if ~isempty(iRefl)
                    xnewp = K*(xnew(:,j)-cx);
                    pp = K*p(:,j);
                    Ap = A*K';
                    bp = b-A*cx;
                    for iR = 1:length(iRefl)
                        aR = Ap(iRefl(iR),:);
                        bR = bp(iRefl(iR));
                        pp = pp-2*(aR*(xnewp+pp)-bR)./(aR*aR')*aR';
                    end
                    p(:,j) = K'*pp;
                end
            end
        end
    end                    
    gnew = fun(xnew, '1');
    
    % Half-step in p
    p = p - epsilon*(K'*K)*gnew/2;     
    
    if verbose, mprogress(t/tau); end;
end

Lnew = fun(xnew, '0');
Hnew = sum(sum(p.^2,1))/2 + Lnew;
dH = Hnew - H;
if dH<0 || rand<exp(-dH)
    accept = true;
    x = xnew; 
    g = gnew;
    L = Lnew;
else
    accept = false;
end


