// Copyright (c) 2015-2018 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#ifndef variational_bayes_mixture_model_hpp
#define variational_bayes_mixture_model_hpp

#include <array>
#include <vector>
#include <algorithm>
#include <numeric>
#include <iterator>
#include <cstddef>
#include <utility>
#include <cassert>
#include <limits>

#include <boost/optional.hpp>
#include <boost/math/special_functions/digamma.hpp>

#include "utils/maths.hpp"
#include "utils/memory_footprint.hpp"
#include "core/models/haplotype_likelihood_array.hpp"

/**
 *
 * This file contains an implementation of the Variational Bayes mixture model
 * used by some genotype models. The notation follows the documentation.
 *
 */

namespace octopus { namespace model {

// Types needed for Variational Bayes model

struct VariationalBayesParameters
{
    double epsilon = 0.05;
    unsigned max_iterations = 1000;
    double save_memory = false;
};

using ProbabilityVector    = std::vector<double>;
using LogProbabilityVector = std::vector<double>;

template <std::size_t K>
using VBAlpha = std::array<double, K>;
template <std::size_t K>
using VBAlphaVector = std::vector<VBAlpha<K>>;

class VBReadLikelihoodArray
{
public:
    using BaseType = HaplotypeLikelihoodArray::LikelihoodVector;

    VBReadLikelihoodArray() = default;

    explicit VBReadLikelihoodArray(const BaseType&);

    VBReadLikelihoodArray(const VBReadLikelihoodArray&)            = default;
    VBReadLikelihoodArray& operator=(const VBReadLikelihoodArray&) = default;
    VBReadLikelihoodArray(VBReadLikelihoodArray&&)                 = default;
    VBReadLikelihoodArray& operator=(VBReadLikelihoodArray&&)      = default;

    ~VBReadLikelihoodArray() = default;

    void operator=(const BaseType&);
    void operator=(std::reference_wrapper<const BaseType>);
    std::size_t size() const noexcept;
    BaseType::const_iterator begin() const noexcept;
    BaseType::const_iterator end() const noexcept;
    double operator[](const std::size_t n) const noexcept;

private:
    const BaseType* likelihoods;
};

template <std::size_t K>
using VBGenotype = std::array<VBReadLikelihoodArray, K>; // One element per haplotype in genotype (i.e. K)
template <std::size_t K>
using VBGenotypeVector = std::vector<VBGenotype<K>>; // Per element per genotype
template <std::size_t K>
using VBReadLikelihoodMatrix = std::vector<VBGenotypeVector<K>>; // One element per sample

using VBTau = std::vector<double>; // One element per read
template <std::size_t K>
using VBResponsabilityVector = std::array<VBTau, K>; // One element per haplotype in genotype (i.e. K)
template <std::size_t K>
using VBResponsabilityMatrix = std::vector<VBResponsabilityVector<K>>; // One element per sample

template <std::size_t K>
struct VBLatents
{
    ProbabilityVector genotype_posteriors;
    LogProbabilityVector genotype_log_posteriors;
    VBAlphaVector<K> alphas;
    VBResponsabilityMatrix<K> responsabilities;
};

// Main VB method

namespace detail {

using VBExpandedLikelihood = std::vector<double>; // One element per genotype
using VBExpandedGenotype = std::vector<VBExpandedLikelihood>; // One element per read
template <std::size_t K>
using VBExpandedGenotypeVector = std::array<VBExpandedGenotype, K>; // One element per haplotype in genotype
template <std::size_t K>
using VBExpandedLikelihoodMatrix = std::vector<VBExpandedGenotypeVector<K>>; // One element per sample

template <std::size_t K>
auto invert(const VBGenotypeVector<K>& likelihoods)
{
    static_assert(K > 0, "K == 0");
    const auto num_genotypes = likelihoods.size();
    assert(num_genotypes > 0);
    const auto num_reads = likelihoods.front().front().size();
    VBExpandedGenotypeVector<K> result {};
    for (std::size_t k {0}; k < K; ++k) {
        result[k] = VBExpandedGenotype(num_reads, VBExpandedLikelihood(num_genotypes));
        for (std::size_t n {0}; n < num_reads; ++n) {
            for (std::size_t g {0}; g < num_genotypes; ++g) {
                result[k][n][g] = likelihoods[g][k][n];
            }
        }
    }
    return result;
}

template <std::size_t K>
auto invert(const VBReadLikelihoodMatrix<K>& matrix)
{
    VBExpandedLikelihoodMatrix<K> result {};
    result.reserve(matrix.size());
    std::transform(std::cbegin(matrix), std::cend(matrix), std::back_inserter(result),
                   [] (const auto& v) { return invert(v); });
    return result;
}

inline ProbabilityVector& exp(const LogProbabilityVector& log_probabilities, ProbabilityVector& result) noexcept
{
    std::transform(std::cbegin(log_probabilities), std::cend(log_probabilities), std::begin(result),
                   [] (const auto lp) noexcept { return std::exp(lp); });
    return result;
}

inline ProbabilityVector exp(const LogProbabilityVector& log_probabilities)
{
    ProbabilityVector result(log_probabilities.size());
    return exp(log_probabilities, result);
}

inline auto sum(const VBAlpha<2>& alpha) noexcept
{
    return alpha[0] + alpha[1];
}

inline auto sum(const VBAlpha<3>& alpha) noexcept
{
    return alpha[0] + alpha[1] + alpha[2];
}

template <std::size_t K>
auto sum(const VBAlpha<K>& alpha) noexcept
{
    using T = typename VBAlpha<K>::value_type;
    return std::accumulate(std::cbegin(alpha), std::cend(alpha), T {0});
}

template <typename T>
inline auto digamma_diff(const T a, const T b)
{
    using boost::math::digamma;
    return digamma(a) - digamma(b);
}

template <typename T>
inline T log_sum_exp(const std::array<T, 1>& logs)
{
    return logs[0];
}

template <typename T>
inline T log_sum_exp(const std::array<T, 2>& logs)
{
    return maths::log_sum_exp(logs[0], logs[1]);
}

template <typename T>
inline T log_sum_exp(const std::array<T, 3>& logs)
{
    return maths::log_sum_exp(logs[0], logs[1], logs[2]);
}

template <typename T, std::size_t K>
T log_sum_exp(const std::array<T, K>& logs)
{
    return maths::log_sum_exp(logs);
}

template <std::size_t K>
auto count_reads(const VBGenotypeVector<K>& likelihoods) noexcept
{
    return likelihoods[0][0].size();
}

template <std::size_t K>
auto count_reads(const VBExpandedGenotypeVector<K>& likelihoods) noexcept
{
    return likelihoods[0].size();
}

template <std::size_t K>
auto marginalise(const ProbabilityVector& distribution, const VBGenotypeVector<K>& likelihoods,
                 const unsigned k, const std::size_t n) noexcept
{
    using T = ProbabilityVector::value_type;
    return std::inner_product(std::cbegin(distribution), std::cend(distribution),
                              std::cbegin(likelihoods), T {0}, std::plus<> {},
                              [k, n] (const auto p, const auto& haplotype_likelihoods) noexcept {
                                  return p * haplotype_likelihoods[k][n];
                              });
}

template <typename T1, typename T2>
auto inner_product(const T1& lhs, const T2& rhs) noexcept
{
    assert(std::distance(std::cbegin(lhs), std::cend(lhs)) == std::distance(std::cbegin(rhs), std::cend(rhs)));
    using T = typename T1::value_type;
    return std::inner_product(std::cbegin(lhs), std::cend(lhs), std::cbegin(rhs), T {0});
}

template <std::size_t K>
auto marginalise(const ProbabilityVector& distribution, const VBExpandedGenotypeVector<K>& likelihoods,
                 const unsigned k, const std::size_t n) noexcept
{
    return inner_product(distribution, likelihoods[k][n]);
}

template <std::size_t K, typename VBLikelihoodVector_>
VBResponsabilityVector<K>
init_responsabilities(const VBAlpha<K>& prior_alphas,
                      const ProbabilityVector& genotype_probabilities,
                      const VBLikelihoodVector_& read_likelihoods)
{
    using T = typename VBAlpha<K>::value_type;
    std::array<T, K> al; // no need to keep recomputing this
    const auto a0 = sum(prior_alphas);
    for (unsigned k {0}; k < K; ++k) {
        al[k] = digamma_diff(prior_alphas[k], a0);
    }
    const auto N = count_reads(read_likelihoods);
    VBResponsabilityVector<K> result {};
    for (auto& tau : result) tau.resize(N);
    std::array<T, K> ln_rho;
    for (std::size_t n {0}; n < N; ++n) {
        for (unsigned k {0}; k < K; ++k) {
            ln_rho[k] = al[k] + marginalise(genotype_probabilities, read_likelihoods, k, n);
        }
        const auto ln_rho_norm = log_sum_exp(ln_rho);
        for (unsigned k {0}; k < K; ++k) {
            result[k][n] = std::exp(ln_rho[k] - ln_rho_norm);
        }
    }
    return result;
}

template <std::size_t K, typename VBLikelihoodMatrix>
VBResponsabilityMatrix<K>
init_responsabilities(const VBAlphaVector<K>& prior_alphas,
                      const ProbabilityVector& genotype_probabilities,
                      const VBLikelihoodMatrix& read_likelihoods)
{
    const auto S = read_likelihoods.size(); // num samples
    VBResponsabilityMatrix<K> result {};
    result.reserve(S);
    for (std::size_t s {0}; s < S; ++s) {
        result.push_back(init_responsabilities(prior_alphas[s], genotype_probabilities, read_likelihoods[s]));
    }
    return result;
}

template <std::size_t K, typename VBLikelihoodVector_>
void update_responsabilities(VBResponsabilityVector<K>& result,
                             const VBAlpha<K>& posterior_alphas,
                             const ProbabilityVector& genotype_probabilities,
                             const VBLikelihoodVector_& read_likelihoods)
{
    using T = typename VBAlpha<K>::value_type;
    std::array<T, K> al;
    const auto a0 = sum(posterior_alphas);
    for (unsigned k {0}; k < K; ++k) {
        al[k] = digamma_diff(posterior_alphas[k], a0);
    }
    const auto N = count_reads(read_likelihoods);
    std::array<T, K> ln_rho;
    for (std::size_t n {0}; n < N; ++n) {
        for (unsigned k {0}; k < K; ++k) {
            ln_rho[k] = al[k] + marginalise(genotype_probabilities, read_likelihoods, k, n);
        }
        const auto ln_rho_norm = log_sum_exp(ln_rho);
        for (unsigned k {0}; k < K; ++k) {
            result[k][n] = std::exp(ln_rho[k] - ln_rho_norm);
        }
    }
}

// same as init_responsabilities but in-place
template <std::size_t K, typename VBLikelihoodMatrix>
void update_responsabilities(VBResponsabilityMatrix<K>& result,
                             const VBAlphaVector<K>& posterior_alphas,
                             const ProbabilityVector& genotype_probabilities,
                             const VBLikelihoodMatrix& read_likelihoods)
{
    const auto S = read_likelihoods.size();
    for (std::size_t s {0}; s < S; ++s) {
        update_responsabilities(result[s], posterior_alphas[s], genotype_probabilities, read_likelihoods[s]);
    }
}

template <typename T>
inline auto sum(const std::vector<T>& values) noexcept
{
    return std::accumulate(std::cbegin(values), std::cend(values), T {});
}

template <std::size_t K>
void update_alpha(VBAlpha<K>& alpha, const VBAlpha<K>& prior_alpha,
                  const VBResponsabilityVector<K>& taus) noexcept
{
    for (unsigned k {0}; k < K; ++k) {
        alpha[k] = prior_alpha[k] + sum(taus[k]);
    }
}

template <std::size_t K>
void update_alphas(VBAlphaVector<K>& alphas, const VBAlphaVector<K>& prior_alphas,
                   const VBResponsabilityMatrix<K>& responsabilities) noexcept
{
    const auto S = alphas.size();
    assert(S == prior_alphas.size() && S == responsabilities.size());
    for (std::size_t s {0}; s < S; ++s) {
        update_alpha(alphas[s], prior_alphas[s], responsabilities[s]);
    }
}

inline auto marginalise(const VBTau& responsabilities, const VBReadLikelihoodArray& likelihoods) noexcept
{
    assert(responsabilities.size() == likelihoods.size()); // num reads
    return inner_product(responsabilities, likelihoods);
}

template <std::size_t K>
auto marginalise(const VBResponsabilityVector<K>& responsabilities,
                 const VBGenotype<K>& read_likelihoods) noexcept
{
    double result {0};
    for (unsigned k {0}; k < K; ++k) {
        result += marginalise(responsabilities[k], read_likelihoods[k]);
    }
    return result;
}

template <std::size_t K>
auto marginalise(const VBResponsabilityMatrix<K>& responsabilities,
                 const VBReadLikelihoodMatrix<K>& read_likelihoods,
                 const std::size_t g) noexcept
{
    double result {0};
    const auto S = read_likelihoods.size(); // num samples
    assert(S == responsabilities.size());
    for (std::size_t s {0}; s < S; ++s) {
        result += marginalise(responsabilities[s], read_likelihoods[s][g]);
    }
    return result;
}

template <std::size_t K>
void update_genotype_log_posteriors(LogProbabilityVector& result,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const VBResponsabilityMatrix<K>& responsabilities,
                                    const VBReadLikelihoodMatrix<K>& read_likelihoods)
{
    const auto G = result.size();
    for (std::size_t g {0}; g < G; ++g) {
        result[g] = genotype_log_priors[g] + marginalise(responsabilities, read_likelihoods, g);
    }
    maths::normalise_logs(result);
}

inline auto entropy(const VBTau& tau) noexcept
{
    using T = VBTau::value_type;
    return -std::accumulate(std::cbegin(tau), std::cend(tau), T {0},
                            [] (const auto curr, const auto t) noexcept { return curr + (t * std::log(t)); });
}

// E [ln q(Z_s)]
template <std::size_t K>
auto sum_entropies(const VBResponsabilityVector<K>& taus) noexcept
{
    using T = VBTau::value_type;
    return std::accumulate(std::cbegin(taus), std::cend(taus), T {0},
                           [] (const auto curr, const auto& tau) noexcept { return curr + entropy(tau); });
}

template <std::size_t K>
auto calculate_evidence_lower_bound(const VBAlphaVector<K>& prior_alphas,
                                    const VBAlphaVector<K>& posterior_alphas,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const ProbabilityVector& genotype_posteriors,
                                    const LogProbabilityVector& genotype_log_posteriors,
                                    const VBResponsabilityMatrix<K>& taus,
                                    const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                    const boost::optional<double> max_posterior_skip = boost::none)
{
    const auto G = genotype_log_priors.size();
    const auto S = log_likelihoods.size();
    double result {0};
    for (std::size_t g {0}; g < G; ++g) {
        if (!max_posterior_skip || genotype_posteriors[g] >= *max_posterior_skip) {
            auto w = genotype_log_priors[g] - genotype_log_posteriors[g];
            for (std::size_t s {0}; s < S; ++s) {
                w += marginalise(taus[s], log_likelihoods[s][g]);
            }
            result += genotype_posteriors[g] * w;
        }
    }
    for (std::size_t s {0}; s < S; ++s) {
        result += (maths::log_beta(posterior_alphas[s]) - maths::log_beta(prior_alphas[s]));
        result += sum_entropies(taus[s]);
    }
    return result;
}

// Main algorithm - single seed

// Starting iteration with given genotype_log_posteriors
template <std::size_t K, typename VBLikelihoodMatrix1, typename VBLikelihoodMatrix2>
VBLatents<K>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBLikelihoodMatrix1& log_likelihoods1,
                      const VBLikelihoodMatrix2& log_likelihoods2,
                      LogProbabilityVector genotype_log_posteriors,
                      const VariationalBayesParameters& params)
{
    assert(!prior_alphas.empty());
    assert(!genotype_log_priors.empty());
    assert(!log_likelihoods1.empty());
    assert(log_likelihoods1.size() == log_likelihoods2.size());
    assert(prior_alphas.size() == log_likelihoods1.size()); // num samples
    assert(log_likelihoods1.front().size() == genotype_log_priors.size()); // num genotypes
    assert(params.max_iterations > 0);
    auto genotype_posteriors = exp(genotype_log_posteriors);
    auto posterior_alphas = prior_alphas;
    auto responsabilities = init_responsabilities<K>(posterior_alphas, genotype_posteriors, log_likelihoods2);
    assert(responsabilities.size() == log_likelihoods1.size()); // num samples
    auto prev_evidence = std::numeric_limits<double>::lowest();
    for (unsigned i {0}; i < params.max_iterations; ++i) {
        update_genotype_log_posteriors(genotype_log_posteriors, genotype_log_priors, responsabilities, log_likelihoods1);
        exp(genotype_log_posteriors, genotype_posteriors);
        update_alphas(posterior_alphas, prior_alphas, responsabilities);
        auto curr_evidence = calculate_evidence_lower_bound(prior_alphas, posterior_alphas, genotype_log_priors,
                                                            genotype_posteriors, genotype_log_posteriors, responsabilities,
                                                            log_likelihoods1, 1e-10);
        if (curr_evidence <= prev_evidence || (curr_evidence - prev_evidence) < params.epsilon) break;
        prev_evidence = curr_evidence;
        update_responsabilities(responsabilities, posterior_alphas, genotype_posteriors, log_likelihoods2);
    }
    return VBLatents<K> {
        std::move(genotype_posteriors), std::move(genotype_log_posteriors),
        std::move(posterior_alphas), std::move(responsabilities)
    };
}

// Not using inverted log likelihoods
template <std::size_t K>
VBLatents<K>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      LogProbabilityVector genotype_log_posteriors,
                      const VariationalBayesParameters& params)
{
    return run_variational_bayes(prior_alphas, genotype_log_posteriors, log_likelihoods,
                                 log_likelihoods, genotype_log_posteriors, params);
}

// Main algorithm - multiple seed

template <std::size_t K>
bool run_vb_with_matrix_inversion(const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                  const VariationalBayesParameters& params,
                                  const std::vector<LogProbabilityVector>& seeds) noexcept
{
    return !params.save_memory;
}

template <std::size_t K>
std::vector<VBLatents<K>>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      const VariationalBayesParameters& params,
                      std::vector<LogProbabilityVector>&& seeds)
{
    std::vector<VBLatents<K>> result {};
    result.reserve(seeds.size());
    if (run_vb_with_matrix_inversion(log_likelihoods, params, seeds)) {
        const auto inverted_log_likelihoods = invert(log_likelihoods);
        for (auto& seed : seeds) {
            result.push_back(detail::run_variational_bayes(prior_alphas, genotype_log_priors,
                                                           log_likelihoods, inverted_log_likelihoods,
                                                           std::move(seed), params));
        }
    } else {
        for (auto& seed : seeds) {
            result.push_back(detail::run_variational_bayes(prior_alphas, genotype_log_priors,
                                                           log_likelihoods,
                                                           std::move(seed), params));
        }
    }
    return result;
}

// lower-bound calculation

template <std::size_t K>
auto calculate_evidence_lower_bound(const VBAlphaVector<K>& prior_alphas,
                                    const LogProbabilityVector& genotype_log_priors,
                                    const VBReadLikelihoodMatrix<K>& log_likelihoods,
                                    const VBLatents<K>& latents)
{
    return calculate_evidence_lower_bound(prior_alphas, latents.alphas, genotype_log_priors,
                                          latents.genotype_posteriors, latents.genotype_log_posteriors,
                                          latents.responsabilities, log_likelihoods);

}

template <std::size_t K>
void check_normalisation(VBLatents<K>& latents)
{
    const auto total_posterior = std::accumulate(std::cbegin(latents.genotype_posteriors), std::cend(latents.genotype_posteriors), 0.0);
    if (total_posterior > 1.0) {
        for (auto& p : latents.genotype_posteriors) p /= total_posterior;
    }
}

template <std::size_t K>
std::pair<VBLatents<K>, double>
get_max_evidence_latents(const VBAlphaVector<K>& prior_alphas,
                         const LogProbabilityVector& genotype_log_priors,
                         const VBReadLikelihoodMatrix<K>& log_likelihoods,
                         std::vector<VBLatents<K>>&& latents)
{
    std::vector<double> seed_evidences(latents.size());
    std::transform(std::cbegin(latents), std::cend(latents), std::begin(seed_evidences),
                   [&] (const auto& seed_latents) {
                       return calculate_evidence_lower_bound(prior_alphas, genotype_log_priors, log_likelihoods, seed_latents);
                   });
    const auto max_itr = std::max_element(std::cbegin(seed_evidences), std::cend(seed_evidences));
    const auto max_idx = std::distance(std::cbegin(seed_evidences), max_itr);
    return std::make_pair(std::move(latents[max_idx]), *max_itr);
}

} // namespace detail

template <std::size_t K>
std::pair<VBLatents<K>, double>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      const VariationalBayesParameters& params,
                      std::vector<LogProbabilityVector> seeds)
{
    assert(!seeds.empty());
    auto latents = detail::run_variational_bayes(prior_alphas, genotype_log_priors, log_likelihoods, params, std::move(seeds));
    auto result = detail::get_max_evidence_latents(prior_alphas, genotype_log_priors, log_likelihoods, std::move(latents));
    detail::check_normalisation(result.first);
    return result;
}

inline VBReadLikelihoodArray::VBReadLikelihoodArray(const BaseType& underlying_likelihoods)
: likelihoods{std::addressof(underlying_likelihoods)} {}

inline void VBReadLikelihoodArray::operator=(const BaseType& other)
{
    likelihoods = std::addressof(other);
}

inline void VBReadLikelihoodArray::operator=(std::reference_wrapper<const BaseType> other)
{
    likelihoods = std::addressof(other.get());
}

inline std::size_t VBReadLikelihoodArray::size() const noexcept
{
    return likelihoods->size();
}

inline VBReadLikelihoodArray::BaseType::const_iterator VBReadLikelihoodArray::begin() const noexcept
{
    return likelihoods->begin();
}

inline VBReadLikelihoodArray::BaseType::const_iterator VBReadLikelihoodArray::end() const noexcept
{
    return likelihoods->end();
}

inline double VBReadLikelihoodArray::operator[](const std::size_t n) const noexcept
{
    return likelihoods->operator[](n);
}

template <std::size_t K>
std::pair<VBLatents<K>, double>
run_variational_bayes(const VBAlphaVector<K>& prior_alphas,
                      const LogProbabilityVector& genotype_log_priors,
                      const VBReadLikelihoodMatrix<K>& log_likelihoods,
                      const VariationalBayesParameters& params,
                      std::vector<LogProbabilityVector> seeds);

template <std::size_t K>
MemoryFootprint
estimate_memory_requirement(const std::vector<SampleName>& samples,
                            const HaplotypeLikelihoodArray& likelihoods,
                            const std::size_t num_genotypes,
                            VariationalBayesParameters params)
{
    std::size_t bytes {};
    for (const auto& sample : samples) {
        bytes += sizeof(VBReadLikelihoodMatrix<K>);
        bytes += sizeof(VBGenotypeVector<K>) * num_genotypes;
        bytes += sizeof(VBResponsabilityMatrix<K>);
        const auto num_likelihoods = likelihoods.num_likelihoods(sample);
        const auto tau_bytes = num_likelihoods * sizeof(VBTau::value_type);
        bytes += tau_bytes * K + sizeof(VBResponsabilityVector<K>);
        if (!params.save_memory) {
            bytes += sizeof(detail::VBExpandedLikelihoodMatrix<K>);
            auto inverse_bytes = sizeof(detail::VBExpandedLikelihood::value_type) * num_genotypes + sizeof(detail::VBExpandedLikelihood);
            inverse_bytes *= num_likelihoods;
            inverse_bytes += sizeof(detail::VBExpandedGenotype);
            bytes += K * inverse_bytes + sizeof(detail::VBExpandedGenotypeVector<K>);
        }
    }
    return MemoryFootprint {bytes};
};

} // namespace model
} // namespace octopus

#endif


/*

CandyPoker
https://github.com/sweeterthancandy/CandyPoker

MIT License

Copyright (c) 2019 Gerry Candy

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

*/
#include "ps/base/cards.h"
#include "ps/support/command.h"
#include "ps/detail/print.h"
#include "ps/eval/class_cache.h"
#include "ps/eval/holdem_class_vector_cache.h"
#include "app/pretty_printer.h"
#include "app/serialization_util.h"
#include "ps/detail/graph.h"
#include "ps/detail/popcount.h"

#include "ps/sim/computer.h"
#include "ps/sim/game_tree.h"
#include "ps/sim/computer_factory.h"
#include "ps/sim/_extra.h"
#include "ps/sim/solver.h"

#include <boost/any.hpp>
#include <numeric>
#include <boost/timer/timer.hpp>

namespace ps{
namespace sim{
        struct PermutationSolverArguments{
                double clamp_epsilon{1e-6};
                size_t grid_size    {11};
                size_t max_popcount {1};
                size_t dbg_use_threads{true};
                // this is stop wasting CPU time
                size_t max_evaluations{1000};

                void EmitDescriptions(SolverDecl::ArgumentVisitor& V)const{
                        V.DeclArgument("clamp-epsilon" , clamp_epsilon,
                                       "used for clamping close to mixed strategies to non-mixed, "
                                       "too small slower convergence");
                        V.DeclArgument("grid-size" , grid_size, "the size of each grid");
                        V.DeclArgument("max-popcount" , max_popcount, "number of freedoms, O(n!)");
                        V.DeclArgument("dbg-use-threads" , dbg_use_threads, "development aid");
                        V.DeclArgument("max-evaluations", max_evaluations, "upper limit of where to fail");
                }
                void Read(bpt::ptree const& args){
                        grid_size       = args.get<double>("grid-size");
                        max_popcount    = args.get<size_t>("max-popcount");
                        clamp_epsilon   = args.get<double>("clamp-epsilon");
                        dbg_use_threads = args.get<size_t>("dbg-use-threads");
                        max_evaluations = args.get<size_t>("max-evaluations");
                }
        };
        /*
         * This is a wrapper around StateType. This solves the problem
         * of computing the counter strategy and EV in observer functions.
         */
        struct MixedSolutionDescription{
                MixedSolutionDescription(size_t player_index_, std::vector<size_t> const& mixed_)
                        :player_index(player_index_),
                        mixed(mixed_)
                {}
                friend std::ostream& operator<<(std::ostream& ostr, MixedSolutionDescription const& self){
                        ostr << "player_index = " << self.player_index;
                        typedef std::vector<size_t>::const_iterator CI0;
                        const char* comma = "";
                        ostr << "mixed" << " = {";
                        for(CI0 iter= self.mixed.begin(), end=self.mixed.end();iter!=end;++iter){
                                ostr << comma << *iter;
                                comma = ", ";
                        }
                        ostr << "}\n";
                        return ostr;
                }
                size_t player_index;
                std::vector<size_t> mixed;
        };

        inline std::vector<MixedSolutionDescription> MakeMixedSolutionDescription(StateType const& S, StateType const& CS){
                /*
                        Here we want to construct a vector which indicate all the mixed strategis.
                        For example if every was push/fold for hero/villian excep A2o for Hero, and
                        67o and 79o for villian, we would have
                                        ({A2o}, {67o, 79o}).
                 */
                std::vector<MixedSolutionDescription> mixed_info;
                for(size_t idx=0;idx!=S.size();++idx){
                        auto const& c = CS[idx][0];
                        auto const& t = S[idx][0];
                        std::vector<size_t> mixed;
                        for(size_t cid=0;cid!=169;++cid){
                                if( t[cid] == c[cid] )
                                        continue;
                                mixed.push_back(cid);
                        }
                        if( mixed.empty() )
                                continue;
                        mixed_info.emplace_back(idx, std::move(mixed));
                }
                return mixed_info;
        }

        struct PermutationSolver : Solver{

                PermutationSolver( PermutationSolverArguments const& args)
                        : args_{args}
                {}
                virtual boost::optional<StateType> Execute(SolverContext& ctx,
                                                           std::shared_ptr<GameTree> const& gt,
                                                           GraphColouring<AggregateComputer> const& AG,
                                                           StateType const& S0)override
                {

                        PS_LOG(trace) << "----------------- PermutationSolver ---------------";

                        assert( args_.grid_size != 1 );

                        Solution Sol = Solution::MakeWithDeps(gt, AG, S0);
                        auto const& S = Sol.S;

                        auto const& Counter = Sol.Counter;

                        std::vector<MixedSolutionDescription> mixed_info = MakeMixedSolutionDescription(S, Counter);
                        #if 0
                        /*
                                Here we want to construct a vector which indicate all the mixed strategis.
                                For example if every was push/fold for hero/villian excep A2o for Hero, and
                                67o and 79o for villian, we would have
                                                ({A2o}, {67o, 79o}).
                         */
                        for(size_t idx=0;idx!=S.size();++idx){
                                auto const& c = Counter[idx][0];
                                auto const& t = S[idx][0];
                                std::vector<size_t> mixed;
                                for(size_t cid=0;cid!=169;++cid){
                                        if( t[cid] == c[cid] )
                                                continue;
                                        mixed.push_back(cid);
                                }
                                if( mixed.empty() )
                                        continue;
                                mixed_info.emplace_back(idx, std::move(mixed));
                        }

                        for(auto const& mi : mixed_info){
                                std::cout << mi << "\n";
                        }
                        #endif

                        using factor_vector_type = std::vector<size_t>;
                        using factor_set_type = std::vector<factor_vector_type>;
                        using factor_family = std::vector<factor_set_type>;

                        std::vector<factor_family> family_vec;

                        for(auto const& mi : mixed_info){

                                factor_vector_type proto(mi.mixed.size(),1);

                                factor_family family;

                                family.push_back(factor_set_type{proto});

                                factor_set_type s;
                                for(size_t idx=0;idx!=mi.mixed.size();++idx){
                                        auto next = proto;
                                        next[idx] = args_.grid_size;
                                        s.push_back(next);
                                }
                                family.push_back(s);

                                family_vec.push_back(std::move(family));

                        }

                        enum{ Debug = 1 };
                        if( Debug ){
                                std::stringstream sstr;
                                for(size_t idx=0;idx!=family_vec.size();++idx){
                                        if(idx != 0 )
                                                sstr << " x ";
                                        sstr << "{";
                                        for(size_t j=0;j!=family_vec[idx].size();++j){
                                                if(j != 0 )
                                                        sstr << ", ";
                                                sstr << "{";
                                                for(size_t k=0;k!=family_vec[idx][j].size();++k){
                                                        if(k != 0 )
                                                                sstr << ", ";
                                                        sstr << detail::to_string(family_vec[idx][j][k]);
                                                }
                                                sstr << "}";

                                        }
                                        sstr << "}";
                                }
                                std::cout << "sstr.str() => " << sstr.str() << "\n"; // __CandyPrint__(cxx-print-scalar,sstr.str())
                        }

                        size_t upper_bound = ( static_cast<size_t>(1) << mixed_info.size() );

                        std::vector<std::vector<factor_set_type> > cross_products;

                        for(size_t level = 0; level <= args_.max_popcount; ++level ){
                                for(size_t mask = 0; mask != upper_bound; ++mask ){
                                        if( detail::popcount(mask) != level )
                                                continue;
                                        cross_products.emplace_back();
                                        for(size_t idx=0;idx!=mixed_info.size();++idx){
                                                size_t cond = !! ( mask & static_cast<size_t>(1) << idx );
                                                cross_products.back().push_back( family_vec[idx][cond] );
                                        }
                                }
                        }

                        if( Debug ){
                                std::cout << "-------------------------------------\n";
                                for(auto const& cp : cross_products ){
                                        std::stringstream sstr;
                                        for(size_t idx=0;idx!=cp.size();++idx){
                                                if( idx != 0 ) sstr << " x ";
                                                sstr << "{";
                                                for(size_t k=0;k!=cp[idx].size();++k){
                                                        if(k != 0 )
                                                                sstr << ", ";
                                                        sstr << detail::to_string(cp[idx][k]);
                                                }
                                                sstr << "}";
                                        }
                                        std::cout << "sstr.str() => " << sstr.str() << "\n"; // __CandyPrint__(cxx-print-scalar,sstr.str())
                                }
                        }

                        using realization_type = std::vector<factor_vector_type>;
                        using realizations_type = std::vector<realization_type>;
                        realizations_type realizations;
                        auto print_realizations = [&](){
                                std::cout << "-------------------------------------\n";
                                for(auto const& realization : realizations){
                                        std::stringstream sstr;
                                        for(size_t idx=0;idx!=realization.size();++idx){
                                                if( idx != 0 ) sstr << ", ";
                                                sstr << detail::to_string(realization[idx]);
                                        }
                                        std::cout << "sstr.str() => " << sstr.str() << "\n"; // __CandyPrint__(cxx-print-scalar,sstr.str())
                                }
                        };
                        print_realizations();
                        for(std::vector<factor_set_type> const& cp : cross_products){
                                realizations_type sub;
                                sub.emplace_back();
                                for(std::vector<factor_vector_type> const & group : cp ){
                                        auto proto = std::move(sub);
                                        for(factor_vector_type const& item : group ){
                                                auto next = proto;
                                                for(auto p : next){
                                                        p.push_back(item);
                                                        sub.push_back(p);
                                                }
                                        }
                                }
                                std::copy(sub.begin(), sub.end(), std::back_inserter(realizations));
                                print_realizations();
                        }

                        if( Debug ){
                                print_realizations();
                        }

                        std::vector<StateType> SV_family;
                        for(auto const& realization : boost::range::unique(boost::range::sort(realizations))){
                                std::vector<StateType> SV;
                                SV.push_back(S);
                                // for each player
                                for(size_t idx=0;idx!=mixed_info.size();++idx){
                                        auto const& mi = mixed_info[idx];
                                        size_t pid = mi.player_index;
                                        // for each mixed card
                                        for(size_t j=0;j!=mi.mixed.size();++j){

                                                // now N states increase N * (num_steps+1) states

                                                size_t cid = mi.mixed[j];
                                                size_t num_steps = realization[idx][j];
                                                auto proto = std::move(SV);
                                                for(size_t k=0;k<=num_steps;++k){
                                                        double pct = 1.0 / num_steps * k;
                                                        for(auto Q : proto ){
                                                                Q[pid][0][cid] =       pct;
                                                                Q[pid][1][cid] = 1.0 - pct;
                                                                SV.push_back(Q);
                                                        }
                                                }
                                        }
                                }
                                std::copy(SV.begin(), SV.end(), std::back_inserter(SV_family));
                        }

                        std::vector<Solution> solution_candidates;

                        PS_LOG(trace) << "SV_family.size() => " << SV_family.size();

                        // do nothing rather than waste computation on something not tangible
                        if( SV_family.size() > args_.max_evaluations ){
                                PS_LOG(warning) << "Failing as we have " << SV_family.size() << " evaluations to perfom";
                                return {};
                        }

                        #if 0
                        for(size_t idx=0;idx!=SV_family.size();++idx){
                                std::cout << "idx => " << idx << " out of " << SV_family.size() << "\n"; // __CandyPrint__(cxx-print-scalar,idx)
                                auto Q = SV_family[idx];
                                computation_kernel::InplaceClamp(Q, args_.clamp_epsilon);
                                auto candidate = Solution::MakeWithDeps(gt, AG, Q);
                                solution_candidates.push_back(candidate);
                        }
                        #else
                        std::vector<std::function<Solution()> > tasks;
                        for(size_t idx=0;idx!=SV_family.size();++idx){
                                auto atom = [&,idx](){
                                        auto Q = SV_family[idx];
                                        computation_kernel::InplaceClamp(Q, args_.clamp_epsilon);
                                        auto candidate = Solution::MakeWithDeps(gt, AG, Q);
                                        return candidate;
                                };
                                tasks.push_back(atom);
                        }
                        if( args_.dbg_use_threads ){
                                std::vector<std::future<Solution> > futs;
                                for(auto& t : tasks ){
                                        //futs.push_back( std::async(std::launch::async, t) );
                                        futs.push_back( std::async(t) );
                                }
                                for(auto& f : futs){
                                        solution_candidates.push_back(f.get());
                                }
                        } else {
                                for(auto& t : tasks ){
                                        solution_candidates.push_back(t());
                                }
                        }
                        #endif

                        std::vector<Pretty::LineItem> dbg;
                        dbg.push_back(std::vector<std::string>{"n", "|.|", "Gamma", "Mixed"});
                        dbg.push_back(Pretty::LineBreak);

                        std::sort(solution_candidates.begin(), solution_candidates.end());

                        for(auto const& sol : solution_candidates){

                                std::vector<std::string> dbg_line;
                                dbg_line.push_back(boost::lexical_cast<std::string>(sol.Level));
                                dbg_line.push_back(boost::lexical_cast<std::string>(sol.Norm));
                                dbg_line.push_back(detail::to_string(sol.Gamma));
                                dbg_line.push_back(detail::to_string(sol.Mixed));
                                dbg.push_back(std::move(dbg_line));
                        }

                        Pretty::RenderTablePretty(std::cout, dbg);

                        if( solution_candidates.empty() )
                               return {};
                        return solution_candidates.front().S;
                        #if 0
                        if( solution_candidates.size() && solution_candidates.front().Level <= 1 ){
                                auto const& best = solution_candidates.front();
                                std::cout << "detail::to_string(best.Gamma) => " << detail::to_string(best.Gamma) << "\n"; // __CandyPrint__(cxx-print-scalar,detail::to_string(best.Gamma))
                                return best.S;
                        } else {
                                ctx.Message("no best :(");
                        }
                        return {};
                        #endif
                }
                virtual std::string StringDescription()const override{
                        std::stringstream sstr;
                        sstr << "PermutationSolver{}";
                        return sstr.str();
                }
        private:
                PermutationSolverArguments args_;
        };

        struct PermutationSolverDecl : SolverDecl{
                virtual void Accept(ArgumentVisitor& V)const override{
                        PermutationSolverArguments proto;
                        proto.EmitDescriptions(V);
                }
                virtual std::shared_ptr<Solver> Make( bpt::ptree const& args)const override
                {
                        PermutationSolverArguments bargs;
                        bargs.Read(args);
                        return std::make_shared<PermutationSolver>(bargs);
                }
        };

        static SolverRegister<PermutationSolverDecl> PermutationSolverReg("permutation");

        struct SinglePermutationSolver : Solver{

                virtual boost::optional<StateType> Execute(SolverContext& ctx,
                                                           std::shared_ptr<GameTree> const& gt,
                                                           GraphColouring<AggregateComputer> const& AG,
                                                           StateType const& S0)override
                {

                        PS_LOG(trace) << "----------------- SinglePermutationSolver ---------------";

                        SequenceConsumer seq;
                        seq.Consume(Solution::MakeWithDeps(gt, AG, S0));
                        for(size_t loop_count=0;;++loop_count){

                                std::cout << "loop_count => " << loop_count << "\n"; // __CandyPrint__(cxx-print-scalar,loop_count)

                                auto opt_sol = seq.AsOptSolution();
                                BOOST_ASSERT( opt_sol );
                                Solution const& Sol = opt_sol.get();

                                auto const& S  = Sol.S;
                                auto const& CS = Sol.Counter;

                                std::vector<MixedSolutionDescription> mixed_info = MakeMixedSolutionDescription(S, CS);

                                std::vector<StateType> candidates;

                                for(auto const& mi : mixed_info ){
                                        for(auto const& cid : mi.mixed ){
                                                auto next = S;
                                                next[mi.player_index][0][cid] = 1.0;
                                                next[mi.player_index][1][cid] = 0.0;
                                                candidates.push_back(next);
                                                next[mi.player_index][0][cid] = 0.0;
                                                next[mi.player_index][1][cid] = 1.0;
                                                candidates.push_back(next);
                                        }
                                }
                                std::cout << "mixed_info.size() => " << mixed_info.size() << "\n"; // __CandyPrint__(cxx-print-scalar,mixed_info.size())
                                std::vector<std::function<Solution()> > tasks;
                                for(auto const& cand : candidates){
                                        auto atom = [&]()->Solution{
                                                auto solution = Solution::MakeWithDeps(gt, AG, cand);
                                                return solution;
                                        };
                                        tasks.emplace_back(atom);
                                }
                                std::cout << "tasks.size() => " << tasks.size() << "\n"; // __CandyPrint__(cxx-print-scalar,tasks.size())
                                std::vector<std::future<Solution> > futs;
                                for(auto& t : tasks ){
                                        futs.push_back( std::async(std::launch::async, t) );
                                }
                                std::vector<Solution> solution_candidates;
                                for(auto& f : futs){
                                        solution_candidates.push_back(f.get());
                                }
                                #if 0
                                std::mutex mtx;
                                auto child = [&]()mutable{
                                        for(;;){

                                                mtx.lock();
                                                if( tasks.empty() ){
                                                        mtx.unlock();
                                                        return;
                                                }
                                                auto fut = std::move(tasks.back());
                                                tasks.pop_back();
                                                mtx.unlock();
                                                fut();
                                        }
                                };
                                std::vector<std::thread> tg;
                                for(size_t idx=0;idx!=std::thread::hardware_concurrency();++idx){
                                        tg.emplace_back(child);
                                }
                                for(auto& _ : tg){
                                        _.join();
                                }
                                #endif

                                bool do_break = true;
                                std::cout << "HEAD " << Sol.Total << "\n"; // __CandyPrint__(cxx-print-scalar,Sol.Total)
                                for(auto& _ : solution_candidates){
                                        switch(seq.Consume(_)){
                                        case SequenceConsumer::Ctrl_Rejected:
                                                break;
                                        case SequenceConsumer::Ctrl_Accepted:
                                                std::cout << "Accepted\n";
                                                do_break = false;
                                                break;
                                        case SequenceConsumer::Ctrl_Perfect:
                                                return seq.AsOptState();
                                        }
                                }

                                if( do_break ){
                                        break;
                                }
                        }
                        return seq.AsOptState();
                }
                virtual std::string StringDescription()const override{
                        std::stringstream sstr;
                        sstr << "SinglePermutationSolver{}";
                        return sstr.str();
                }
        };

        struct SinglePermutationSolverDecl : SolverDecl{
                virtual void Accept(ArgumentVisitor& V)const override{
                }
                virtual std::shared_ptr<Solver> Make( bpt::ptree const& args)const override
                {
                        return std::make_shared<SinglePermutationSolver>();
                }
        };

        static SolverRegister<SinglePermutationSolverDecl> SinglePermutationSolverReg("single-permutation");

} // end namespace sim
} // end namespace ps


#include <fstream>
#include <utility>
#include <vector>

#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>
#include <boost/log/trivial.hpp>
#include <bitset>

#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/graph/default_device.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/command_line_flags.h"

using tensorflow::Flag;
using tensorflow::Tensor;
using tensorflow::Status;
using tensorflow::string;
using tensorflow::int32;

#include "tensorflow.h"

using namespace std;
using namespace boost;

namespace ml {
    namespace images {

        Inception::Inception(std::string graph_file_name, std::string labels_file_name) {
            tensorflow::GraphDef graph_def;
            Status status = ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("unable to load compute graph at %s, %s") % graph_file_name % status.error_message();
            }
            status = tensorflow::NewSession(tensorflow::SessionOptions(), &session);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("failure creating new session, %s") % status.error_message();
            }
            status = session->Create(graph_def);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("unable to create session for compute graph at %s, %s") % graph_file_name % status.error_message();
            }
            ifstream file(labels_file_name);
            if (!file) {
                BOOST_LOG_TRIVIAL(error) << format("Labels file %s not found.") % labels_file_name;
            }
            labels.clear();
            string line;
            while (getline(file, line)) {
                labels.push_back(line);
            }
            found_label_count = labels.size();
            const int padding = 16;
            while (labels.size() % padding) {
                labels.emplace_back();
            }
        }

        Inception::~Inception() {
            Status session_close = session->Close();
            if (!session_close.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not close session, %s") % session_close.error_message();
            }
        }

        vector<inception_score> Inception::classify(memory_block *block, int width, int height) {
            vector<inception_score> result;

            auto root = tensorflow::Scope::NewRootScope();
            using namespace ::tensorflow::ops;

            tensorflow::StringPiece data(block->memory, block->size);
            Tensor input(tensorflow::DT_STRING, tensorflow::TensorShape());
            (&input)->scalar<string>()() = data.ToString();

            auto reader = Placeholder(root.WithOpName("input"), tensorflow::DataType::DT_STRING);
            const int wanted_channels = 3;
            tensorflow::Output image_reader;

            switch(block->type) {
                case png:
                    image_reader = DecodePng(root.WithOpName("png_reader"), reader, DecodePng::Channels(wanted_channels));
                    break;
                case gif:
                    image_reader = Squeeze(root.WithOpName("squeeze_first_dim"), DecodeGif(root.WithOpName("gif_reader"), reader));
                    break;
                case bmp:
                    image_reader = DecodeBmp(root.WithOpName("bmp_reader"), reader);
                    break;
                case jpg:
                    image_reader = DecodeJpeg(root.WithOpName("jpeg_reader"), reader, DecodeJpeg::Channels(wanted_channels));
                case unkown:
                    BOOST_LOG_TRIVIAL(error) << "could not determine image type to analyze, aborting...";
                    return result;
            }

            auto float_caster = Cast(root.WithOpName("float_caster"), reader, tensorflow::DT_FLOAT);
            auto dims_expander = ExpandDims(root, float_caster, 0);
            auto resized = ResizeBilinear(root, dims_expander, Const(root.WithOpName("size"), {height, width}));

            Div(root.WithOpName("normalized"), Sub(root, resized, {input_mean}), {input_std});

            tensorflow::GraphDef image_graph;
            Status status = root.ToGraphDef(&image_graph);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not create image_graph definition, %s") % status.error_message();
                return result;
            }

            session.reset(tensorflow::NewSession(tensorflow::SessionOptions()));

            status = session->Create(image_graph);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not create image_graph, %s") % status.error_message();
                return result;
            }

            vector<Tensor> outputs;
            vector<pair<string, tensorflow::Tensor>> inputs = { {"input", input}, };
            status = session->Run({inputs}, {"normalized"}, {}, &outputs);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("failure during analysis, %s") % status.error_message();
                return result;
            }

            string output_name = "top_k";
            TopK(root.WithOpName(output_name), outputs[0], min(5, static_cast<int>(found_label_count)));

            tensorflow::GraphDef label_graph;
            status = root.ToGraphDef(&label_graph);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not create label_graph definition, %s") % status.error_message();
                return result;
            }

            session.reset(tensorflow::NewSession(tensorflow::SessionOptions()));

            status = session->Create(label_graph);
            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not create label_graph, %s") % status.error_message();
                return result;
            }

            std::vector<Tensor> out_tensors;
            status = session->Run({}, {output_name + ":0", output_name + ":1"}, {}, &out_tensors));

            if (!status.ok()) {
                BOOST_LOG_TRIVIAL(error) << format("could not create execute run, %s") % status.error_message();
                return result;
            }

            tensorflow::TTypes<float>::Flat scores_flat = out_tensors[0].flat<float>();
            tensorflow::TTypes<int32>::Flat indices_flat = out_tensors[1].flat<int32>();

            for (int pos = 0; pos < min(5, static_cast<int>(found_label_count)); ++pos) {
                const int label_index = indices_flat(pos);
                const float score = scores_flat(pos);
                result.push_back(inception_score{labels[label_index], label_index, score});
            }

            return result;
        }

    }
}


#include "traffic/traffic_info.hpp"
#include "traffic/speed_groups.hpp"

#include "indexer/classificator_loader.hpp"

#include "platform/platform.hpp"

#include "base/assert.hpp"
#include "base/logging.hpp"
#include "base/math.hpp"

#include "std/cstdint.hpp"
#include "std/map.hpp"
#include "std/sstream.hpp"
#include "std/string.hpp"
#include "std/vector.hpp"

#include "pyhelpers/module_version.hpp"
#include "pyhelpers/vector_list_conversion.hpp"
#include "pyhelpers/vector_uint8.hpp"

#include <boost/python.hpp>
#include <boost/python/suite/indexing/map_indexing_suite.hpp>
#include <boost/python/suite/indexing/vector_indexing_suite.hpp>

namespace
{
using namespace boost::python;

struct SegmentSpeeds
{
  SegmentSpeeds() = default;
  SegmentSpeeds(double weightedSpeed, double weightedRefSpeed, double weight)
    : m_weightedSpeed(weightedSpeed), m_weightedRefSpeed(weightedRefSpeed), m_weight(weight)
  {
  }

  double m_weightedSpeed = 0;
  double m_weightedRefSpeed = 0;
  double m_weight = 0;
};

using SegmentMapping = map<traffic::TrafficInfo::RoadSegmentId, SegmentSpeeds>;

string SegmentSpeedsRepr(SegmentSpeeds const & v)
{
  ostringstream ss;
  ss << "SegmentSpeeds("
     << " weighted_speed=" << v.m_weightedSpeed << " weighted_ref_speed=" << v.m_weightedRefSpeed
     << " weight=" << v.m_weight << " )";
  return ss.str();
}

traffic::TrafficInfo::Coloring TransformToSpeedGroups(SegmentMapping const & segmentMapping)
{
  double const kEps = 1e-9;
  traffic::TrafficInfo::Coloring result;
  for (auto const & kv : segmentMapping)
  {
    double const ws = kv.second.m_weightedSpeed;
    double const wrs = kv.second.m_weightedRefSpeed;
    double const w = kv.second.m_weight;
    if (base::AlmostEqualAbs(w, 0.0, kEps))
    {
      LOG(LWARNING, ("A traffic segment has zero weight."));
      continue;
    }
    double const u = ws / w;
    double const v = wrs / w;
    bool const uz = base::AlmostEqualAbs(u, 0.0, kEps);
    bool const vz = base::AlmostEqualAbs(v, 0.0, kEps);
    if (uz && vz)
    {
      result[kv.first] = traffic::SpeedGroup::TempBlock;
    }
    else if (vz)
    {
      LOG(LWARNING, ("A traffic segment has zero reference speed."));
      continue;
    }
    else
    {
      double p = 100.0 * u / v;
      p = base::clamp(p, 0.0, 100.0);
      result[kv.first] = traffic::GetSpeedGroupByPercentage(p);
    }
  }
  return result;
}

string RoadSegmentIdRepr(traffic::TrafficInfo::RoadSegmentId const & v)
{
  ostringstream ss;
  ss << "RoadSegmentId(" << v.m_fid << ", " << v.m_idx << ", " << int(v.m_dir) << ")";
  return ss.str();
}

boost::python::list GenerateTrafficKeys(string const & mwmPath)
{
  vector<traffic::TrafficInfo::RoadSegmentId> result;
  traffic::TrafficInfo::ExtractTrafficKeys(mwmPath, result);
  return std_vector_to_python_list(result);
}

vector<uint8_t> GenerateTrafficValues(vector<traffic::TrafficInfo::RoadSegmentId> const & keys,
                                      boost::python::dict const & segmentMappingDict,
                                      uint8_t useTempBlock)
{
  SegmentMapping segmentMapping;
  boost::python::list mappingKeys = segmentMappingDict.keys();
  for (size_t i = 0; i < len(mappingKeys); ++i)
  {
    object curArg = segmentMappingDict[mappingKeys[i]];
    if (curArg)
      segmentMapping[extract<traffic::TrafficInfo::RoadSegmentId>(mappingKeys[i])] =
          extract<SegmentSpeeds>(segmentMappingDict[mappingKeys[i]]);
  }

  traffic::TrafficInfo::Coloring const knownColors = TransformToSpeedGroups(segmentMapping);
  traffic::TrafficInfo::Coloring coloring;
  traffic::TrafficInfo::CombineColorings(keys, knownColors, coloring);

  vector<traffic::SpeedGroup> values(coloring.size());

  size_t i = 0;
  for (auto const & kv : coloring)
  {
    ASSERT_EQUAL(kv.first, keys[i], ());
    if (useTempBlock == 0 && kv.second == traffic::SpeedGroup::TempBlock)
      continue;

    values[i] = kv.second;
    ++i;
  }
  ASSERT_EQUAL(i, values.size(), ());

  vector<uint8_t> buf;
  traffic::TrafficInfo::SerializeTrafficValues(values, buf);
  return buf;
}

vector<uint8_t> GenerateTrafficValuesFromList(boost::python::list const & keys,
                                              boost::python::dict const & segmentMappingDict)
{
  vector<traffic::TrafficInfo::RoadSegmentId> keysVec =
      python_list_to_std_vector<traffic::TrafficInfo::RoadSegmentId>(keys);

  return GenerateTrafficValues(keysVec, segmentMappingDict, 1 /* useTempBlock */);
}

vector<uint8_t> GenerateTrafficValuesFromBinary(vector<uint8_t> const & keysBlob,
                                                boost::python::dict const & segmentMappingDict,
                                                uint8_t useTempBlock = 1)
{
  vector<traffic::TrafficInfo::RoadSegmentId> keys;
  traffic::TrafficInfo::DeserializeTrafficKeys(keysBlob, keys);

  return GenerateTrafficValues(keys, segmentMappingDict, useTempBlock);
}

void LoadClassificator(string const & classifPath)
{
  GetPlatform().SetResourceDir(classifPath);
  classificator::Load();
}
}  // namespace

BOOST_PYTHON_MODULE(pytraffic)
{
  using namespace boost::python;
  scope().attr("__version__") = PYBINDINGS_VERSION;

  // Register the to-python converters.
  to_python_converter<vector<uint8_t>, vector_uint8t_to_str>();
  vector_uint8t_from_python_str();

  class_<SegmentSpeeds>("SegmentSpeeds", init<double, double, double>())
      .def("__repr__", &SegmentSpeedsRepr)
      .def_readwrite("weighted_speed", &SegmentSpeeds::m_weightedSpeed)
      .def_readwrite("weighted_ref_speed", &SegmentSpeeds::m_weightedRefSpeed)
      .def_readwrite("weight", &SegmentSpeeds::m_weight)
  ;

  class_<traffic::TrafficInfo::RoadSegmentId>("RoadSegmentId", init<uint32_t, uint16_t, uint8_t>())
      .def("__repr__", &RoadSegmentIdRepr)
      .add_property("fid", &traffic::TrafficInfo::RoadSegmentId::GetFid)
      .add_property("idx", &traffic::TrafficInfo::RoadSegmentId::GetIdx)
      .add_property("dir", &traffic::TrafficInfo::RoadSegmentId::GetDir)
  ;

  class_<std::vector<traffic::TrafficInfo::RoadSegmentId>>("RoadSegmentIdVec")
      .def(vector_indexing_suite<std::vector<traffic::TrafficInfo::RoadSegmentId>>());

  enum_<traffic::SpeedGroup>("SpeedGroup")
      .value("G0", traffic::SpeedGroup::G0)
      .value("G1", traffic::SpeedGroup::G1)
      .value("G2", traffic::SpeedGroup::G2)
      .value("G3", traffic::SpeedGroup::G3)
      .value("G4", traffic::SpeedGroup::G4)
      .value("G5", traffic::SpeedGroup::G5)
      .value("TempBlock", traffic::SpeedGroup::TempBlock)
      .value("Unknown", traffic::SpeedGroup::Unknown)
  ;

  def("load_classificator", LoadClassificator);
  def("generate_traffic_keys", GenerateTrafficKeys);
  def("generate_traffic_values_from_list", GenerateTrafficValuesFromList);
  def("generate_traffic_values_from_binary", GenerateTrafficValuesFromBinary,
      (arg("keysBlob"), arg("segmentMappingDict"), arg("useTempBlock") = 1));
}