#include <Eigen/Core>
#include <chrono>
#include <cstdlib>
#include <exception>
// #include <kids/toltec/toltec.h>
#include <thread>
#include <tula/config.h>
#include <tula/container.h>
#include <tula/formatter/matrix.h>
#include <tula/formatter/ptr.h>
#include <tula/grppiex.h>
#include <tula/mpi.h>
#include <yaml-cpp/yaml.h>

namespace predefs {

// we need to be careful about the int type used here as
// we may have very large array in the memory.

// check eigen index type
static_assert(std::is_same_v<std::ptrdiff_t, Eigen::Index>,
              "UNEXPECTED EIGEN INDEX TYPE");
using index_t = std::ptrdiff_t;
using shape_t = Eigen::Matrix<index_t, 2, 1>;
using data_t = double;
} // namespace predefs

/**
 * @brief The BlockAccessor struct
 * This provides methods to access array in block notions
 */
template <typename Array, auto n_dims_ = 2> struct BlockAccessor {
    using array_t = Array;
    using index_t = predefs::index_t;
    using bins_t = Eigen::Matrix<index_t, 2, Eigen::Dynamic>;
    constexpr static const index_t n_dims = n_dims_;
    static_assert(n_dims > 0, "NDIMS HAS TO BE POSITIVE");
    static_assert(n_dims <= 2, "NDIMS > 2 IS NOT IMPLEMENTED");
    constexpr static const index_t left = 0;
    constexpr static const index_t right = 1;
    BlockAccessor(array_t array_) : array(std::move(array_)) {}
    array_t array;
    std::array<bins_t, n_dims> all_bins;
    /**
     * @brief Return the bins for dim.
     */
    template <index_t dim> constexpr auto bins() -> decltype(auto) {
        return all_bins[dim];
    }
    /**
     * @brief Return the number of bins for dim.
     */
    template <index_t dim> constexpr index_t n_bins() const {
        return bins<dim>().cols();
    }
    /**
     * @brief Set bins with bin matrix of shape (2, nbins).
     */
    template <index_t dim> void set_bins(bins_t bins_) {
        bins<dim>() = std::move(bins_);
    }
    /**
     * @brief Set bins with vector of pairs of (left, right).
     */
    template <index_t dim>
    void set_bins(const std::vector<std::pair<index_t, index_t>> &bins_) {
        set_bins(eigen_utils::tomat(bins_));
    }
    /**
     * @brief Set bins with vector of bin edges of size (nbins + 1).
     */
    template <index_t dim>
    void set_bins(const std::vector<index_t> &bin_edges) {
        static_assert(dim > 0 && dim < n_dims, "INVAID DIM");
        auto edges = eigen_utils::asvec(bin_edges);
        index_t n_bins = edges.size() - 1;
        bins<dim>().resize(2, n_bins);
        bins<dim>().row(left) = edges.head(n_bins);
        bins<dim>().row(right) = edges.tail(n_bins);
    }

    /**
     * @brief Returns number of blocks
     */
    constexpr index_t n_blocks() const {
        return std::reduce(all_bins.begin(), all_bins.end(), 1,
                           [](const auto &lhs, const auto &rhs) {
                               return lhs.cols() * rhs.cols();
                           });
    }
    template <index_t n_dims__ = n_dims, REQUIRES_V(n_dims__ == 1)>
    auto block(index_t i) {
        constexpr decltype(auto) bins_ = bins<0>();
        assert(i < bins_.cols());
        auto i0 = bins_.coeff(left, i);
        auto ni = bins_.coeff(right, i) - i0;
        return array.segment(i0, ni);
    }
    template <index_t n_dims__ = n_dims, REQUIRES_V(n_dims__ == 2)>
    auto block(index_t i, index_t j) {
        constexpr decltype(auto) ibins = bins<0>();
        assert(i < ibins.cols());
        auto i0 = ibins.coeff(left, i);
        auto ni = ibins.coeff(right, i) - i0;

        constexpr decltype(auto) jbins = bins<1>();
        assert(j < jbins.cols());
        auto j0 = jbins.coeff(left, j);
        auto nj = jbins.coeff(right, j) - j0;
        return array.block(i0, j0, ni, nj);
    }

    friend std::ostream &operator<<(std::ostream &os, const BlockAccessor &b) {
        return os << fmt::format("BlockAccessor(array{} all_bins={})", b.array,
                                 b.all_bins);
    }
};

/**
 * @brief The YamlConfig struct
 * This is a thin wrapper around YAML::Node.
 */
struct YamlConfig {
    using key_t = std::string;
    using value_t = YAML::Node;
    using storage_t = YAML::Node;
    YamlConfig() = default;
    explicit YamlConfig(storage_t node) : m_node(std::move(node)) {}
    std::string dump_to_str() const { return fmt::format("{}", m_node); }
    static auto load_from_str(std::string s) {
        return YamlConfig(YAML::Load(s));
    }
    friend std::ostream &operator<<(std::ostream &os,
                                    const YamlConfig &config) {
        return os << config.dump_to_str();
    }

    template <typename T> auto get_typed(const key_t &key) const {
        return m_node[key].as<T>(key);
    }
    template <typename T> auto get_typed(const key_t &key, T &&defval) const {
        decltype(auto) node = m_node[key];
        if (node.IsDefined() && !node.IsNull()) {
            return node.as<T>(key);
        }
        return FWD(defval);
    }

    auto get_str(const key_t &key) const { return get_typed<std::string>(key); }
    auto get_str(const key_t &key, const std::string &defval) const {
        return get_typed<std::string>(key, std::string(defval));
    }

    decltype(auto) operator[](const key_t &key) const { return m_node[key]; }

    bool has(const key_t &key) const { return m_node[key].IsDefined(); }
    template <typename T> bool has_typed(const key_t &key) const {
        if (!has(key)) {
            return false;
        }
        try {
            get_typed<T>(key);
        } catch (YAML::BadConversion) {
            return false;
        }
        return true;
    }

private:
    storage_t m_node;
};

/// @brief The mixin class for config handling.
template <typename Derived> struct ConfigMixin {
    using config_t = YamlConfig;

private:
    using Self = ConfigMixin<Derived>;
    config_t m_config;
    struct derived_has_validate_config {
        define_has_member_traits(Derived, validate_config);
        constexpr static auto value = has_validate_config::value;
    };

public:
    ConfigMixin() = default;
    template <typename... Args> ConfigMixin(Args &&... args) {
        set_config(FWD(args)...);
    }
    const config_t &config() { return m_config; }
    void set_config(config_t config, bool validate = true) {
        if (validate) {
            if constexpr (derived_has_validate_config::value) {
                if (auto opt_errors = Derived::validate_config(config);
                    opt_errors.has_value()) {
                    SPDLOG_ERROR("invalid config:\n{}\nerrors: {}", config,
                                 opt_errors.value());
                } else {
                    SPDLOG_TRACE("set config validated");
                }
            } else {
                SPDLOG_WARN(
                    "set config validation requested but no validator found");
            }
        } else {
            SPDLOG_TRACE("set config without validation");
        }
        m_config = std::move(config);
    }
};

/// @brief Interface base class
template <typename Derived, typename IO_> struct InterfaceBase {
    using IO = IO_;
    IO io;
};

/// @ brief Interface to TolTEC detector network
struct ToltecInterface : InterfaceBase<ToltecInterface, int> {
    using Base = InterfaceBase<ToltecInterface, int>;
};

/// @ brief Interface to HWP
struct HwpInterface : InterfaceBase<HwpInterface, int> {
    using Base = InterfaceBase<HwpInterface, int>;
};

/// @ brief Interface to LMT
struct LmtInterface : InterfaceBase<LmtInterface, int> {
    using Base = InterfaceBase<LmtInterface, int>;
};

/**
 * @brief The InterfaceRegistry struct
 * This is a helper class to access the actual interface implementation.
 */
struct InterfaceRegistry {
    enum class Interface;

    // dispatcher helper
    template <typename Case, typename... Rest> struct dispatcher_impl {
        template <Interface interface>
        using io_t = meta::switch_t<interface, Case, Rest...>;
        using variant_t =
            std::variant<typename Case::type, typename Rest::type...>;
    };
    /////////////////////////////////////////////
    // Register interface classes here
    META_ENUM(Interface, int, Toltec, Hwp, Lmt);
    using dispatcher =
        dispatcher_impl<meta::case_t<Interface::Toltec, ToltecInterface>,
                        meta::case_t<Interface::Hwp, HwpInterface>,
                        meta::case_t<Interface::Lmt, LmtInterface>>;
    /////////////////////////////////////////////
    template <Interface interface> using io_t = dispatcher::io_t<interface>;
    using variant_t = dispatcher::variant_t;
};

/**
 * @brief The Coordinator struct
 * This wraps around the config object and provides
 * high level methods in various ways to setup the MPI runtime
 * with node-local and cross-node environment.
 */
struct Coordinator : ConfigMixin<Coordinator> {
    using config_t = ConfigMixin::config_t;
    using comm_t = mpi_utils::comm;
    using index_t = predefs::index_t;
    using shape_t = predefs::shape_t;
    constexpr static index_t not_a_worker = -1;

    Coordinator(const comm_t &comm_, const comm_t &comm_local_,
                MPI_Win &win_local_)
        : ConfigMixin<Coordinator>{}, comm(comm_), comm_local(comm_local_),
          win_local(win_local_) {}
    const comm_t &comm;       // This is all, including cross-node
    const comm_t &comm_local; // node-local comm
    MPI_Win &win_local;       // window object for local shm

    constexpr auto n_nodes() const { return comm.n_nodes(); }
    constexpr auto node_index() const { return comm.node_index(); }
    const auto &node_name() const { return comm.node_name(); }

    // global mpi context
    constexpr auto rank() const { return comm.rank(); }
    constexpr auto master_rank() const { return 0; }
    constexpr auto is_master() const { return rank() == master_rank(); }

    // local mpi context
    constexpr auto rank_local() const { return comm_local.rank(); }
    constexpr auto master_rank_local() const { return 0; }
    constexpr auto is_master_local() const {
        return rank_local() == master_rank_local();
    }
    constexpr auto worker_index_local() const {
        if (is_master_local()) {
            return not_a_worker;
        }
        return rank_local() -
            ((rank_local() > master_rank_local()) ? index_t(1) : 0);
    }
    constexpr auto n_workers_local() const { return comm_local.size() - 1; }

    /**
     * Sync the config.
     * This is done to the global comm.
     */
    void sync() {

        auto config_str = comm.bcast_str(
            master_rank(), [&]() { return config().dump_to_str(); });

        if (is_master()) {
            SPDLOG_TRACE("config:\n{}", config());
        } else {
            // here we just disable the validation since it is already done
            set_config(config_t::load_from_str(config_str), false);
            SPDLOG_TRACE("synced config from master");
        }
        comm.barrier();
        // collect inputs
        const auto &inputs = collect_inputs();
        if (is_master_local()) {
            SPDLOG_TRACE("inputs: {}", inputs);
        }
    }

    /**
     * @brief The Observation struct
     * This represents a single observation that contains a set of data items
     */
    struct Observation : ConfigMixin<Observation> {
        /**
         * @brief The DataItem struct
         * This represent a single data item that belongs to a particular
         * observation
         */
        struct DataItem : ConfigMixin<DataItem> {

            DataItem(config_t config_)
                : ConfigMixin<DataItem>{std::move(config_)},
                  interface(config().get_str("interface")),
                  filepath(config().get_str("filepath")) {
                // initalize io
                //                 meta::switch_invoke<InterfaceRegistry::InterfaceKind>(
                //                     [&](auto interfacekind_) {
                //                         constexpr auto interfacekind =
                //                             DECAY(interfacekind_)::value;
                //                         io = InterfaceRegistry::template
                //                         io_t<interfacekind>();
                //                     },
                //                     interface_kind());
            }
            std::string interface{};
            std::string filepath{};
            friend std::ostream &operator<<(std::ostream &os,
                                            const DataItem &d) {
                return os << fmt::format("DataItem(interface={} filepath={})",
                                         d.interface, d.filepath);
            }
            InterfaceRegistry::variant_t io;
            static auto validate_config(const config_t &config)
                -> std::optional<std::string> {
                std::vector<std::string> missing_keys;
                for (const auto &key : {"interface", "filepath"}) {
                    if (!config.has(key)) {
                        missing_keys.push_back(key);
                    }
                }
                if (missing_keys.empty()) {
                    return std::nullopt;
                }
                return fmt::format("missing keys={}", missing_keys);
            }
            index_t buffer_size() const { return 4880 * 500; }
            shape_t buffer_shape() const {
                shape_t s;
                s << 4880, 500;
                return s;
            }
        };
        Observation(config_t config_)
            : ConfigMixin<Observation>{config_}, name{config().get_str(
                                                     "name", "unnamed")} {
            // initialize the data_items
            //             auto node = config()["data_items"];
            //             assert(node.IsSequence());
            //             for (std::size_t i = 0; i < node.size(); ++i) {
            //                 data_items.emplace_back(config_t(node[i]));
            //             }
        }
        std::string name;
        // std::vector<DataItem> data_items{};
        static auto validate_config(const config_t &config)
            -> std::optional<std::string> {
            if (config.has("data_items")) {
                return std::nullopt;
            }
            return fmt::format("missing key={}", "data_items");
        }
        friend std::ostream &operator<<(std::ostream &os,
                                        const Observation &obs) {
            return os << fmt::format("Observation(name={})", obs.name);
        }
    };

    // io_buffer
    using payloads_buffer_data_t = predefs::data_t;
    std::size_t io_buffer_size() { return 4000 * 4880 * 10; }

    constexpr auto n_inputs() const { return m_inputs.size(); }

private:
    using input_t = Observation;
    using payload_t = Observation::DataItem;
    // this is all the inputs
    std::vector<input_t> m_inputs;
    // this is the current payload
    std::vector<payload_t> m_payloads;

    // collect input data
    const std::vector<input_t> &collect_inputs() {
        m_inputs.clear();
        // this is run for all processes so that everyone has the necessary info
        // to collect payloads
        auto index = node_index();
        SPDLOG_TRACE("get input data for node={} name={}", index, node_name());
        std::vector<input_t> results;
        auto c = config()["io"]["inputs"];
        assert(c.IsSequence());
        for (std::size_t i = 0; i < c.size(); i++) {
            if (index == (meta::size_cast<decltype(index)>(i) % n_nodes())) {
                results.emplace_back(config_t{c[i]});
            }
        }
        m_inputs = std::move(results);

        if (is_master_local()) {
            // report number of obs to process
            SPDLOG_DEBUG("collected n_inputs={} on node={} name={}",
                         m_inputs.size(), index, node_name());
        }
        return m_inputs;
    }

public:
    // collect payload
    const std::vector<payload_t> &collect_payloads(std::size_t input_index) {
        m_payloads.clear();
        if (is_master_local()) {
            // master local does not get payload
        } else {
            // construct payload
            // distribute inputs through out worker ranks
            auto index = worker_index_local();
            SPDLOG_TRACE("get payload for input={} worker={} local_rank={}",
                         m_inputs.at(input_index), index, rank_local());
            std::vector<payload_t> results;
            auto c = m_inputs.at(input_index).config()["data_items"];
            assert(c.IsSequence());
            for (std::size_t i = 0; i < c.size(); i++) {
                if (index ==
                    (meta::size_cast<decltype(index)>(i) % n_workers_local())) {
                    results.emplace_back(config_t{c[i]});
                }
            }
            m_payloads = std::move(results);
            SPDLOG_DEBUG("collected n_payloads={} on worker={} local_rank{} "
                         "node={} name={}",
                         m_payloads.size(), index, rank_local(), node_index(),
                         node_name());
        }
        return m_payloads;
    }

    using payloads_buffer_t = mpi_utils::Span<payloads_buffer_data_t, index_t>;
    using payloads_buffer_accessor_t =
        BlockAccessor<payloads_buffer_t::mat_t, 2>;
    constexpr static auto payloads_buffer_time_axis = 0;
    constexpr static auto payloads_buffer_interface_axis = 1;
    auto allocate_payloads_buffer() {
        // get total size of current payloads, and allocate the shm
        // compute payloads buffer size for each node
        shape_t buffer_shape_all = shape_t::Zero();
        auto shape_concat = [](const auto &lhs, const auto &rhs) {
            shape_t r;
            // concat columns
            r(0) = std::max(lhs(0), rhs(0));
            r(1) = lhs(1) + rhs(1);
            return r;
        };
        {
            shape_t buffer_shape = std::transform_reduce(
                m_payloads.begin(), m_payloads.end(), shape_t::Zero().eval(),
                shape_concat,
                [](const auto &v) -> shape_t { return v.buffer_shape(); });
            SPDLOG_TRACE("payloads buffer size={} worker_index={}",
                         buffer_shape, worker_index_local());
            {
                // shape concat
                MPI_Allreduce(&buffer_shape(0), &buffer_shape_all(0), 1,
                              MPI_UTILS_DECLTYPE(buffer_shape(0)), MPI_MAX,
                              comm_local);
                MPI_Allreduce(&buffer_shape(1), &buffer_shape_all(1), 1,
                              MPI_UTILS_DECLTYPE(buffer_shape(1)), MPI_SUM,
                              comm_local);
            }
        }
        if (is_master_local()) {
            SPDLOG_TRACE("payloads buffer shape all {}", buffer_shape_all);
        }
        // allocate memory and initialize accessor
        auto mem = payloads_buffer_t::allocate_shared(master_rank_local(),
                                                      buffer_shape_all.prod(),
                                                      comm_local, &win_local);
        payloads_buffer_accessor_t buffer{mem.asmat(buffer_shape_all)};
        // set up row bins. This is done according to the
        // buffer.set_bins<payloads_buffer_time_axis>();
        std::map<index_t, index_t> map_playload_to_block{};
        return buffer;
    }
};

int main(int argc, char **argv) {

    mpi_utils::env mpi_env(argc, argv); // RAII MPI env
    mpi_env.set_exception_on_error();
    mpi_utils::comm comm; // MPI_COMM_WORLD
    mpi_utils::logging_init(comm);
    mpi_utils::pprint_node_ranks(0, comm, [&](const auto &node_ranks) {
        SPDLOG_INFO("{}", mpi_env);
        SPDLOG_INFO("{}", node_ranks);
    });

    mpi_utils::comm comm_shm{comm.split_shared()};
    {
        // check if shm comm is valid.
        const auto master_rank = 0;
        const auto is_master = comm.rank() == master_rank;
        // check env
        // we need to the io procs and the master on the same node
        if (is_master) {
            if (mpi_env.memory_model !=
                mpi_utils::env::WinMemoryModel::Unified) {
                SPDLOG_ERROR("unified memory model is required, found {}",
                             mpi_env.memory_model);
                MPI_Abort(comm, EXIT_FAILURE);
            }
            if (comm_shm.size() != comm.size()) {
                SPDLOG_ERROR("shared memory access across all processes is "
                             "required, found {}/{}",
                             comm_shm.size(), comm.size());
                MPI_Abort(comm, EXIT_FAILURE);
            }
        }
    }

    // io config
    using coordinator_t = Coordinator;
    using index_t = coordinator_t::index_t;
    using config_t = coordinator_t::config_t;

    MPI_Win win_shm;
    coordinator_t co{comm, comm_shm, win_shm};
    if (co.is_master()) {
        const char *test_config{
            R"(---
description: |-
    This file is provided as an example of the citlali config file.
io:
    inputs:
      - name: obs1
        data_items:
          - interface: toltec0
            filepath: "toltec0.nc"
          - interface: toltec1
            filepath: "toltec1.nc"
          - interface: toltec2
            filepath: "toltec2.nc"
          - interface: lmt
            filepath: "lmt.nc"
          - interface: hwp
            filepath: "hwp.nc"
        cal_items:
          - name: beammap
            filepath: "beammap.nc"
          - name: detector_maskk
            indices: [::10]
      - name: obs2
        data_items:
          - interface: toltec0
            filepath: "toltec0.nc"
          - interface: toltec1
            filepath: "toltec1.nc"
          - interface: toltec2
            filepath: "toltec2.nc"
          - interface: lmt
            filepath: "lmt.nc"
          - interface: hwp
            filepath: "hwp.nc"
        cal_items:
          - name: beammap
            filepath: "beammap.nc"
          - name: detector_maskk
            indices: [::10]
    time_chunking:
        enabled: yes
        method:
            value: "fixed"
            choises: ["hold_signal", "fixed_length"]
            description: |-
                hold_signal: Define chunks using the telescope hold signal.
                fixed_length: Each chunk has specified length.
        parameters:
           hold_signal:
                value: 1
                description: |-
                    value: Continuous samples with the hold signal of this value will be recongnized as one chunk.
           fixed_length:
                value: 10
                unit: second
                description: |-
                    value: The length of the time chunk.
                    unit: The unit of said value, e.g., sample, second, etc.
reduce:
    map_making:
)"};
        SPDLOG_TRACE("yaml config in:\n{}", test_config);
        co.set_config(config_t::load_from_str(test_config));
        /// VALIDATE_CONFIG(config, )
        /// for all types that consumes config.
    }

    // make sure all ranks have access to the config.
    co.sync();
    auto n_obs = co.n_inputs();
    for (std::size_t i = 0; i < n_obs; ++i) {
        const auto &payloads = co.collect_payloads(i);
        SPDLOG_TRACE("payloads: {}", payloads);
        auto buffer = co.allocate_payloads_buffer();
        if (co.is_master_local()) {
            buffer.array(0) = 100;
        }
        MPI_Win_sync(win_shm);
        SPDLOG_TRACE("payloads buffer {}", buffer);
        co.comm.barrier();
        MPI_Abort(comm, 0);
    }

    // coordinator settings

    int arg{0};

    if (co.is_master()) {
        arg = 5;
        SPDLOG_TRACE("load config wait={}", arg);
        std::this_thread::sleep_for(std::chrono::seconds(arg));
        MPI_Bcast(&arg, 1, MPI_INT, 0, comm);
    } else {
        MPI_Bcast(&arg, 1, MPI_INT, 0, comm);
        SPDLOG_TRACE("got config wait={}", arg);
    }

    // intialize the shared memory buffer for io
    //     auto coor =
    //         mpi_utils::Span<coordinator_t::io_buffer_data_t>::allocate_shared(
    //             co.master_rank(), co.io_buffer_size(), comm_shm, &win);

    if (co.is_master_local()) {
        //
        // using citlali_co_t = int;
        // citlali_co_t cco(co.config());
        // allocate all buffers map, coadded map. per map product.
        // cco.allocate_buffers(telescope_data);
        auto ex = tula::grppi_utils::Mode::omp;
        SPDLOG_TRACE("use grppi ex {}", ex);
        grppi::pipeline(
            tula::grppi_utils::dyn_ex(ex),
            [&]() -> std::optional<int> {
                static int rtc = 0;
                // returns upon each recieved rtc
                std::this_thread::sleep_for(std::chrono::seconds(arg));
                MPI_Win_sync(win_shm);
                //  SPDLOG_TRACE("io_buffer{}",
                //              io_buffer.asvec().head(comm.size()));
                return ++rtc;
            },
            [&](int rtc) {
                // cco.process(rtc);
                SPDLOG_TRACE("processing rtc={}", rtc);
                std::this_thread::sleep_for(std::chrono::seconds(arg));
            });
    } else {
        SPDLOG_TRACE("reading rtc={}", 0);
        std::this_thread::sleep_for(std::chrono::seconds(arg));
        SPDLOG_TRACE("reading rtc={} done", 0);
        auto rank = comm_shm.rank();
        // io_buffer.data[rank] = rank;
        MPI_Win_sync(win_shm);
    }
    MPI_Win_free(&win_shm);
    return EXIT_SUCCESS;
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T.Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SHIFT_LEFT_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_SHIFT_LEFT_HPP_INCLUDED
#include <boost/simd/detail/overload.hpp>
#include <boost/simd/meta/hierarchy/simd.hpp>
#include <boost/simd/function/simd/bitwise_cast.hpp>
#include <boost/simd/detail/assert_utils.hpp>
#include <boost/dispatch/meta/as_integer.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;

  BOOST_DISPATCH_OVERLOAD(shift_left_
                         , (typename A0,typename A1,typename X)
                         , bd::cpu_
                         , bs::pack_<bd::floating_<A0>,X>
                         , bd::scalar_<bd::integer_<A1>>
                         )
  {
    BOOST_FORCEINLINE A0 operator()( const A0& a0, const A1& a1) const BOOST_NOEXCEPT
    {
      using int_type = bd::as_integer_t<A0, signed>;
      BOOST_ASSERT_MSG(assert_good_shift<A0>(a1), "shift_left: a shift is out of range");
      return bitwise_cast<A0>( shift_left(bitwise_cast<int_type>(a0), a1) );
    }
  };
} } }

#endif


/* Copyright 2020 Gopinath Chennupati, Raviteja Vangara, Namita Kharat, Erik Skau and Boian Alexandrov,
Triad National Security, LLC. All rights reserved
This program was produced under U.S. Government contract 89233218CNA000001 for Los Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC for the U.S. Department of Energy/National Nuclear Security Administration. All rights in the program are reserved by Triad National Security, LLC, and the U.S. Department of Energy/National Nuclear Security Administration. The Government is granted for itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license in this material to reproduce, prepare derivative works, distribute copies to the public, perform publicly and display publicly, and to permit others to do so.
*/

#ifndef DISTNMF_DISTREORDER_HPP_
#define DISTNMF_DISTREORDER_HPP_

#include <unistd.h>
#include <armadillo>
#include <string>
#include "../planc-master/common/utils.hpp"
#include "../planc-master/common/distutils.hpp"
#include "../planc-master/distnmf/mpicomm.hpp"

/**
 * File name formats
 */

namespace planc {

template <class MATTYPE>
class DistReOrder {
 private:
  const MPICommunicator& m_mpicomm;
  int m_ownedm;
  int m_ownedn;
  int m_globalm;
  int m_globaln;

  int m_pr;
  int m_pc;
  int m_k;

  UROWVEC localMaxIDx; // vector of max element indices from each proc
  ROWVEC localMax;     // vector of max elements from each proc

  UROWVEC globalMaxIDx; // vector of max element indices from all proc(s)
  ROWVEC globalMax;     // vector of max elements from all proc(s)

  UROWVEC finalIDx;     //vector of final max element indices;
                        //still contains the local indices
  UROWVEC globalIDx;    //Contains the global row indexes for m_k columns
  UVEC sortIDx;         //sort indices of W
  UVEC sharedSortIDx;   //rank 0 initialized with same sortIDx for pr * pc
  UVEC globalSortIDx;   // global sort indices of W

  UMAT globalIMx;       // matrix of final max-val indices across all proc(s)
  MAT globalM;          // matrix of final max elements across all proc(s)

  MAT W, H;             //Both the low rank factors
  MAT new_W, new_H;     //reordered low rank factors

  /**
   * Allocates matrices and vectors
   */
  void allocateData() {
    // Init vectors
    /*
    DISTPRINTINFO("k::"<< this->k <<"::localm::"<< this->m_ownedm <<"::localn::"
                    << this->m_ownedn << "::globalm::" << this->m_globalm <<
                    "::globaln::"<< this->m_globaln <<"::MPI_SIZE::"<< MPI_SIZE);
    */
    localMax.zeros(this->m_k);
    localMaxIDx.zeros(this->m_k);
    globalMaxIDx.zeros(this->m_pr * this->m_pc * this->m_k);
    globalMax.zeros(this->m_pr * this->m_pc * this->m_k);
    finalIDx.zeros(this->m_k);
    globalIDx.zeros(this->m_k);
    sortIDx.zeros(this->m_k);
    sharedSortIDx.zeros(this->m_pr * this->m_pc, 1);
    globalSortIDx.zeros(this->m_k);

    // Init matrices
    globalIMx.zeros(this->m_pr * this->m_pc, this->m_k);
    globalM.zeros(this->m_pr * this->m_pc, this->m_k);

    // low-rank factor matrices
    new_W.zeros(this->W.n_rows, this->W.n_cols);
    new_H.zeros(this->H.n_rows, this->H.n_cols);
  }

  void freeMatrices() {
    localMax.clear();
    localMaxIDx.clear();
    globalMaxIDx.clear();
    globalMax.clear();
    finalIDx.clear();
    globalIDx.clear();
    globalIMx.clear();
    sortIDx.clear();
    globalSortIDx.clear();
    globalM.clear();
    W.clear();
    H.clear();
    new_W.clear();
    new_H.clear();
  }

 public:
    DistReOrder<MATTYPE>(const MAT &leftlowrankfactor, const MAT &rightlowrankfactor,
            const MPICommunicator& communicator, const int k)
        : m_mpicomm(communicator) {
        assert(leftlowrankfactor.n_cols == rightlowrankfactor.n_cols);
        this->m_k = k;
        this->m_pr = NUMROWPROCS;
        this->m_pc = NUMCOLPROCS;
        //Assign the low-rank factors
        this->W = leftlowrankfactor;
        this->H = rightlowrankfactor;
        this->m_ownedm  = leftlowrankfactor.n_rows;
        this->m_ownedn = rightlowrankfactor.n_cols;
        allocateData();
        PRINTROOT("distreorder()::constructor succesful");
    }

    ~DistReOrder() {
      //freeMatrices();
    }

    void finalGlobalSortIDx() {
        /* local Max and IDx */
        localMax = arma::max(this->W,0);
        localMaxIDx = arma::index_max(this->W, 0) + this->m_ownedm;
        //TODO! Need to take care of the topk maxima and indices here
        /* Global Max and IDx */
        MPITIC;  // gather globalMax, globalMaxIDx
        int sendcnt = this->m_k;
        int recvcnt = this->m_k;
        globalMax.zeros();
        MPI_Gather(localMax.memptr(), sendcnt, MPI_DOUBLE, globalMax.memptr(),
                recvcnt, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        globalMaxIDx.zeros();
        MPI_Gather(localMaxIDx.memptr(), sendcnt, MPI_DOUBLE, globalMaxIDx.memptr(),
                recvcnt, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        double temp = MPITOC;  // gather globalMax, globalMaxIDx
        //TODO! Need to report time later
        /* finalIDx, globalM, globalIMx and globalIDx */
        if(MPI_RANK == 0)   {
            globalIMx = arma::reshape(globalMaxIDx.t(), this->m_k,
                                (this->m_pr * this->m_pc)).t();
            globalM = arma::reshape(globalMax, this->m_k,
                                (this->m_pr * this->m_pc)).t();
            //globalM.print("globalM = ");
            //globalIMx.print("globalIMx = ");
            finalIDx = arma::index_max(globalM, 0);
            for(int ki = 0; ki < this->m_k; ki++)   {
                globalIDx(ki) = globalIMx(finalIDx(ki),ki);
            }
            //globalIDx.print("globalIDx = ");
            sortIDx = arma::sort_index(globalIDx);
            sortIDx.print("Rank 0 sortIDx = ");
            sharedSortIDx = arma::repmat(sortIDx, this->m_pr * this->m_pc, 1);
            //H.print("H in rank 0 = ");
            //globalSortIDx.print("Rank 0 globalSortIDx = ");
        }
        //Scatter the sortIDx to all the procs in MPI_COMM_WORLD
        globalSortIDx.zeros();
        MPI_Scatter(this->sharedSortIDx.memptr(), sendcnt, MPI_DOUBLE, this->globalSortIDx.memptr(),
                    recvcnt, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    }

    void reorderW()   {
        finalGlobalSortIDx();
        MPI_Barrier(MPI_COMM_WORLD);
        //Now is the time for reordering the columns in W
        this->sortIDx.print("local sortIDx = ");
        this->globalSortIDx.print("globalSortIDX = ");
        //W.print("local W = ");
        //new_W.zeros();
        for(int col = 0; col < this->m_k; col++)    {
            //W is stored as row-major order
            new_W.col(col) = W.col(globalSortIDx(col));
        }
        if(MPI_RANK == 0) { new_W.print("local new_W = "); std::cout<<"current k "<<this->m_k<<std::endl; }
    }

    void reorderH()   {
        /*
        if(arma::sum(sortIDx) == 0){
            finalGlobalSortIDx();
        }*/
        MPI_Barrier(MPI_COMM_WORLD);
        //Now is the time for reordering the rows in H
        //new_H.zeros();
        for(int col = 0; col < this->m_k; col++)    {
            //In H also, swapping the cols, because, it is in column-major order
            new_H.col(col) = H.col(this->globalSortIDx(col));
        }
    }

    /// Returns the left low rank factor matrix W
    MAT getLeftLowRankFactor() { return new_W; }
    /// Returns the right low rank factor matrix H
    MAT getRightLowRankFactor() { return new_H; }

  }; //class DistReOrder

}  // namespace planc

// run with mpi run 3.
void testDistClust(char argc, char* argv[]) {
  planc::MPICommunicator mpicomm(argc, argv);
  planc::DistReOrder<MAT> dc(arma::randn(10, 4),
                  arma::randu(4,8), mpicomm, 8);;

  dc.reorderW();
}

#endif  // DISTNMF_DISTREORDER_HPP_


/*
 * This file is part of bogus, a C++ sparse block matrix library.
 *
 * Copyright 2013 Gilles Daviet <gdaviet@gmail.com>
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

#ifndef BOGUS_EIGEN_BLOCK_CONTAINERS_HPP
#define BOGUS_EIGEN_BLOCK_CONTAINERS_HPP

#include <Eigen/StdVector>
#include "../Utils/CppTools.hpp"

namespace bogus {

// We do not have to use the specialized allocator if the size is dynamic or not a multiple of 16 bytes
template< typename Scalar, int Rows, int Cols, int Options, int MaxRows, int MaxCols >
struct ResizableSequenceContainer< Eigen::Matrix<Scalar, Rows, Cols, Options, MaxRows, MaxCols> >
{
	typedef Eigen::Matrix<Scalar, Rows, Cols, Options, MaxRows, MaxCols> BlockType ;
	typedef typename TypeSwapIf<
		Rows == Eigen::Dynamic || Cols == Eigen::Dynamic || 0 != ( static_cast< std::size_t >( Rows * Cols * sizeof( Scalar ) ) & 0xf ),
		std::vector< BlockType, Eigen::aligned_allocator<BlockType> >,
		std::vector< BlockType > >::First Type ;
} ;

} //namespace bogus

#endif


#include <boost/math/concepts/real_type_concept.hpp>


/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#ifndef __itkEvaluateTractogramDirectionsFilter_cpp
#define __itkEvaluateTractogramDirectionsFilter_cpp

#include "itkEvaluateTractogramDirectionsFilter.h"
#include <itkImageRegionConstIterator.h>
#include <itkImageRegionIterator.h>
#include <itkImageDuplicator.h>
#include <boost/progress.hpp>

#define _USE_MATH_DEFINES
#include <math.h>

namespace itk {

template< class PixelType >
EvaluateTractogramDirectionsFilter< PixelType >
::EvaluateTractogramDirectionsFilter():
    m_ReferenceImageSet(nullptr),
    m_IgnoreMissingDirections(false),
    m_Eps(0.0001),
    m_UseInterpolation(false)
{
    this->SetNumberOfOutputs(1);
}

template< class PixelType >
itk::Vector<PixelType, 3> EvaluateTractogramDirectionsFilter< PixelType >::GetItkVector(double point[3])
{
    itk::Vector<PixelType, 3> itkVector;
    itkVector[0] = point[0];
    itkVector[1] = point[1];
    itkVector[2] = point[2];
    return itkVector;
}

template< class PixelType >
vnl_vector_fixed<PixelType, 3> EvaluateTractogramDirectionsFilter< PixelType >::GetVnlVector(double point[3])
{
    vnl_vector_fixed<PixelType, 3> vnlVector;
    vnlVector[0] = point[0];
    vnlVector[1] = point[1];
    vnlVector[2] = point[2];
    return vnlVector;
}

template< class PixelType >
vnl_vector_fixed<PixelType, 3> EvaluateTractogramDirectionsFilter< PixelType >::GetVnlVector(Vector<PixelType,3>& vector)
{
    vnl_vector_fixed<PixelType, 3> vnlVector;
    vnlVector[0] = vector[0];
    vnlVector[1] = vector[1];
    vnlVector[2] = vector[2];
    return vnlVector;
}

template< class PixelType >
itk::Point<PixelType, 3> EvaluateTractogramDirectionsFilter< PixelType >::GetItkPoint(double point[3])
{
    itk::Point<PixelType, 3> itkPoint;
    itkPoint[0] = point[0];
    itkPoint[1] = point[1];
    itkPoint[2] = point[2];
    return itkPoint;
}

template< class PixelType >
void EvaluateTractogramDirectionsFilter< PixelType >::GenerateData()
{
    if (m_Tractogram.IsNull() || m_ReferenceImageSet.IsNull())
        return;

    if (m_UseInterpolation)
        MITK_INFO << "Using trilinear interpolation";
    else
        MITK_INFO << "Using nearest neighbor interpolation";

    if (m_IgnoreMissingDirections)
        MITK_INFO << "Ignoring missing directions";
    else
        MITK_INFO << "Penalizing missing directions";

    // angular error image
    typename OutputImageType::Pointer outputImage = OutputImageType::New();
    outputImage->SetOrigin( m_ReferenceImageSet->GetElement(0)->GetOrigin() );
    outputImage->SetRegions( m_ReferenceImageSet->GetElement(0)->GetLargestPossibleRegion() );
    outputImage->SetSpacing( m_ReferenceImageSet->GetElement(0)->GetSpacing() );
    outputImage->SetDirection( m_ReferenceImageSet->GetElement(0)->GetDirection() );
    outputImage->Allocate();
    outputImage->FillBuffer(0.0);

    DoubleImageType::Pointer counterImage = DoubleImageType::New();
    counterImage->SetOrigin( m_ReferenceImageSet->GetElement(0)->GetOrigin() );
    counterImage->SetRegions( m_ReferenceImageSet->GetElement(0)->GetLargestPossibleRegion() );
    counterImage->SetSpacing( m_ReferenceImageSet->GetElement(0)->GetSpacing() );
    counterImage->SetDirection( m_ReferenceImageSet->GetElement(0)->GetDirection() );
    counterImage->Allocate();
    counterImage->FillBuffer(0.0);

    std::vector< DoubleImageType::Pointer > directionFound;
    for (int i=0; i<m_ReferenceImageSet->Size(); i++)
    {
        DoubleImageType::Pointer check = DoubleImageType::New();
        check->SetOrigin( m_ReferenceImageSet->GetElement(0)->GetOrigin() );
        check->SetRegions( m_ReferenceImageSet->GetElement(0)->GetLargestPossibleRegion() );
        check->SetSpacing( m_ReferenceImageSet->GetElement(0)->GetSpacing() );
        check->SetDirection( m_ReferenceImageSet->GetElement(0)->GetDirection() );
        check->Allocate();
        check->FillBuffer(90);
        directionFound.push_back(check);
    }

    if (m_MaskImage.IsNull())
    {
        m_MaskImage = UCharImageType::New();
        m_MaskImage->SetOrigin( outputImage->GetOrigin() );
        m_MaskImage->SetRegions( outputImage->GetLargestPossibleRegion() );
        m_MaskImage->SetSpacing( outputImage->GetSpacing() );
        m_MaskImage->SetDirection( outputImage->GetDirection() );
        m_MaskImage->Allocate();
        m_MaskImage->FillBuffer(1);
    }

    m_MeanAngularError = 0.0;
    m_MedianAngularError = 0;
    m_MaxAngularError = 0.0;
    m_MinAngularError = itk::NumericTraits<float>::max();
    m_VarAngularError = 0.0;
    m_AngularErrorVector.clear();

    float minSpacing = 1;
    if(outputImage->GetSpacing()[0]<outputImage->GetSpacing()[1] && outputImage->GetSpacing()[0]<outputImage->GetSpacing()[2])
        minSpacing = outputImage->GetSpacing()[0];
    else if (outputImage->GetSpacing()[1] < outputImage->GetSpacing()[2])
        minSpacing = outputImage->GetSpacing()[1];
    else
        minSpacing = outputImage->GetSpacing()[2];
    FiberBundleType::Pointer fiberBundle = m_Tractogram->GetDeepCopy();
    fiberBundle->ResampleFibers(minSpacing/10);

    MITK_INFO << "Evaluating tractogram";
    vtkSmartPointer<vtkPolyData> fiberPolyData = fiberBundle->GetFiberPolyData();
    boost::progress_display disp( fiberBundle->GetNumFibers() );
    for( int i=0; i<fiberBundle->GetNumFibers(); i++ )
    {
        vtkCell* cell = fiberPolyData->GetCell(i);
        int numPoints = cell->GetNumberOfPoints();
        vtkPoints* points = cell->GetPoints();

        if (numPoints<2)
            continue;

        for( int j=0; j<numPoints; j++)
        {
            double* temp = points->GetPoint(j);
            itk::Point<PixelType, 3> vertex = GetItkPoint(temp);
            itk::Vector<PixelType> v = GetItkVector(temp);

            itk::Vector<PixelType, 3> dir(3);
            if (j<numPoints-1)
                dir = GetItkVector(points->GetPoint(j+1))-v;
            else
                dir = v-GetItkVector(points->GetPoint(j-1));

            vnl_vector_fixed< PixelType, 3 > fiberDir = GetVnlVector(dir);
            fiberDir.normalize();

            itk::Index<3> idx;
            itk::ContinuousIndex<float, 3> contIndex;
            m_MaskImage->TransformPhysicalPointToIndex(vertex, idx);
            m_MaskImage->TransformPhysicalPointToContinuousIndex(vertex, contIndex);

            if (!m_UseInterpolation)    // use nearest neighbour interpolation
            {
                if (!m_MaskImage->GetLargestPossibleRegion().IsInside(idx) || m_MaskImage->GetPixel(idx)<=0)
                    continue;

                double angle = 90;
                int usedIndex = -1;
                for (int k=0; k<m_ReferenceImageSet->Size(); k++)     // and each test direction
                {
                    vnl_vector_fixed< PixelType, 3 > refDir = m_ReferenceImageSet->GetElement(k)->GetPixel(idx).GetVnlVector();
                    if (refDir.magnitude()>m_Eps)  // normalize if not null
                        refDir.normalize();
                    else
                        continue;

                    // calculate angle between directions
                    double tempAngle = acos(fabs(dot_product(refDir, fiberDir)))*180.0/M_PI;
                    directionFound.at(k)->SetPixel(idx, tempAngle);

                    if (tempAngle < angle)
                    {
                        angle = tempAngle;
                        usedIndex = k;
                    }
                }
                if (usedIndex>=0)
                    directionFound.at(usedIndex)->SetPixel(idx, -1);
                else if (m_IgnoreMissingDirections)
                    angle = 0;

                counterImage->SetPixel(idx, counterImage->GetPixel(idx)+1);
                outputImage->SetPixel(idx, outputImage->GetPixel(idx)+angle);
                continue;
            }

            double frac_x = contIndex[0] - idx[0];
            double frac_y = contIndex[1] - idx[1];
            double frac_z = contIndex[2] - idx[2];
            if (frac_x<0)
            {
                idx[0] -= 1;
                frac_x += 1;
            }
            if (frac_y<0)
            {
                idx[1] -= 1;
                frac_y += 1;
            }
            if (frac_z<0)
            {
                idx[2] -= 1;
                frac_z += 1;
            }

            // use trilinear interpolation
            itk::Index<3> newIdx;
            for (int x=0; x<2; x++)
            {
                frac_x = 1-frac_x;
                for (int y=0; y<2; y++)
                {
                    frac_y = 1-frac_y;
                    for (int z=0; z<2; z++)
                    {
                        frac_z = 1-frac_z;

                        newIdx[0] = idx[0]+x;
                        newIdx[1] = idx[1]+y;
                        newIdx[2] = idx[2]+z;

                        double frac = frac_x*frac_y*frac_z;

                        // is position valid?
                        if (!m_MaskImage->GetLargestPossibleRegion().IsInside(newIdx) || m_MaskImage->GetPixel(newIdx)<=0)
                            continue;

                        double angle = 90;
                        int usedIndex = -1;
                        for (int k=0; k<m_ReferenceImageSet->Size(); k++)     // and each test direction
                        {
                            vnl_vector_fixed< PixelType, 3 > refDir = m_ReferenceImageSet->GetElement(k)->GetPixel(newIdx).GetVnlVector();
                            if (refDir.magnitude()>m_Eps)  // normalize if not null
                                refDir.normalize();
                            else
                                continue;

                            // calculate angle between directions
                            double tempAngle = acos(fabs(dot_product(refDir, fiberDir)))*180.0/M_PI;
                            directionFound.at(k)->SetPixel(newIdx, tempAngle);

                            if (tempAngle < angle)
                            {
                                angle = tempAngle;
                                usedIndex = k;
                            }
                        }
                        if (usedIndex>=0)
                            directionFound.at(usedIndex)->SetPixel(newIdx, -1);
                        else if (m_IgnoreMissingDirections)
                            angle = 0;

                        counterImage->SetPixel(newIdx, counterImage->GetPixel(newIdx)+1);
                        outputImage->SetPixel(newIdx, outputImage->GetPixel(newIdx)+frac*angle);
                    }
                }
            }
        }
        ++disp;
    }

    ImageRegionIterator< OutputImageType > oit(outputImage, outputImage->GetLargestPossibleRegion());
    ImageRegionIterator< UCharImageType > mit(m_MaskImage, m_MaskImage->GetLargestPossibleRegion());
    ImageRegionIterator< DoubleImageType > cit(counterImage, counterImage->GetLargestPossibleRegion());

    if (!m_IgnoreMissingDirections)
        MITK_INFO << "Creatings statistics and accounting for missing directions";
    else
        MITK_INFO << "Creatings statistics";
    boost::progress_display disp2(outputImage->GetLargestPossibleRegion().GetNumberOfPixels());
    while( !oit.IsAtEnd() )
    {
        if ( mit.Get()>0)
        {
            if ( cit.Get()>m_Eps )
                oit.Set(oit.Get()/cit.Get());

            if (!m_IgnoreMissingDirections)
            {
                int missingCount = 0;

                if ( cit.Get()>m_Eps )
                    missingCount++;

                for (int i=0; i<directionFound.size(); i++)
                {
                    if ( directionFound.at(i)->GetPixel(oit.GetIndex())>0 && m_ReferenceImageSet->GetElement(i)->GetPixel(oit.GetIndex()).GetVnlVector().magnitude()>m_Eps )
                    {
                        oit.Set(oit.Get()+directionFound.at(i)->GetPixel(oit.GetIndex()));
                        missingCount++;
                    }
                }
                if (missingCount>1)
                    oit.Set(oit.Get()/missingCount);
            }

            if (oit.Get()>m_MaxAngularError)
                m_MaxAngularError = oit.Get();
            if (oit.Get()<m_MinAngularError)
                m_MinAngularError = oit.Get();
        }

        ++oit;
        ++cit;
        ++mit;
        ++disp2;
    }

    this->SetNthOutput(0, outputImage);
}

}

#endif // __itkEvaluateTractogramDirectionsFilter_cpp


#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <Eigen/Dense>

#include <algorithm>
#include <cstdint>
#include <limits>
#include <map>
#include <memory>
#include <random>
#include <stdexcept>
#include <utility>
#include <vector>

namespace py = pybind11;

namespace ote {

const double EPS = 1e-8;
const double EPS2 = 1e-5;
const double EPS3 = 1e-3;

inline int32_t sign(float x) {
  if (fabs(x) < EPS) {
    throw std::logic_error("computing sign of ~0");
  }
  if (x > 0) return 1;
  return -1;
}

class OTEstimators {
 public:
  using NumPyFloatArray = py::array_t<float, py::array::c_style>;
  using NumPyIntArray = py::array_t<int32_t, py::array::c_style>;

  using EigenVector = Eigen::Matrix<float, 1, Eigen::Dynamic, Eigen::RowMajor>;
  using EigenMatrix =
      Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
  using Matrix = Eigen::Map<EigenMatrix>;

  OTEstimators() : stage(0) {}

  void load_vocabulary(NumPyFloatArray points) {
    if (stage != 0) {
      throw std::logic_error(
          "load_vocabulary() should be called once in the beginning");
    }
    stage = 1;
    py::buffer_info buf = points.request();
    if (buf.ndim != 2) {
      throw std::logic_error(
          "load_vocabulary() expects a two-dimensional NumPy array");
    }
    auto n = buf.shape[0];
    auto d = buf.shape[1];
    dictionary = std::make_unique<Matrix>(static_cast<float *>(buf.ptr), n, d);
    auto cmin = std::numeric_limits<float>::max();
    auto cmax = std::numeric_limits<float>::min();
    for (ssize_t i = 0; i < n; ++i) {
      for (ssize_t j = 0; j < d; ++j) {
        cmin = std::min(cmin, (*dictionary)(i, j));
        cmax = std::max(cmax, (*dictionary)(i, j));
      }
    }
    auto delta = cmax - cmin;
    cmin -= delta;
    std::random_device rd;
    std::mt19937_64 gen(rd());
    std::uniform_real_distribution<float> shift_gen(0.0, delta);
    std::vector<std::pair<float, float>> bounding_box;
    for (ssize_t i = 0; i < d; ++i) {
      auto s = shift_gen(gen);
      bounding_box.push_back(std::make_pair(cmin + s, cmax + s));
    }
    std::vector<int32_t> all;
    for (ssize_t i = 0; i < n; ++i) {
      all.push_back(i);
    }
    leaf.resize(n);
    build_quadtree(all, bounding_box, 0, -1);
    num_queries = 0;
    marked.resize(parents.size());
    for (auto &x : marked) {
      x = -1;
    }
    node_id.resize(parents.size());
  }

  void load_dataset(
      const std::vector<std::vector<std::pair<int32_t, float>>> &dataset) {
    if (stage != 1) {
      throw std::logic_error(
          "load_dataset() should be called once after calling "
          "load_vocabulary()");
    }
    stage = 2;
    if (dataset.empty()) {
      throw std::logic_error("the dataset can't be empty");
    }
    for (auto &measure : dataset) {
      check_measure(measure);
      dataset_embedding.push_back(compute_embedding(measure));
    }

    raw_dataset = dataset;
    for (auto &measure : raw_dataset) {
      std::sort(measure.begin(), measure.end());
    }

    means.resize(dataset.size(), dictionary->cols());
    for (size_t i = 0; i < dataset.size(); ++i) {
      means.row(i) =
          dictionary->row(dataset[i][0].first) * dataset[i][0].second;
      for (size_t j = 1; j < dataset[i].size(); ++j) {
        means.row(i) +=
            dictionary->row(dataset[i][j].first) * dataset[i][j].second;
      }
    }
    query_mean.resize(dictionary->cols());
    distances.resize(dataset.size());
  }

  void means_rank(const std::vector<std::pair<int32_t, float>> &query,
                  NumPyIntArray input_ids, NumPyIntArray output_ids,
                  NumPyFloatArray output_scores, bool to_sort) {
    check_stage();
    check_measure(query);
    check_input_output_arrays(input_ids, output_ids, output_scores);
    query_mean = dictionary->row(query[0].first) * query[0].second;
    for (size_t j = 1; j < query.size(); ++j) {
      query_mean =
          query_mean + dictionary->row(query[j].first) * query[j].second;
    }
    auto input_ids_buf = input_ids.request();
    auto input_ids_raw = static_cast<int32_t *>(input_ids_buf.ptr);
    int32_t k1 = input_ids_buf.shape[0];
    for (int32_t i = 0; i < k1; ++i) {
      float score = (query_mean - means.row(input_ids_raw[i])).squaredNorm();
      distances[i] = std::make_pair(score, input_ids_raw[i]);
    }
    select_topk_aux(k1, output_ids, output_scores, to_sort);
  }

  void overlap_rank(const std::vector<std::pair<int32_t, float>> &query,
                    NumPyIntArray input_ids, NumPyIntArray output_ids,
                    NumPyFloatArray output_scores, bool to_sort) {
    check_stage();
    check_measure(query);
    auto query_copy = query;
    std::sort(query_copy.begin(), query_copy.end());
    auto input_ids_buf = input_ids.request();
    auto input_ids_raw = static_cast<int32_t *>(input_ids_buf.ptr);
    int32_t k1 = input_ids_buf.shape[0];
    for (int32_t i = 0; i < k1; ++i) {
      auto &point = raw_dataset[input_ids_raw[i]];
      float score = 0.0;
      size_t qp = 0;
      size_t dp = 0;

      while (qp < query_copy.size() && dp < point.size()) {
        if (query_copy[qp].first < point[dp].first) {
          ++qp;
        } else if (query_copy[qp].first > point[dp].first) {
          ++dp;
        } else {
          score += 1;
          ++dp;
          ++qp;
        }
      }
      distances[i] = std::make_pair(-score, input_ids_raw[i]);
    }
    select_topk_aux(k1, output_ids, output_scores, to_sort);
  }

  void quadtree_rank(const std::vector<std::pair<int32_t, float>> &query,
                     NumPyIntArray input_ids, NumPyIntArray output_ids,
                     NumPyFloatArray output_scores, bool to_sort) {
    check_stage();
    check_measure(query);
    check_input_output_arrays(input_ids, output_ids, output_scores);
    auto query_embedding = compute_embedding(query);
    auto input_ids_buf = input_ids.request();
    auto input_ids_raw = static_cast<int32_t *>(input_ids_buf.ptr);
    int32_t k1 = input_ids_buf.shape[0];
    for (int32_t i = 0; i < k1; ++i) {
      auto &point_embedding = dataset_embedding[input_ids_raw[i]];
      float score = 0.0;
      size_t qp = 0;
      size_t dp = 0;
      while (qp < query_embedding.size() || dp < point_embedding.size()) {
        if (qp == query_embedding.size()) {
          score += point_embedding[dp].second;
          ++dp;
        } else if (dp == point_embedding.size()) {
          score += query_embedding[qp].second;
          ++qp;
        } else if (query_embedding[qp].first < point_embedding[dp].first) {
          score += query_embedding[qp].second;
          ++qp;
        } else if (point_embedding[dp].first < query_embedding[qp].first) {
          score += point_embedding[dp].second;
          ++dp;
        } else {
          score +=
              fabs(query_embedding[qp].second - point_embedding[dp].second);
          ++qp;
          ++dp;
        }
      }
      distances[i] = std::make_pair(score, input_ids_raw[i]);
    }
    select_topk_aux(k1, output_ids, output_scores, to_sort);
  }

  void flowtree_rank(const std::vector<std::pair<int32_t, float>> &query,
                     NumPyIntArray input_ids, NumPyIntArray output_ids,
                     NumPyFloatArray output_scores, bool to_sort) {
    check_stage();
    check_measure(query);
    check_input_output_arrays(input_ids, output_ids, output_scores);
    auto input_ids_buf = input_ids.request();
    auto input_ids_raw = static_cast<int32_t *>(input_ids_buf.ptr);
    int32_t k1 = input_ids_buf.shape[0];
    for (int32_t i = 0; i < k1; ++i) {
      auto cur_id = input_ids_raw[i];
      auto score = flowtree_query(query, raw_dataset[cur_id]);
      distances[i] = std::make_pair(score, cur_id);
    }
    select_topk_aux(k1, output_ids, output_scores, to_sort);
  }

  void select_topk(NumPyIntArray input_ids, NumPyFloatArray input_scores,
                   NumPyIntArray output_ids, NumPyFloatArray output_scores,
                   bool to_sort) {
    check_input_output_arrays(input_ids, input_scores, output_ids,
                              output_scores);
    auto input_scores_buf = input_scores.request();
    auto input_ids_buf = input_ids.request();
    auto input_scores_raw = static_cast<float *>(input_scores_buf.ptr);
    auto input_ids_raw = static_cast<int32_t *>(input_ids_buf.ptr);
    int32_t k1 = input_ids_buf.shape[0];
    for (int32_t i = 0; i < k1; ++i) {
      distances[i] = std::make_pair(input_scores_raw[i], input_ids_raw[i]);
    }
    select_topk_aux(k1, output_ids, output_scores, to_sort);
  }

 private:
  std::vector<int32_t> parents;
  std::vector<int32_t> leaf;
  std::vector<int32_t> marked;
  int32_t num_queries;
  std::vector<int32_t> node_id;
  std::vector<int32_t> id_node;
  std::vector<std::vector<int32_t>> subtree;
  std::vector<std::vector<std::pair<float, int32_t>>> excess;
  std::vector<float> delta_node;
  std::unique_ptr<Matrix> dictionary;
  std::vector<int32_t> unleaf;
  std::vector<std::vector<std::pair<int32_t, float>>> dataset_embedding;
  std::vector<std::vector<std::pair<int32_t, float>>> raw_dataset;
  EigenMatrix means;
  EigenVector query_mean;
  std::vector<std::pair<float, int32_t>> distances;
  int32_t stage;

  void build_quadtree(const std::vector<int32_t> &subset,
                      const std::vector<std::pair<float, float>> &bounding_box,
                      int32_t depth, int32_t parent) {
    int32_t node_id(parents.size());
    parents.push_back(parent);
    if (subset.size() == 1) {
      leaf[subset[0]] = node_id;
      return;
    }
    int32_t d = dictionary->cols();
    std::vector<float> mid(d);
    for (int32_t i = 0; i < d; ++i) {
      mid[i] = (bounding_box[i].first + bounding_box[i].second) / 2.0;
    }
    std::map<std::vector<uint8_t>, std::vector<int32_t>> parts;
    for (auto ind : subset) {
      std::vector<uint8_t> code((d + 7) / 8, 0);
      for (int32_t i = 0; i < d; ++i) {
        if ((*dictionary)(ind, i) > mid[i]) {
          code[i / 8] |= 1 << (i % 8);
        }
      }
      parts[code].push_back(ind);
    }
    std::vector<std::pair<float, float>> new_bounding_box(d);
    for (const auto &part : parts) {
      for (int32_t i = 0; i < d; ++i) {
        uint8_t bit = (part.first[i / 8] >> (i % 8)) & 1;
        if (bit) {
          new_bounding_box[i] = std::make_pair(mid[i], bounding_box[i].second);
        } else {
          new_bounding_box[i] = std::make_pair(bounding_box[i].first, mid[i]);
        }
      }
      build_quadtree(part.second, new_bounding_box, depth + 1, node_id);
    }
  }

  std::vector<std::pair<int32_t, float>> compute_embedding(
      const std::vector<std::pair<int32_t, float>> &a) {
    std::vector<std::pair<int32_t, float>> result;
    for (auto x : a) {
      auto id = leaf[x.first];
      int32_t level = 0;
      while (id != -1) {
        ++level;
        id = parents[id];
      }
      id = leaf[x.first];
      while (id != -1) {
        --level;
        result.push_back(std::make_pair(id, x.second / (1 << level)));
        id = parents[id];
      }
    }
    std::sort(result.begin(), result.end());
    std::vector<std::pair<int32_t, float>> ans;
    for (auto x : result) {
      if (ans.empty() || ans.back().first != x.first) {
        ans.push_back(x);
      } else {
        ans.back().second += x.second;
      }
    }
    return ans;
  }

  float flowtree_query(const std::vector<std::pair<int32_t, float>> &a,
                       const std::vector<std::pair<int32_t, float>> &b) {
    int32_t num_nodes = 0;
    id_node.clear();
    for (auto x : a) {
      auto id = leaf[x.first];
      while (id != -1) {
        if (marked[id] != num_queries) {
          id_node.push_back(id);
          node_id[id] = num_nodes++;
        }
        marked[id] = num_queries;
        id = parents[id];
      }
    }
    for (auto x : b) {
      auto id = leaf[x.first];
      while (id != -1) {
        if (marked[id] != num_queries) {
          id_node.push_back(id);
          node_id[id] = num_nodes++;
        }
        marked[id] = num_queries;
        id = parents[id];
      }
    }
    if (static_cast<int32_t>(subtree.size()) < num_nodes) {
      subtree.resize(num_nodes);
    }
    for (int32_t i = 0; i < num_nodes; ++i) {
      subtree[i].clear();
    }
    for (int32_t i = 0; i < num_nodes; ++i) {
      int32_t u = parents[id_node[i]];
      if (u != -1) {
        subtree[node_id[u]].push_back(i);
      }
    }
    if (static_cast<int32_t>(excess.size()) < num_nodes) {
      excess.resize(num_nodes);
    }
    delta_node.assign(num_nodes, 0.0);
    unleaf.resize(num_nodes);
    for (auto x : a) {
      delta_node[node_id[leaf[x.first]]] += x.second;
      unleaf[node_id[leaf[x.first]]] = x.first;
    }
    for (auto x : b) {
      delta_node[node_id[leaf[x.first]]] -= x.second;
      unleaf[node_id[leaf[x.first]]] = x.first;
    }
    float res = run_query(0, node_id[0]);
    if (!excess[node_id[0]].empty()) {
      float unassigned = 0.0;
      for (auto x : excess[node_id[0]]) {
        unassigned += x.first;
      }
      if (unassigned > EPS2) {
        throw std::logic_error("too much unassigned flow");
      }
    }
    ++num_queries;
    return res;
  }

  float run_query(int32_t depth, int32_t nd) {
    float res = 0.0;
    for (auto x : subtree[nd]) {
      res += run_query(depth + 1, x);
    }
    excess[nd].clear();
    if (subtree[nd].empty()) {
      if (fabs(delta_node[nd]) > EPS) {
        excess[nd].push_back(std::make_pair(delta_node[nd], unleaf[nd]));
      }
    } else {
      for (auto x : subtree[nd]) {
        if (excess[x].empty()) {
          continue;
        }
        bool same = false;
        if (excess[nd].empty()) {
          same = true;
        } else if (sign(excess[x][0].first) == sign(excess[nd][0].first)) {
          same = true;
        }
        if (same) {
          for (auto y : excess[x]) {
            excess[nd].push_back(y);
          }
        } else {
          while (!excess[x].empty() && !excess[nd].empty()) {
            auto u = excess[nd].back();
            auto v = excess[x].back();

            float dist =
                (dictionary->row(u.second) - dictionary->row(v.second)).norm();
            if (fabs(u.first + v.first) < EPS) {
              excess[nd].pop_back();
              excess[x].pop_back();
              res += dist * fabs(u.first);
            } else if (fabs(u.first) < fabs(v.first)) {
              excess[nd].pop_back();
              excess[x].back().first += u.first;
              res += dist * fabs(u.first);
            } else {
              excess[x].pop_back();
              excess[nd].back().first += v.first;
              res += dist * fabs(v.first);
            }
          }
          if (!excess[x].empty()) {
            excess[x].swap(excess[nd]);
          }
        }
      }
    }
    return res;
  }

  void select_topk_aux(int32_t k1, NumPyIntArray output_ids,
                       NumPyFloatArray output_scores, bool to_sort) {
    auto output_ids_buf = output_ids.request();
    auto output_ids_raw = static_cast<int32_t *>(output_ids_buf.ptr);
    auto output_scores_buf = output_scores.request();
    auto output_scores_raw = static_cast<float *>(output_scores_buf.ptr);
    int32_t k2 = output_ids_buf.shape[0];
    std::nth_element(distances.begin(), distances.begin() + k2 - 1,
                     distances.begin() + k1);
    if (to_sort) {
      std::sort(distances.begin(), distances.begin() + k2);
    }
    for (int32_t i = 0; i < k2; ++i) {
      output_scores_raw[i] = distances[i].first;
      output_ids_raw[i] = distances[i].second;
    }
  }

  void check_measure(const std::vector<std::pair<int32_t, float>> &measure) {
    float sum = 0.0;
    auto n = dictionary->rows();
    for (auto &atom : measure) {
      if (atom.first < 0 || atom.first >= n) {
        throw std::logic_error("invalid index in the measure");
      }
      if (atom.second < -EPS) {
        throw std::logic_error("negative mass");
      }
      sum += atom.second;
    }
    if (fabs(sum - 1.0) > EPS3) {
      throw std::logic_error("the masses don't sum to 1");
    }
  }

  void check_stage() {
    if (stage != 2) {
      throw std::logic_error(
          "need to call load_vocabulary() and load_dataset() first");
    }
  }

  template <typename T>
  void check_dimension(T x) {
    auto buf = x.request();
    if (buf.ndim != 1) {
      throw std::logic_error(
          "input_ids, output_ids, output_scores must be one-dimensional");
    }
  }

  template <typename T>
  ssize_t get_length(T x) {
    return x.request().shape[0];
  }

  void check_input_output_arrays(NumPyIntArray input_ids,
                                 NumPyIntArray output_ids,
                                 NumPyFloatArray output_scores) {
    check_dimension(input_ids);
    check_dimension(output_ids);
    check_dimension(output_scores);
    auto l1 = get_length(input_ids);
    auto l2 = get_length(output_ids);
    auto l3 = get_length(output_scores);
    if (l2 != l3) {
      throw std::logic_error(
          "output_ids and output_scores must be of the same length");
    }
    if (l2 > l1) {
      throw std::logic_error(
          "output_ids and output_scores must be no longer than input_ids");
    }
    auto buf = static_cast<int32_t *>(input_ids.request().ptr);
    for (ssize_t i = 0; i < l1; ++i) {
      auto val = buf[i];
      if (val < 0 || val >= static_cast<int32_t>(raw_dataset.size())) {
        throw std::logic_error("input_ids contain an invalid index");
      }
    }
  }

  void check_input_output_arrays(NumPyIntArray input_ids,
                                 NumPyFloatArray input_scores,
                                 NumPyIntArray output_ids,
                                 NumPyFloatArray output_scores) {
    check_input_output_arrays(input_ids, output_ids, output_scores);
    check_dimension(input_scores);
    if (get_length(input_ids) != get_length(input_scores)) {
      throw std::logic_error(
          "input_ids and input_scores must be of the same length");
    }
  }
};
}  // namespace ote

PYBIND11_MODULE(ot_estimators, m) {
  using ote::OTEstimators;
  py::class_<OTEstimators>(m, "OTEstimators")
      .def(py::init<>())
      .def("load_vocabulary", &OTEstimators::load_vocabulary)
      .def("load_dataset", &OTEstimators::load_dataset)
      .def("means_rank", &OTEstimators::means_rank)
      .def("overlap_rank", &OTEstimators::overlap_rank)
      .def("quadtree_rank", &OTEstimators::quadtree_rank)
      .def("flowtree_rank", &OTEstimators::flowtree_rank)
      .def("select_topk", &OTEstimators::select_topk);
}


#include <boost/graph/chrobak_payne_drawing.hpp>


#pragma once

#include <ros/ros.h>
#include <pluginlib/class_list_macros.h>
#include <gazebo/gazebo.hh>
#include <gazebo/common/Plugin.hh>
#include <gazebo/common/common.hh>
#include <gazebo/physics/physics.hh>
#include <gazebo/util/system.hh>
#include <ignition/math/Vector3.hh>
#include <ignition/math/Pose3.hh>
//#include "IActuator.hpp"
#include "ISee.hpp"
#include "IViaPoints.hpp"
#include "SphericalWrapping.hpp"
#include "CylindricalWrapping.hpp"
#include "MeshWrapping.hpp"
#include <boost/numeric/odeint.hpp>
#include <boost/bind.hpp>
#include <math.h>
#include <map>
#include <vector>
#include "cardsflow_gazebo/muscle/MuscPID.hpp"
#include "Actuator.h"
#include "std_srvs/Trigger.h"

#include <Eigen/Core>
#include <Eigen/Dense>
#include <mutex>
#include <common_utilities/CommonDefinitions.h>

#ifdef ENABLE_LOGGING
#include <nlohmann/json.hpp>
#endif

#include <ignition/transport.hh>
#include <ignition/math.hh>
#include <ignition/msgs.hh>

//#define ENABLE_LOGGING
#define springEncoderTicksPerMeter(meter) (meter*10000.0)

namespace cardsflow_gazebo {

	using namespace std;
	using namespace boost::numeric::odeint;
	using namespace gazebo;
	using namespace Eigen;

	struct MuscInfo{
		string name;
		vector<ViaPointInfo> viaPoints;
		Actuator actuator;
//		Motor motor;
//		Gear gear;
//		Spindle spindle;
		SEE see;
	};

	class IMuscle{

	public:
		IMuscle(physics::ModelPtr parent_model);

        ////////////////////
		/// \brief The Init function.
		///
		/// This function initializes the plugin.
		/// \param[in] muscInfo contains info about via points, motor, gear, spindle and see of the muscles
		void Init(MuscInfo &muscInfo);
		void Update(ros::Time &time, ros::Duration &period );
        double getMuscleLength();
        double getInitialTendonLength();
        double getTendonLength();
        double getMuscleForce();
        double getTendonVelocity(); // in m/s
        bool saveDataService(std_srvs::Trigger::Request &req,std_srvs::Trigger::Response &res);
		string name;
		vector<IViaPointsPtr> viaPoints;
		double cmd = 0;
		bool pid_control = true;
        bool dummy = true;
        boost::shared_ptr<MuscPID> PID;
        ros::ServiceServer save_srv;

        struct{
            double position = 0;
            double velocity = 0;
            double displacement = 0;
        }feedback;

        vector<string> link_names;

		ignition::math::Vector3d momentArm;
        physics::ModelPtr parent_model;
        physics::LinkPtr parent_link;
        boost::shared_ptr<map<string, Matrix4d>> world_to_link_transform;
	private:
        ros::NodeHandlePtr nh;
        ignition::transport::Node node;
        ignition::msgs::Marker markerMsg;
    public:
        int muscleID;
        string muscleName;
        ofstream *file;
		ISee see;
//		IActuator actuator;
		Actuator motor;

    private:
//        IActuator::state_type x;
		//actuatorForce ist the force generated by the motor
		double actuatorForce = 0;
    public:
		//Motorforce is the force getting applied onto the first Viapoint
		double muscleForce = 0;
    private:
		//muscleLength describes the TendonLength from the first Viapoint to the last viapoint. (TendonLength outside the myoMotor)
        double muscleLength = 0;
		// prevMuscleLength is needded to calculate the actual angVel of the motor
		double prevMuscleLength;
		//tendonLength describes the total TendonLength from then Motor until the last viapoint. (TendonLength excluding the tendon coiled up on the motor)
        double tendonLength;
		//initial TendonLength describes the initial total tendonlength
        double initialTendonLength;
		//actual angVel
		double sim_angVel;
        bool firstUpdate;
		double sinParm = 0;
		std::vector<double> springConsts = {0.0, 0.03247} ;

		void initViaPoints( MuscInfo &myoMuscle );
		void calculateTendonForceProgression();
		std::map<uint64_t, std::map<string,double>> log;
        std::map<string,double> status;

		mutex mux;
	};
}


/*
 * Copyright 2012 CarCV Development Team
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/objdetect/objdetect.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/ml/ml.hpp>

#include <boost/lexical_cast.hpp>
#include <boost/filesystem.hpp>

#include <iostream>
#include <fstream>
#include <cstdlib>

#include "ocr_train.hpp"
#include "ocr_recognizer.hpp"

using namespace cv;
using namespace std;

namespace fs = boost::filesystem;

int main() {

	//1. training

	string samples_str = "samples";
	string responses_str = "responses";
	vector<float> responses = loadArray< vector<float> >(responses_str);
	vector< vector<float> > samples = load2DArray<float>(samples_str);

	cout << "Samples: " << samples.size() << endl;
	cout << "Responses: " << responses.size() << endl;

	//test: convert vector to Mat
	Mat_<float> samples_mat = vector2Mat2D(samples);
	Mat_<float> responses_mat = vector2Mat1D(responses);

	cout << "Samples_mat: " << samples_mat.size() << endl;
	cout << "Responses_mat: " << responses_mat.size() << endl;

	//train
        CvKNearest kn_model(samples_mat, responses_mat);

	//2. testing
	//Loads test image with numbers
	Mat img = imread("test.jpg");
	Mat img_original;
	img.copyTo(img_original);

	//Grayscale
	Mat gray;
	cvtColor(img, gray, CV_BGR2GRAY);

	//Builds a threshold making a binary image
	Mat thresh;
	adaptiveThreshold(gray, thresh, 255, 1, 1, 11, 2);

	//Output
	Mat out(img.size(), CV_32F, Scalar(0, 1));

	//Find contours
	vector<vector<Point> > contours;
	vector<Vec4i> hierarchy;
	findContours(thresh, contours, hierarchy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);

	double cntArea;
	Mat contourMat;
	for (vector< vector<Point> >::iterator it = contours.begin();
		 it != contours.end(); ++it ) {

		vector<Point> cnt = *it;
		cntArea = contourArea(cnt);

		if (cntArea<50) {
			continue;
		}

		Rect boundbox = boundingRect(cnt);
		int x = boundbox.x;
		int y = boundbox.y;
		int width = boundbox.width;
		int height = boundbox.height;

		if (height<25) { //If it isn't high enough
			continue;
		}

		Point px(x,y);
		Point py(x+width, y+height);
		Scalar color(0, 0, 255); //red

		img.copyTo(contourMat);
		rectangle(contourMat, px, py, color, 2); //draws a rectangle on top of img

		Mat roi;
		roi = thresh(boundbox); //todo could be reverse (x,y) (y,x), am not sure

		//resize it to 10x10
		Mat roismall;
		Size_<float> smallSize(100,1); //todo: make sure no data is getting lost
		resize(roi, roismall, smallSize);

		//convert to float:
		roismall.convertTo(roismall, CV_32F);

		//todo ""  roismall = roismall.reshape((1,100))
		//roismall = np.float32(roismall)

		Mat_<float> results, neighborResponses, dist;

		//cout << "Roismall_mat: " << roismall.size() << endl;

		float retval = kn_model.find_nearest(roismall, 1, results, neighborResponses, dist);
		float number_res = results.at<float>(0,0);

		string str_result = boost::lexical_cast<string>(number_res);

		Point start_point(x,y+height);
		putText(out, str_result, start_point, 0, 1, color);
	}

	namedWindow("img");
	namedWindow("out");
	imshow("img", img);
	imshow("out", out);

	waitKey(0);

	destroyAllWindows();
	return 0;

}