function [positionMatrix probabilityVector res] = sampleMetropolis(par, nSample)
% sample from arbitrary distribution
%
% Syntax: [pointMatrix, probabilityVector] = sampleMetropolis(fHandle, ...
%            nSample, initX, initCov)
%
% THIS IS NO USER FUNCTION

% The elk-library: convex geometry applied to crystallization modeling.
%   Copyright (C) 2013 Alexander Reinhold
%
% This program is free software: you can redistribute it and/or modify it 
%   under the terms of the GNU General Public License as published by the 
%   Free Software Foundation, either version 3 of the License, or (at your 
%   option) any later version.
%
% This program is distributed in the hope that it will be useful, but 
%   WITHOUT ANY WARRANTY; without even the implied warranty of 
%   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU 
%   General Public License for more details.
%
% You should have received a copy of the GNU General Public License along 
%   with this program.  If not, see <http://www.gnu.org/licenses/>.

%% Algorithm Overview (Metropolis)
%
% Source: H. Haario, "An Adaptive Metropolis Algorithm"; Bernoulli, Vol. 7,
%         2001.
%
% The Metropolis algorith can create samples proportional to any
%   distribution:
fHandle = par.metroHandle;
%   without any a priory knowledge. The distribution density should be
%   reasonably connected (e.g. two well seperated normal distributions
%   might not be sampled properly). The algorithm uses a random walk or
%   Markov chain to generate these sample points. Therefore, the user must 
%   provide a reasonable initial state for the Markov chain:
metroX = par.metroInit;
%   The corresponding distribution density sould be large for that point.
%   Additionally, a guess for the covariance matrix of sampled points is
%   required:
metroCov = par.metroCov;
%   which must reflect the range of the distribution but is not required to
%   be accurate. We can direclty calculate the initial value of of the
%   point in the Markov chain that is used to sample upcoming points:
metroVal = fHandle(metroX);
%
%% Phase 1 (Metropolis)
% In an initial step of the algorithm:
metroState = 1;
% the Markov chain is evaluated for a small amount of steps (10):
nNonAdaptive = par.metroNonAdaptive;
%   This is done to allow a reasonable estimate for the covariance matrix
%   in the next step.
%
%% Phase 2 (Metropolis)
% The inital covariance matrix is substituted by one calculated from the
%   already sampled points and the adaption rules (see source) is applied.
%   This phase may be hold for additional sample points:
nNoHistory = par.metroNoHistory;
%   which might be required when the initial guess of the covariance matrix
%   is too bad or many points are sampled so that the history should avoid
%   the slightly poorer adaptive phase of the algorithm.
%
%% Phase 3 (Metropolis)
% The adaption rules still apply and a history of:
nHistory = par.metroHistory;
if nHistory < nNonAdaptive
    nHistory = nNonAdaptive;
end
%   sample points is taken from which an autocorrelation function is
%   computed up to a lag of:
nLag = par.metroLag;
%   Based on this autocorrelation function, we estimate how many steps in
%   the Markov chain are required so that we have an independent point.
%   This time is assumed to be proportional to the time until which the
%   last coordinate has taken a negative value. This value is stored as:
%     nStepsLastZero
%
%% Phase 4 (Metropolis)
% We now take samples, every:
%     nSampleDelay = nStepsLastZero * nSampleDelayFraction
%   step of the Markov chain where the fration (typically 2) is specified
%   by the user:
nSampleDelayFraction = par.metroSpacing;
%   During this phase, all samples are taken and the acceptance rate is
%   evaluated which is the ratio of accepted points vs. steps in the Markov
%   chain.
% This phase has also been parallelized, allowing to generate:
nSampleBatch = par.metroSampleBatch;
%   samples at once. Therefore, the X-value of the Markov chain is copied
%   nSampleBatch times to yield nSampleBatch seperate Markov chains that
%   run in parallel. The last step reduces the amount of chain to obtain
%   exactly nSample points at the end.
%
%% Computing the probability density
% to provide the probability distribution for the sampled points, we need
%   to integrate the given (non-normalized) distribution function.
%   Therefore, we can use the estimated mean and covariance that was used
%   for the markov steps so that the integral is detemined by Monte Carlo
%   integration with normally sampled points.
relTol = par.metroRelTol;
nIntegralBatch = par.metroIntegralBatch;
% The user might, however specify the integral himself avoiding this
% procedure:
metroIntegral = par.metroIntegral;

%% Init algorithm
% for the following, we use a state indicating:
%   1 - no adaption, no history, no sampling
%   2 - with adaption, no history, no sampling
%   3 - with adaption, with history, no sampling
%   4 - no adaption, no history, sampling
%   0 - algorithm finished

nDim = length(par.metroInit);

% initial variables
nMetroIter = 0;
nAdaptCounter = 0;
nAccepted = 0;
metroCholCov = chol(metroCov);
historyMatrix = nan(nDim, nHistory);
% historyCount = 0;
sampleMatrix = nan(nDim, nSample);
sampleValueVector = nan(1, nSample);

% to be save any parameter was introduced properly, we delete it so it
% cannot be used
verbose = par.verbose;
clear par

%% Start algorithm
cumulativeTime = zeros(1, 5);tic;
nBatch = 1;
if verbose,disp('Started: Metropolis algorithm.');end
while metroState
    % increase iteration count
    nMetroIter = nMetroIter + nBatch;

    % sample new metro point (adapted lines from sampleNormal to avoid
    %   overhead)
    %! Dev: to avoid rare problems involving valid hC-vectors, a try catch
    %    is added.
    for i = 1:10
        try
            metroNewX = metroCholCov*randn(nDim, nBatch) + metroX;
            metroNewVal = fHandle(metroNewX);
            break
        catch
            % well, accept it
        end
    end
    
    % accept (or reject)
    acceptFilter = rand(1, nBatch) < ...
        min(ones(1, nBatch), metroNewVal./metroVal);
    nAccepted = nAccepted + sum(acceptFilter);
    metroX(:, acceptFilter) = metroNewX(:, acceptFilter);
    metroVal(acceptFilter) = metroNewVal(acceptFilter);
    
%     if rand < min([1, metroNewVal / metroVal])
%         nAccepted = nAccepted + 1;
%         metroX = metroNewX;
%         metroVal = metroNewVal;
%     end
    
    % update history (also used to evaluate mean and cov after phase 2)
    if metroState <= 3
        historyMatrix(:, mod(nMetroIter-1, nHistory)+1) = metroX;
    end
    
    % adapt covariance matrix
    if metroState == 2 || metroState == 3
        nAdaptCounter = nAdaptCounter + 1;
        p1 = (nAdaptCounter - 1) / nAdaptCounter;
        p2 = 1/nAdaptCounter;
        % update mean
        metroPreMean = metroMean;
        metroMean = p1*metroMean + p2*metroX;
        % update cov
        metroCov = p1 * metroCov + covFactor*p2*(...
            nAdaptCounter * (metroPreMean*(metroPreMean')) - ...
            (nAdaptCounter + 1)*(metroMean*(metroMean')) + ...
            metroX*metroX' + covMargin);
        % update cholesky factorization for random sampling
        try
            metroCholCov = chol(metroCov);
        catch
            error('elk:distribution:numericalError', ['Adaptive part of ' ...
                'Metropolis algorithm failed. Usually, your guess for the ' ...
                'initial point or the initial covariance is bad. Change ' ...
                'them, or if you are already close, try increasing the ' ...
                'number of non-adaptive cycles.']);
        end
    end
    
    % add sample
    if metroState == 4 && iStepElapsedForSampling == nSampleDelay
        sampleMatrix(:, iSample + (1:nBatch)) = metroX;
        sampleValueVector(iSample + (1:nBatch)) = metroVal;
        iSample = iSample + nBatch;
        if verbose,disp(['Metropolis: sample ' num2str(iSample) ' generated.']);end
        % at the end, ensure we take exactly the requested number of
        % samples
        nBatch = min(nBatch, nSample - iSample);
        metroX = metroX(:, 1:nBatch);
        metroVal = metroVal(1:nBatch);
        iStepElapsedForSampling = 1;
    elseif metroState == 4
        iStepElapsedForSampling = iStepElapsedForSampling + 1;
    end
    
    % switch through metro states // events
    if metroState == 1 && nMetroIter == nNonAdaptive
        if verbose,disp(['Metropolis: state 1 finished after ' ...
                num2str(nMetroIter) ' steps.']);end
        cumulativeTime(1) = toc;
        % initialize covariance matrix
        metroMean = mean(historyMatrix(:, 1:nMetroIter), 2);
        % the following factor is provided by the paper:
        covFactor = 2.4^2 / nDim;
        metroCov = covFactor * cov(historyMatrix(:, 1:nMetroIter)');
        
        % display result
        if verbose > 1
            disp('So far, the estimated mean is:')
            disp(metroMean')
            disp('.. the estimated covariance is:')
            disp(metroCov)
            disp('.. the determinant caclucates to:')
            disp(det(metroCov))
            disp('..the acceptance rate is:')
            disp(nAccepted/nMetroIter)
        end
        
        metroCholCov = chol(metroCov);
        % the following margin is also suggested by the paper:
        covMargin = 1e-6 * diag(diag(metroCov));
        % the covariance uses now
        nAdaptCounter = nMetroIter;
        
        % state change
        metroState = 2;
    end
    if metroState == 2 && nMetroIter == nNonAdaptive + nNoHistory
        if verbose,disp('Metropolis: state 2 finished');end
        cumulativeTime(2) = toc;
        % state change
        nMetroIter = 0;
        metroState = 3;
    end
    if metroState == 3 && nMetroIter == nHistory
        if verbose,disp('Metropolis: state 3 finished');end
        cumulativeTime(3) = toc;
        % evaluate history
        metroAcMatrix = computeAcMatrix(historyMatrix, nLag);
        % the autocorrelation must have reached negative values for every 
        %   coordinate, we will need the step in which this condition
        %   becomes true:
        lastZeroVector = nan(1, nDim);
        for iDim = 1:nDim
            thisLagTime = find(metroAcMatrix(iDim, :) < 0, 1, 'first');
            if isempty(thisLagTime)
                lastZeroVector(iDim) = nLag;           
            else
                lastZeroVector(iDim) = thisLagTime;
            end
        end
        
        % check if everything is sane
        if mean(lastZeroVector) < 0.9*nLag
            % set parameters for next step
            nStepsLastZero = mean(lastZeroVector);
            nSampleDelay = round(nStepsLastZero * nSampleDelayFraction);
            iSample = 0;
        else
            if verbose > 1
            disp('For the following error, the estimated mean is:')
            disp(metroMean')
            disp('And the estimated covariance is:')
            disp(metroCov)
            end
            error('elk:distribution:adaptiveMetropolis', ...
                ['The adaptive Metropolis algorithm did not converge ' ...
                 'properly. Look, if metroInit and metroCov are ' ...
                 'appropriately chosen.']);
        end
        
        % change to batch mode
        nBatch = min(nSampleBatch, nSample);
        metroX = metroX * ones(1, nBatch);
        metroVal = metroVal * ones(1, nBatch);
        
        % state change
        nMetroIter = 0;
        nAccepted = 0;
        iStepElapsedForSampling = 0;
        metroState = 4;
        if verbose,disp(['Metropolis: generating samples every ' ...
                num2str(nSampleDelay) ' step.']);end
        
    end
    if metroState == 4 && iSample >= nSample
        if verbose,disp('Metropolis: state 4 finished');end
        cumulativeTime(4) = toc;
        metroState = 0;
    end

end
% position matrix is straight forward
positionMatrix = sampleMatrix;

% for the probability values, we require the normalizing factor which makes
%   the input density distribution to a probability distribution.
if ~isempty(metroIntegral)
    y = metroIntegral;
    yErr = 0;
else
    [y yErr] = integrateNormal(fHandle, metroMean, metroCov/covFactor, ...
        relTol, nIntegralBatch, verbose);
end
probabilityVector = sampleValueVector / y;
cumulativeTime(5) = toc;

%% assign some additional output
res.acMatrix = metroAcMatrix;
res.acceptanceRate = nAccepted / nMetroIter;
res.rejectionRate = 1 - res.acceptanceRate;
res.sampleDelay = nSampleDelay;
res.integralValue = y;
res.integralRelErr = yErr;
res.mean = metroMean;
res.cov = metroCov;
res.cumulativeTime = cumulativeTime;

end

function acMatrix = computeAcMatrix(dataMatrix, nLag)
    acMatrix = dataMatrix(:, 1:nLag);
    meanVector = mean(dataMatrix, 2);
%     meanVector = dataMatrix(:, 1);
%     figure(101);cla;hold on
    
    for iDim = 1:size(dataMatrix, 1)
        thisDataVector = dataMatrix(iDim, :) - meanVector(iDim);
        
        for iK = 1:nLag
            acMatrix(iDim, iK) = thisDataVector(1:(end-(iK-1))) * ...
                thisDataVector(iK:end)';
        end
        acMatrix(iDim, :) = acMatrix(iDim, :) / acMatrix(iDim, 1);
        
        %! debug code
%         plot(acMatrix(iDim, :))
    end
%     size(dataMatrix)
    
    %! Debug:
%     figure(102),cla,plot(dataMatrix(1,:), dataMatrix(2,:))
    %/
end

function [y yErr] = integrateNormal(fHandle, parMean, parCov, relTol, nBatch, verbose)

intStatus = 1;
nNew = nBatch;
valueVector = [];
probVector = [];

if verbose,disp(['Metropolis: phase 5 started.']);end
relTol = relTol
while intStatus
    % sample new points
    disD = obtainDistributionDiscrete('normal', 'nSample', nNew, ...
        'mean', parMean, 'cov', parCov);
    
    % extract probability
    probVector = [probVector disD.probabilityVector];
    
    % add value for samples
    valueVector = [valueVector fHandle(disD.positionMatrix)];    
    
    % integral
    nSample = length(valueVector);
    y = 1/nSample * sum(valueVector ./ probVector);
    
    % error
    yErr = sqrt(1/nSample/(nSample-1) * sum((...
                valueVector ./ probVector - y).^2));
    yErr = yErr / y;
    
    % decide if finished and estimate how many samples should be sampled in
    % the next step
    if yErr < relTol
        intStatus = 0;
    else
        % Example: if we have taken 1000 samples and obtain an error of
        %   0.04 but require an error of 0.01, we need 4 times more
        %   accuracy. This means we need 4^2=16 times the number of 
        %   samples (15 times more). To ensure, we obtain the required 
        %   error we demand 0.9 times the original relative tolerance
        nNew = round(nSample * (yErr/(0.9*relTol))^2 - nSample);
        nNew = min(nNew, nBatch);
        if verbose,disp(['Metropolis: integration error now at ' ...
                num2str(yErr) ', using ' num2str(nNew) ...
                ' new sample points']);
        end
    end
end


end