/*  Multimodal Image Segmentation Tool (MIST)  */
/*  Eelke Visser  */

/*  Copyright (c) 2016 University of Oxford  */

/*  CCOPYRIGHT  */

#include "gibbsshapemodel.h"

using namespace std;
using namespace NEWMAT;

GibbsShapeModel::GibbsShapeModel(const Shape &shape, const std::vector<std::string> &modalitynames, int profilepoints, double profilespacing,
                                 bool usenormalisationmasks) :
    ShapeModel(shape, modalitynames, profilepoints, profilespacing, usenormalisationmasks)
{
    std::random_device rd;
    Reseed(rd());
}

void GibbsShapeModel::Reseed(std::mt19937::result_type seed)
{
    BOOST_LOG_TRIVIAL(info) << "Reseeding MT for shape model. Seed = " << seed;
    m_randomEngine.seed(seed);
}

void GibbsShapeModel::UseGibbs(bool use)
{
    m_useGibbs = use;
}

void GibbsShapeModel::SetGibbsIterations(int iterations)
{
    if (iterations <= m_gibbsBurnIn)
        throw ModelException("Total number of iterations must be larger than number of burn-in iterations");

    m_gibbsIters = iterations;
}

void GibbsShapeModel::SetGibbsBurnIn(int burnin)
{
    if (burnin >= m_gibbsIters)
        throw ModelException("Number of burn-in iterations must be smaller than total number of iterations");

    m_gibbsBurnIn = burnin;
}

ColumnVector GibbsShapeModel::FitImplementation(const std::unordered_map<string, std::vector<ColumnVector> > &data,
                                              ColumnVector &cilower, ColumnVector &ciupper, double cilevel)
{
    // TODO: Add m_useGibbs to class + serialization code
    if (m_useGibbs)
        return DoGibbs(data, cilower, ciupper, cilevel);

    return ShapeModel::FitImplementation(data, cilower, ciupper, cilevel);
}

std::vector<double> GibbsShapeModel::GetConditional(const std::vector<int> &deltas, int vert) const
{
    int steps = m_vertexModels[0]->GetNumberOfSteps();

    ColumnVector p(deltas.size());

    for (int i = 0; i < deltas.size(); i++)
        p(i + 1) = m_profileSpacing * ((steps - 1) / 2.0 - deltas[i]);

    return GetConditionalContinuous(p, vert);
}

ColumnVector GibbsShapeModel::DoGibbs(const std::unordered_map<string, std::vector<ColumnVector> > &data,
                                      ColumnVector &cilower, ColumnVector &ciupper, double cilevel)
{
    std::vector<std::vector<double> > profilelogprobs = GetAllDeltaLikelihoods(data);
    int steps = m_vertexModels[0]->GetNumberOfSteps();

    std::vector<ColumnVector> chain(m_gibbsIters - m_gibbsBurnIn);
    ColumnVector p(m_shape->GetNumberOfVertices());
    p = 0.0;

    for (int v = 0; v < m_shape->GetNumberOfVertices(); v++)
    {
        const std::vector<double> &dl = profilelogprobs[v];

        p(v + 1) = m_profileSpacing * ((steps - 1) / 2.0 - (std::max_element(dl.begin(), dl.end()) - dl.begin()));
        BOOST_LOG_TRIVIAL(debug) << "Using initial displacement " << p(v + 1) << " for vertex " << v;
    }

    for (int iter = 0; iter < m_gibbsIters; iter++)
    {
        BOOST_LOG_TRIVIAL(info) << "Gibbs iteration " << iter
                                << (iter < m_gibbsBurnIn ? " (burn-in)" : "") <<  " ...";

        for (int v = 0; v < m_shape->GetNumberOfVertices(); v++)
        {
            BOOST_LOG_TRIVIAL(debug) << "Old p(" << v + 1 << ") = " << p(v + 1) << " (Vertex " << v << ")";

            std::vector<double> logprobs = GetConditionalContinuous(p, v);

            for (auto lp = logprobs.begin(), plp = profilelogprobs[v].begin(); lp != logprobs.end();)
                *lp++ += *plp++;

            p(v + 1) = GenerateSample(logprobs);
            BOOST_LOG_TRIVIAL(debug) << "New p(" << v + 1 << ") = " << p(v + 1) << " (Vertex " << v << ")";
        }

        if (iter >= m_gibbsBurnIn)
            chain[iter - m_gibbsBurnIn] = p;
    }

    std::vector<std::vector<int> > hists(m_shape->GetNumberOfVertices(), std::vector<int>(steps, 0));
    for (const auto &t : chain)
    {
        for (int v = 0; v < m_shape->GetNumberOfVertices(); v++)
        {
            double val = - t(v + 1) / m_profileSpacing + (steps - 1) / 2.0;

            int bin = static_cast<int>(val + 0.5);

            if (bin < 0)
                bin = 0;

            if (bin >= steps)
                bin = steps - 1;

            hists[v][bin]++;
        }
    }

    ColumnVector displacements(m_shape->GetNumberOfVertices());
    ColumnVector cil(m_shape->GetNumberOfVertices());
    ColumnVector ciu(m_shape->GetNumberOfVertices());

    displacements = 0.0;
    for (int v = 0; v < m_shape->GetNumberOfVertices(); v++)
    {
        BOOST_LOG_TRIVIAL(debug) << "Histogram for vertex " << v << ":";

        int sum = 0;
        int maxval = 0;
        int maxind = 0;
        int ciuind = -1;
        int cilind = 0;

        for (std::size_t i = 0; i < steps; i++)
        {
            BOOST_LOG_TRIVIAL(debug) << "Bin " << i << ": " << hists[v][i];

            if (hists[v][i] > maxval)
            {
                maxval = hists[v][i];
                maxind = i;
            }

            sum += hists[v][i];

            if (sum < (m_gibbsIters - m_gibbsBurnIn) * (0.5 - cilevel / 2))
                ciuind = i;

            if (cilind == 0 && sum > (m_gibbsIters - m_gibbsBurnIn) * (0.5 + cilevel / 2))
                cilind = i;
        }

        displacements(v + 1) = m_profileSpacing * ((steps - 1) / 2.0 - maxind);
        cil(v + 1) = m_profileSpacing * ((steps - 1) / 2.0 - cilind - 0.5);
        ciu(v + 1) = m_profileSpacing * ((steps - 1) / 2.0 - ciuind - 0.5);
    }

    cilower = cil;
    ciupper = ciu;

    return displacements;
}

double GibbsShapeModel::GenerateSample(const std::vector<double> &logprobs)
{
    int steps = m_vertexModels[0]->GetNumberOfSteps();

    auto max = std::max_element(logprobs.cbegin(), logprobs.cend());
    std::vector<double> probs(logprobs.size());
    std::transform(logprobs.cbegin(), logprobs.cend(), probs.begin(), [&](double p){ return std::exp(p - *max); });
    std::vector<double> cumprobs(probs.size() + 1);
    cumprobs[0] = 0.0;
    std::partial_sum(probs.cbegin(), probs.cend(), cumprobs.begin() + 1);

    std::uniform_real_distribution<> dist(0.0, *(cumprobs.cend() - 1));
    double rn = dist(m_randomEngine);
    auto upper = std::find_if(cumprobs.cbegin(), cumprobs.cend(), [&](double p){ return p >= rn; });

    double sample;
    if (upper != cumprobs.cbegin())
    {
        auto lower = upper - 1;
        sample = (lower - cumprobs.cbegin()) + (rn - *lower) / (*upper - *lower);
    }
    else
        sample = 0.0;

    // Needed to center cumprobs probability regions around probs bin centres
    sample -= 0.5;

    return m_profileSpacing * ((steps - 1) / 2.0 - sample);
}
