#include "timing.h"
#include "pyp-topics.hh"
#include "contexts_corpus.hh"

//Dict const *dict;

//#include <boost/date_time/posix_time/posix_time_types.hpp>
void PYPTopics::sample_corpus(const Corpus& corpus, int samples,
                              int freq_cutoff_start, int freq_cutoff_end,
                              int freq_cutoff_interval,
                              int max_contexts_per_document,
                              F temp_start, F temp_end) {
  Timer timer;
  //dict = &((ContextsCorpus*) &corpus)->dict();

  if (!m_backoff.get()) {
    m_word_pyps.clear();
    m_word_pyps.push_back(PYPs());
  }

  std::cerr << "\n Training with " << m_word_pyps.size()-1 << " backoff level"
    << (m_word_pyps.size()==2 ? ":" : "s:") << std::endl;

  for (int i=0; i<(int)m_word_pyps.size(); ++i)
  {
    m_word_pyps.at(i).reserve(m_num_topics);
    for (int j=0; j<m_num_topics; ++j)
      m_word_pyps.at(i).push_back(new PYP<int>(0.01, 1.0, m_seed));
  }
  std::cerr << std::endl;

  m_document_pyps.reserve(corpus.num_documents());
  for (int j=0; j<corpus.num_documents(); ++j)
    m_document_pyps.push_back(new PYP<int>(0.01, 1.0, m_seed));

  m_topic_p0 = 1.0/m_num_topics;
  m_term_p0 = 1.0/(F)m_backoff->terms_at_level(m_word_pyps.size()-1);
  //m_term_p0 = 1.0/corpus.num_types();
  m_backoff_p0 = 1.0/corpus.num_documents();

  std::cerr << " Documents: " << corpus.num_documents() << " Terms: "
    << corpus.num_types() << std::endl;

  int frequency_cutoff = freq_cutoff_start;
  std::cerr << " Context frequency cutoff set to " << frequency_cutoff << std::endl;

  timer.Reset();
  // Initialisation pass
  int document_id=0, topic_counter=0;
  for (Corpus::const_iterator corpusIt=corpus.begin();
       corpusIt != corpus.end(); ++corpusIt, ++document_id) {
    m_corpus_topics.push_back(DocumentTopics(corpusIt->size(), 0));

    int term_index=0;
    for (Document::const_iterator docIt=corpusIt->begin();
         docIt != corpusIt->end(); ++docIt, ++term_index) {
      topic_counter++;
      Term term = *docIt;

      // sample a new_topic
      //int new_topic = (topic_counter % m_num_topics);
      int freq = corpus.context_count(term);
      int new_topic = -1;
      if (freq > frequency_cutoff
          && (!max_contexts_per_document || term_index < max_contexts_per_document)) {
        //new_topic = sample(document_id, term);
        //new_topic = document_id % m_num_topics;
        new_topic = (int) (rnd() * m_num_topics);

        // add the new topic to the PYPs
        increment(term, new_topic);

        if (m_use_topic_pyp) {
          F p0 = m_topic_pyp.prob(new_topic, m_topic_p0);
          int table_delta = m_document_pyps[document_id].increment(new_topic, p0);
          if (table_delta)
            m_topic_pyp.increment(new_topic, m_topic_p0);
        }
        else m_document_pyps[document_id].increment(new_topic, m_topic_p0);
      }

      m_corpus_topics[document_id][term_index] = new_topic;
    }
  }
  std::cerr << "  Initialized in " << timer.Elapsed() << " seconds\n";

  int* randomDocIndices = new int[corpus.num_documents()];
  for (int i = 0; i < corpus.num_documents(); ++i)
	  randomDocIndices[i] = i;

  if (num_jobs < max_threads)
    num_jobs = max_threads;
  int job_incr = (int) ( (float)m_document_pyps.size() / float(num_jobs) );

  // Sampling phase
  for (int curr_sample=0; curr_sample < samples; ++curr_sample) {
    if (freq_cutoff_interval > 0 && curr_sample != 1
        && curr_sample % freq_cutoff_interval == 1
        && frequency_cutoff > freq_cutoff_end) {
      frequency_cutoff--;
      std::cerr << "\n Context frequency cutoff set to " << frequency_cutoff << std::endl;
    }

    F temp = 1.0 / (temp_start - curr_sample*(temp_start-temp_end)/samples);
    std::cerr << "\n  -- Sample " << curr_sample << " (T=" << temp << ") "; std::cerr.flush();

    // Randomize the corpus indexing array
    int tmp;
    int processed_terms=0;
    /*
    for (int i = corpus.num_documents()-1; i > 0; --i)
    {
        //i+1 since j \in [0,i] but rnd() \in [0,1)
    	int j = (int)(rnd() * (i+1));
      assert(j >= 0 && j <= i);
     	tmp = randomDocIndices[i];
    	randomDocIndices[i] = randomDocIndices[j];
    	randomDocIndices[j] = tmp;
    }
    */

    // for each document in the corpus
    int document_id;
    for (int i=0; i<corpus.num_documents(); ++i) {
    	document_id = randomDocIndices[i];

      // for each term in the document
      int term_index=0;
      Document::const_iterator docEnd = corpus.at(document_id).end();
      for (Document::const_iterator docIt=corpus.at(document_id).begin();
           docIt != docEnd; ++docIt, ++term_index) {
        if (max_contexts_per_document && term_index > max_contexts_per_document)
          break;

        Term term = *docIt;

        int freq = corpus.context_count(term);
        if (freq < frequency_cutoff)
          continue;

        processed_terms++;

        // remove the prevous topic from the PYPs
        int current_topic = m_corpus_topics[document_id][term_index];
        // a negative label mean that term hasn't been sampled yet
        if (current_topic >= 0) {
          decrement(term, current_topic);

          int table_delta = m_document_pyps[document_id].decrement(current_topic);
          if (m_use_topic_pyp && table_delta < 0)
            m_topic_pyp.decrement(current_topic);
        }

        // sample a new_topic
        int new_topic = sample(document_id, term, temp);
        //std::cerr << "TERM: " << dict->Convert(term) << " (" << term << ") " << " Old Topic: "
        //  << current_topic << " New Topic: " << new_topic << "\n" << std::endl;

        // add the new topic to the PYPs
        m_corpus_topics[document_id][term_index] = new_topic;
        increment(term, new_topic);

        if (m_use_topic_pyp) {
          F p0 = m_topic_pyp.prob(new_topic, m_topic_p0);
          int table_delta = m_document_pyps[document_id].increment(new_topic, p0);
          if (table_delta)
            m_topic_pyp.increment(new_topic, m_topic_p0);
        }
        else m_document_pyps[document_id].increment(new_topic, m_topic_p0);
      }
      if (document_id && document_id % 10000 == 0) {
        std::cerr << "."; std::cerr.flush();
      }
    }
    std::cerr << " ||| LLH= " << log_likelihood();

    if (curr_sample != 0 && curr_sample % 10 == 0) {
    //if (true) {
      std::cerr << " ||| time=" << (timer.Elapsed() / 10.0) << " sec/sample" << std::endl;
      timer.Reset();
      std::cerr << "     ... Resampling hyperparameters (";

      // resample the hyperparamters
      F log_p=0.0;
      if (max_threads == 1)
      {
        std::cerr << "1 thread)" << std::endl; std::cerr.flush();
        log_p += hresample_topics();
        log_p += hresample_docs(0, m_document_pyps.size());
      }
      else
      { //parallelize
        std::cerr << max_threads << " threads, " << num_jobs << " jobs)" << std::endl; std::cerr.flush();

        WorkerPool<JobReturnsF, F> pool(max_threads);
        int i=0, sz = m_document_pyps.size();
        //documents...
        while (i <= sz - 2*job_incr)
        {
          JobReturnsF job = boost::bind(&PYPTopics::hresample_docs, this, i, i+job_incr);
          pool.addJob(job);
          i += job_incr;
        }
        //  do all remaining documents
        JobReturnsF job = boost::bind(&PYPTopics::hresample_docs, this, i,sz);
        pool.addJob(job);

        //topics...
        JobReturnsF topics_job = boost::bind(&PYPTopics::hresample_topics, this);
        pool.addJob(topics_job);

        log_p += pool.get_result(); //blocks

      }

      if (m_use_topic_pyp) {
        m_topic_pyp.resample_prior(rnd);
        log_p += m_topic_pyp.log_restaurant_prob();
      }

      std::cerr.precision(10);
      std::cerr << " ||| LLH=" << log_likelihood() << " ||| resampling time=" << timer.Elapsed() << " sec" << std::endl;
      timer.Reset();

      int k=0;
      std::cerr << "Topics distribution: ";
      std::cerr.precision(2);
      for (PYPs::iterator pypIt=m_word_pyps.front().begin();
           pypIt != m_word_pyps.front().end(); ++pypIt, ++k) {
        if (k % 5 == 0) std::cerr << std::endl << '\t';
        std::cerr << "<" << k << ":" << pypIt->num_customers() << ","
          << pypIt->num_types() << "," << m_topic_pyp.prob(k, m_topic_p0) << "> ";
      }
      std::cerr.precision(10);
      std::cerr << std::endl;
    }
  }
  delete [] randomDocIndices;
}

PYPTopics::F PYPTopics::hresample_docs(int start, int end)
{
  int resample_counter=0;
  F log_p = 0.0;
  assert(start >= 0);
  assert(end >= 0);
  assert(start <= end);
  for (int i=start; i < end; ++i)
  {
    m_document_pyps[i].resample_prior(rnd);
    log_p += m_document_pyps[i].log_restaurant_prob();
    if (resample_counter++ % 5000 == 0) {
      std::cerr << "."; std::cerr.flush();
    }
  }
  return log_p;
}

PYPTopics::F PYPTopics::hresample_topics()
{
  F log_p = 0.0;
  for (std::vector<PYPs>::iterator levelIt=m_word_pyps.begin();
      levelIt != m_word_pyps.end(); ++levelIt) {
    for (PYPs::iterator pypIt=levelIt->begin();
        pypIt != levelIt->end(); ++pypIt) {

      pypIt->resample_prior(rnd);
      log_p += pypIt->log_restaurant_prob();
    }
    std::cerr << log_p << std::endl;
  }
  return log_p;
}

PYPTopics::F PYPTopics::log_likelihood() const
{
  F log_p = 0.0;

  // LLH of topic term distribution
  size_t i=0;
  for (std::vector<PYPs>::const_iterator levelIt=m_word_pyps.begin();
      levelIt != m_word_pyps.end(); ++levelIt, ++i) {
    for (PYPs::const_iterator pypIt=levelIt->begin();
        pypIt != levelIt->end(); ++pypIt, ++i) {
      log_p += pypIt->log_restaurant_prob();

      if (i == m_word_pyps.size()-1)
        log_p += (pypIt->num_tables() * -log(m_backoff->terms_at_level(i)));
      else
        log_p += (pypIt->num_tables() * log(m_term_p0));
    }
  }
  std::cerr << " TERM LLH: " << log_p << " "; //std::endl;

  // LLH of document topic distribution
  for (size_t i=0; i < m_document_pyps.size(); ++i) {
    log_p += m_document_pyps[i].log_restaurant_prob();
    if (!m_use_topic_pyp) log_p += (m_document_pyps[i].num_tables() * m_topic_p0);
  }
  if (m_use_topic_pyp) {
    log_p += m_topic_pyp.log_restaurant_prob();
    log_p += (m_topic_pyp.num_tables() * log(m_topic_p0));
  }

  return log_p;
}

void PYPTopics::decrement(const Term& term, int topic, int level) {
  //std::cerr << "PYPTopics::decrement(" << term << "," << topic << "," << level << ")" << std::endl;
  int table_delta = m_word_pyps.at(level).at(topic).decrement(term);
  if (table_delta && m_backoff.get()) {
    Term backoff_term = (*m_backoff)[term];
    if (!m_backoff->is_null(backoff_term))
      decrement(backoff_term, topic, level+1);
  }
}

void PYPTopics::increment(const Term& term, int topic, int level) {
  //std::cerr << "PYPTopics::increment(" << term << "," << topic << "," << level << ")" << std::endl;
  int table_delta = m_word_pyps.at(level).at(topic).increment(term, word_pyps_p0(term, topic, level));

  if (table_delta && m_backoff.get()) {
    Term backoff_term = (*m_backoff)[term];
    if (!m_backoff->is_null(backoff_term))
      increment(backoff_term, topic, level+1);
  }
}

int PYPTopics::sample(const DocumentId& doc, const Term& term, F inv_temp) {
  // First pass: collect probs
  F sum=0.0;
  std::vector<F> sums;
  for (int k=0; k<m_num_topics; ++k) {
    F p_w_k = prob(term, k);

    F topic_prob = m_topic_p0;
    if (m_use_topic_pyp) topic_prob = m_topic_pyp.prob(k, m_topic_p0);

    //F p_k_d = m_document_pyps[doc].prob(k, topic_prob);
    F p_k_d = m_document_pyps[doc].unnormalised_prob(k, topic_prob);

    F prob = p_w_k*p_k_d;
    /*
    if (prob < 0.0) { std::cerr << "\n\n" << prob << " " << p_w_k << " " << p_k_d << std::endl; assert(false); }
    if (prob > 1.0) { std::cerr << "\n\n" << prob << " " << p_w_k << " " << p_k_d << std::endl; assert(false); }
    assert (pow(prob, inv_temp) >= 0.0);
    assert (pow(prob, inv_temp) <= 1.0);
    */
    sum += pow(prob, inv_temp);
    sums.push_back(sum);
  }
  // Second pass: sample a topic
  F cutoff = rnd() * sum;
  for (int k=0; k<m_num_topics; ++k) {
    if (cutoff <= sums[k])
      return k;
  }
  assert(false);
}

PYPTopics::F PYPTopics::word_pyps_p0(const Term& term, int topic, int level) const {
  //for (int i=0; i<level+1; ++i) std::cerr << "  ";
  //std::cerr << "PYPTopics::word_pyps_p0(" << term << "," << topic << "," << level << ")" << std::endl;

  F p0 = m_term_p0;
  if (m_backoff.get()) {
    //static F fudge=m_backoff_p0; // TODO

    Term backoff_term = (*m_backoff)[term];
    //std::cerr << "T: " << term << " BO: " << backoff_term << std::endl;
    if (!m_backoff->is_null(backoff_term)) {
      assert (level < m_backoff->order());
      //p0 = (1.0/(F)m_backoff->terms_at_level(level))*prob(backoff_term, topic, level+1);
      p0 = m_term_p0*prob(backoff_term, topic, level+1);
      p0 = prob(backoff_term, topic, level+1);
    }
    else
      p0 = (1.0/(F) m_backoff->terms_at_level(level));
      //p0 = m_term_p0;
  }
  //for (int i=0; i<level+1; ++i) std::cerr << "  ";
  //std::cerr << "PYPTopics::word_pyps_p0(" << term << "," << topic << "," << level << ") = " << p0 << std::endl;
  return p0;
}

PYPTopics::F PYPTopics::prob(const Term& term, int topic, int level) const {
  //for (int i=0; i<level+1; ++i) std::cerr << "  ";
  //std::cerr << "PYPTopics::prob(" << dict->Convert(term) << "," << topic << "," << level << ")" << std::endl;

  F p0 = word_pyps_p0(term, topic, level);
  F p_w_k = m_word_pyps.at(level).at(topic).prob(term, p0);

  /*
  for (int i=0; i<level+1; ++i) std::cerr << "  ";
  std::cerr << "PYPTopics::prob(" << dict->Convert(term) << "," << topic << "," << level << ") = " << p_w_k << std::endl;
  for (int i=0; i<level+1; ++i) std::cerr << "  ";
  m_word_pyps.at(level).at(topic).debug_info(std::cerr);
  */
  return p_w_k;
}

int PYPTopics::max_topic() const {
  if (!m_use_topic_pyp)
    return -1;

  F current_max=0.0;
  int current_topic=-1;
  for (int k=0; k<m_num_topics; ++k) {
    F prob = m_topic_pyp.prob(k, m_topic_p0);
    if (prob > current_max) {
      current_max = prob;
      current_topic = k;
    }
  }
  assert(current_topic >= 0);
  return current_topic;
}

std::pair<int,PYPTopics::F> PYPTopics::max(const DocumentId& doc) const {
  //std::cerr << "PYPTopics::max(" << doc << "," << term << ")" << std::endl;
  // collect probs
  F current_max=0.0;
  int current_topic=-1;
  for (int k=0; k<m_num_topics; ++k) {
    //F p_w_k = prob(term, k);

    F topic_prob = m_topic_p0;
    if (m_use_topic_pyp)
      topic_prob = m_topic_pyp.prob(k, m_topic_p0);

    F prob = 0;
    if (doc < 0) prob = topic_prob;
    else         prob = m_document_pyps[doc].prob(k, topic_prob);

    if (prob > current_max) {
      current_max = prob;
      current_topic = k;
    }
  }
  assert(current_topic >= 0);
  assert(current_max >= 0);
  return std::make_pair(current_topic, current_max);
}

std::pair<int,PYPTopics::F> PYPTopics::max(const DocumentId& doc, const Term& term) const {
  //std::cerr << "PYPTopics::max(" << doc << "," << term << ")" << std::endl;
  // collect probs
  F current_max=0.0;
  int current_topic=-1;
  for (int k=0; k<m_num_topics; ++k) {
    F p_w_k = prob(term, k);

    F topic_prob = m_topic_p0;
    if (m_use_topic_pyp)
      topic_prob = m_topic_pyp.prob(k, m_topic_p0);

    F p_k_d = 0;
    if (doc < 0) p_k_d = topic_prob;
    else         p_k_d = m_document_pyps[doc].prob(k, topic_prob);

    F prob = (p_w_k*p_k_d);
    if (prob > current_max) {
      current_max = prob;
      current_topic = k;
    }
  }
  assert(current_topic >= 0);
  assert(current_max >= 0);
  return std::make_pair(current_topic,current_max);
}

std::ostream& PYPTopics::print_document_topics(std::ostream& out) const {
  for (CorpusTopics::const_iterator corpusIt=m_corpus_topics.begin();
       corpusIt != m_corpus_topics.end(); ++corpusIt) {
    int term_index=0;
    for (DocumentTopics::const_iterator docIt=corpusIt->begin();
         docIt != corpusIt->end(); ++docIt, ++term_index) {
      if (term_index) out << " ";
      out << *docIt;
    }
    out << std::endl;
  }
  return out;
}

std::ostream& PYPTopics::print_topic_terms(std::ostream& out) const {
  for (PYPs::const_iterator pypsIt=m_word_pyps.front().begin();
       pypsIt != m_word_pyps.front().end(); ++pypsIt) {
    int term_index=0;
    for (PYP<int>::const_iterator termIt=pypsIt->begin();
         termIt != pypsIt->end(); ++termIt, ++term_index) {
      if (term_index) out << " ";
      out << termIt->first << ":" << termIt->second;
    }
    out << std::endl;
  }
  return out;
}


//---------------------------------------------------------------------------//
//  MIT License
//
//  Copyright (c) 2020-2021 Mikhail Komarov <nemo@nil.foundation>
//  Copyright (c) 2020-2021 Nikita Kaskov <nemo@nil.foundation>

//
//  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.
//---------------------------------------------------------------------------//

#ifndef FILECOIN_STORAGE_PROOFS_CORE_MERKLE_BUILDERS_HPP
#define FILECOIN_STORAGE_PROOFS_CORE_MERKLE_BUILDERS_HPP

#include <vector>

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

#include <nil/filecoin/storage/proofs/core/sector.hpp>

namespace nil {
    namespace filecoin {
        namespace merkletree {
            template<typename MerkleTreeType>
            inline std::size_t get_base_tree_size(sector_size_type sector_size) {
                std::uint64_t base_tree_leaves =
                    sector_size / MerkleTreeType::hash_type::digest_bits / get_base_tree_count<MerkleTreeType>();

                return get_merkle_tree_len(base_tree_leaves, MerkleTreeType::base_arity);
            }

            template<typename MerkleTreeType>
            inline std::size_t get_base_tree_leafs(std::size_t base_tree_size) {
                return get_merkle_tree_leafs(base_tree_size, MerkleTreeType::base_arity);
            }

            // Create a DiskTree from the provided config(s), each representing a 'base' layer tree with
            // 'base_tree_len' elements.
            template<typename MerkleTreeType>
            DiskTree<typename MerkleTreeType::hash_type, MerkleTreeType::base_arity, MerkleTreeType::sub_tree_arity,
                     MerkleTreeType::top_tree_arity>
                create_disk_tree(std::size_t base_tree_len, const std::vector<StoreConfig> &configs) {
                std::size_t base_tree_leafs = get_merkle_tree_leafs(base_tree_len, MerkleTreeType::base_arity);

                if (MerkleTreeType::top_tree_arity > 0) {
                    BOOST_ASSERT_MSG(MerkleTreeType::sub_tree_arity > 0, "Invalid top arity specified without sub arity");

                    return DiskTree::from_sub_tree_store_configs(base_tree_leafs, configs);
                } else if (MerkleTreeType::sub_tree_arity > 0) {
                    BOOST_ASSERT_MSG(!configs.empty(), "Cannot create sub-tree with a single tree config");

                    return DiskTree::from_store_configs(base_tree_leafs, configs);
                } else {
                    BOOST_ASSERT_MSG(configs.size() == 1, "Invalid tree-shape specified");
                    DiskStore store = DiskStore::new_from_disk(base_tree_len, MerkleTreeType::base_arity, configs[0]);

                    return DiskTree::from_data_store(store, base_tree_leafs);
                }
            }

            // Create an LCTree from the provided config(s) and replica(s), each representing a 'base' layer tree with
            // 'base_tree_len' elements.
            template<typename MerkleTreeType>
            LCTree<typename MerkleTreeType::hash_type, MerkleTreeType::base_arity, MerkleTreeType::sub_tree_arity,
                   MerkleTreeType::top_tree_arity>
                create_lc_tree(std::size_t base_tree_len, const std::vector<StoreConfig> &configs,
                               const ReplicaConfig &replica_config) {
                std::size_t base_tree_leafs = get_merkle_tree_leafs(base_tree_len, MerkleTreeType::base_arity);

                if (MerkleTreeType::top_tree_arity > 0) {
                    BOOST_ASSERT_MSG(MerkleTreeType::sub_tree_arity > 0, "Invalid top arity specified without sub arity");

                    return LCTree::from_sub_tree_store_configs_and_replica(base_tree_leafs, configs, replica_config);
                } else if (MerkleTreeType::sub_tree_arity > 0) {
                    BOOST_ASSERT_MSG(!configs.empty(), "Cannot create sub-tree with a single tree config");

                    return LCTree::from_store_configs_and_replica(base_tree_leafs, configs, replica_config);
                } else {
                    BOOST_ASSERT_MSG(configs.size() == 1, "Invalid tree-shape specified");
                    LCStore store =
                        LCStore::new_from_disk_with_reader(base_tree_len, MerkleTreeType::base_arity, configs[0],
                                                           ExternalReader::new_from_path(replica_config.path), );

                    return LCTree::from_data_store(store, base_tree_leafs);
                }
            }

            // Given base tree configs and optionally a replica_config, returns
            // either a disktree or an lctree, specified by Tree.
            template<typename MerkleTreeType>
            MerkleTreeType create_tree(std::size_t base_tree_len, const std::vector<StoreConfig> &configs,
                                       boost::optional<ReplicaConfig> replica_config) {

                std::size_t base_tree_leafs = get_base_tree_leafs<MerkleTreeType>(base_tree_len);
                std::vector<MerkleTreeType> trees(configs.size());
                for (int i = 0; i < configs.size(); i++) {
                    typename MerkleTreeType::Store store = typename MerkleTreeType::Store::new_with_config(
                        base_tree_len, MerkleTreeType::base_arity, configs[i]);
                    if (const auto Some(lc_store) = Any::downcast_mut::<merkletree::store::LevelCacheStore
                                                             <typename MerkleTreeType::hash_type::digest_type,
                        std::fs::File >> (store)) {
                        BOOST_ASSERT_MSG(replica_config, "Cannot create LCTree without replica paths");
                        lc_store.set_external_reader(ExternalReader::new_from_config(&replica_config, i));
                    }

                    if (configs.size() == 1) {
                        return MerkleTreeType::from_data_store(store, base_tree_leafs);
                    } else {
                        trees.push_back(
                            MerkleTreeType<typename MerkleTreeType::hash_type, MerkleTreeType::Store, MerkleTreeType::arity, 0,
                                           0>::from_data_store(store, base_tree_leafs));
                    }
                }

                BOOST_ASSERT_MSG(MerkleTreeType::top_tree_arity > 0 || MerkleTreeType::sub_tree_arity > 0,
                    "Cannot have a sub/top tree without more than 1 config");
                if (MerkleTreeType::top_tree_arity > 0) {
                    BOOST_ASSERT_MSG(MerkleTreeType::sub_tree_arity > 0, "Invalid top arity specified without sub arity");

                    return MerkleTreeType<
                        typename MerkleTreeType::hash_type, MerkleTreeType::Store, MerkleTreeType::base_arity, MerkleTreeType::sub_tree_arity, MerkleTreeType::top_tree_arity>::
                        from_sub_trees_as_trees(trees);
                } else {
                    BOOST_ASSERT_MSG(!configs.empty(), "Cannot create sub-tree with a single tree config");

                    return MerkleTreeType::from_trees(trees);
                }
            }

            template<typename MerkleTreeType>
            MerkleTreeType create_base_merkle_tree(boost::optional<StoreConfig> config, std::size_t size,
                                                   const std::vector<std::uint8_t> &data) {
                assert(data.size == NODE_SIZE * size);

                BOOST_ASSERT_MSG(is_merkle_tree_size_valid(size, MerkleTreeType::base_arity),
                    "Invalid merkle tree size given the arity");

                const auto f = [&](std::size_t i) {
                    // TODO Replace `expect()` with `context()` (problem is the parallel iterator)
                    std::vector<std::uint8_t> d = data_at_node(data, i);
                    // TODO/FIXME: This can panic. FOR NOW, const auto's leave this since we're experimenting with
                    // optimization paths. However, we need to ensure that bad input will not lead to a panic
                    // that isn't caught by the FPS API.
                    // Unfortunately, it's not clear how to perform this error-handling in the parallel
                    // iterator case.
                    return typename MerkleTreeType::hash_type::digest_type(d);
                };

                if (config) {
                    merkle::MerkleTree <typename MerkleTreeType::hash_type::digest_type,
                        <typename MerkleTreeType::hash_type>::Function, MerkleTreeType::Store, MerkleTreeType::base_arity,
                        MerkleTreeType::sub_tree_arity, MerkleTreeType::top_tree_arity,
                        > ::from_par_iter_with_config((0..size).into_par_iter().map(f), x);
                } else {
                    merkle::MerkleTree:: <typename MerkleTreeType::hash_type::digest_type,
                        <typename MerkleTreeType::hash_type>::Function, MerkleTreeType::Store, MerkleTreeType::base_arity,
                        MerkleTreeType::sub_tree_arity, MerkleTreeType::top_tree_arity,
                        > ::from_par_iter((0..size).into_par_iter().map(f))
                }

                return MerkleTreeType::from_merkle(tree);
            }

            /// Construct a new level cache merkle tree, given the specified
            /// config.
            ///
            /// Note that while we don't need to pass both the data AND the
            /// replica path (since the replica file will contain the same data),
            /// we pass both since we have access from all callers and this avoids
            /// reading that data from the replica_config here.
            template<typename Hash, std::size_t BaseTreeArity>
            LCMerkleTree create_base_lcmerkle_tree(const StoreConfig &config,
                                                   std::size_t size,
                                                   const std::vector<std::uint8_t> &data,
                                                   const ReplicaConfig &replica_config) {
                BOOST_ASSERT_MSG(is_merkle_tree_size_valid(size, BaseTreeArity), "Invalid merkle tree size given the arity");
                BOOST_ASSERT_MSG(data.size() == Hash::digest_bits / CHAR_BIT, "Invalid data length for merkle tree");

                const auto f = [&](std::size_t i) {
                    const auto d = data_at_node(&data, i);
                    return Hash::digest_type(d);
                };

                LCMerkleTree<Hash, BaseTreeArity> lc_tree =
                    LCMerkleTree<Hash, BaseTreeArity>::try_from_iter_with_config((0..size).map(f), config);

                lc_tree.set_external_reader_path(&replica_config.path);

                return lc_tree;
            }

            // Given a StoreConfig, generate additional ones with appended numbers
            // to uniquely identify them and return the results.  If count is 1,
            // the original config is not modified.
            std::vector<StoreConfig> split_config(const StoreConfig &config, std::size_t count) {
                if (count == 1) {
                    return {config};
                }

                std::vector<StoreConfig> configs(count);
                for (int i = 0; i < count; i++) {
                    configs.push_back(StoreConfig::from_config(&config, std::format("{}-{}", config.id, i), None));
                    configs[i].rows_to_discard = config.rows_to_discard;
                }

                return configs;
            }

            // Given a StoreConfig, generate additional ones with appended numbers
            // to uniquely identify them and return the results.  If count is 1,
            // the original config is not modified.
            //
            // Useful for testing, where there the config may be None.
            std::vector<boost::optional<StoreConfig>> split_config_wrapped(boost::optional<StoreConfig> config,
                                                                           std::size_t count) {
                if (count == 1) {
                    return {config};
                }

                if (config) {
                    std::vector<boost::optional<StoreConfig>> configs(count);
                    for (int i = 0; i < count; i++) {
                        configs.push_back(Some(StoreConfig::from_config(&c, std::format("{}-{}", c.id, i), None, )));
                    }
                    return configs;
                } else {
                    return {};
                }
            }

            // Given a StoreConfig, replica path and tree_width (leaf nodes),
            // append numbers to each StoreConfig to uniquely identify them and
            // return the results along with a ReplicaConfig using calculated
            // offsets into the single replica path specified for later use with
            // external readers.  If count is 1, the original config is not
            // modified.
            std::tuple<std::vector<StoreConfig>, ReplicaConfig>
                split_config_and_replica(StoreConfig config, const boost::filesystem::path &replica_path,
                                         std::size_t sub_tree_width,    // nodes, not bytes
                                         std::size_t count) {
                if (count == 1) {
                    return std::make_tuple({config}, {replica_path, {0}});
                }

                std::vector<StoreConfig> configs(count);
                std::vector<std::size_t> replica_offsets(count);

                for (int i = 0; i < count; i++) {
                    configs.push_back(StoreConfig::from_config(&config, std::format("{}-{}", config.id, i), None));
                    configs[i].rows_to_discard = config.rows_to_discard;

                    replica_offsets.push_back(i * sub_tree_width * NODE_SIZE);
                }

                return std::make_tuple(configs, {replica_path, replica_offsets});
            }

            template<typename MerkleTreeType>
            std::size_t get_base_tree_count() {
                if (MerkleTreeType ::top_tree_arity == 0 && MerkleTreeType::sub_tree_arity == 0) {
                    return 1;
                }

                if (MerkleTreeType ::top_tree_arity > 0) {
                    assert(MerkleTreeType::sub_tree_arity != 0);

                    return MerkleTreeType::top_tree_arity * MerkleTreeType::sub_tree_arity;
                } else {
                    return MerkleTreeType::sub_tree_arity;
                }
            }
        }    // namespace merkletree
    }    // namespace filecoin
}    // namespace nil

#endif


#define BOOST_TEST_MODULE "test_read_steepest_descent_simulator"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif
#include <test/util/make_empty_input.hpp>
#include <mjolnir/input/read_simulator.hpp>

BOOST_AUTO_TEST_CASE(read_steepest_descent_simulator)
{
    mjolnir::LoggerManager::set_default_logger("test_read_steepest_descent_simulator.log");

    using real_type = double;
    using traits_type = mjolnir::SimulatorTraits<real_type, mjolnir::UnlimitedBoundary>;
    auto root = mjolnir::test::make_empty_input();

    using namespace toml::literals;
    const auto v = u8R"(
        type            = "SteepestDescent"
        precision       = "double"
        boundary_type   = "Unlimited"
        delta           = 0.1
        step_limit      = 100
        save_step       = 10
        threshold       = 0.0
    )"_toml;
    root.as_table()["simulator"] = v;

    {
        const auto sim = mjolnir::read_integrator_type<traits_type>(root, v);
        BOOST_TEST(static_cast<bool>(sim));

        const auto sdsim = dynamic_cast<
            mjolnir::SteepestDescentSimulator<traits_type>*>(sim.get());
        BOOST_TEST(static_cast<bool>(sdsim));

        sim->initialize();
        for(std::size_t i=0; i<99; ++i)
        {
            BOOST_TEST(sim->step()); // check it can step
        }
        // at the last (100-th) step, it returns false to stop the simulation.
        BOOST_TEST(!sim->step());
        sim->finalize();
    }
    {
        const auto sim = mjolnir::read_simulator<traits_type,
              mjolnir::VelocityVerletIntegrator<traits_type>>(root, v);
        BOOST_TEST(static_cast<bool>(sim));

        const auto sdsim = dynamic_cast<
            mjolnir::SteepestDescentSimulator<traits_type>*>(sim.get());
        BOOST_TEST(static_cast<bool>(sdsim));

        sim->initialize();
        for(std::size_t i=0; i<99; ++i)
        {
            BOOST_TEST(sim->step()); // check it can step
        }
        // at the last (100-th) step, it returns false to stop the simulation.
        BOOST_TEST(!sim->step());
        sim->finalize();
    }

}


/*
 * BowDescriptor.hpp
 *
 *  Created on: Oct 14, 2011
 *      Author: lbossard
 */

#ifndef FEATURES_BOWDESCRIPTOR_HPP_
#define FEATURES_BOWDESCRIPTOR_HPP_

#include <vector>

#include <boost/noncopyable.hpp>

#include <opencv2/core/core.hpp>
#include <opencv2/flann/flann.hpp>

namespace vision
{
namespace features
{

class BowExtractor : boost::noncopyable
{
public:
    typedef int WordId;
    typedef cv::Mat_<int> RowHistograms;

    explicit BowExtractor(const cv::Mat_<float>& vocabulary,
                          cvflann::flann_algorithm_t=cvflann::FLANN_INDEX_LINEAR);
    virtual ~BowExtractor();

    void match(const cv::Mat_<float>& features,
            std::vector<WordId>& words);

    void match(
            const cv::Mat_<float>& features,
            cv::Mat_<WordId>& indices,
            cv::Mat_<float>& dists,
            unsigned int knn);

    void sumPool( const std::vector<WordId>& words,
            const std::vector<cv::Point>& locations,
            const cv::Rect& roi,
            const cv::Mat_<uchar>& mask,
            const unsigned int levels,
            RowHistograms& histograms
            ) const;

    void createHistogram(
                const std::vector<WordId>& words,
                const std::vector<cv::Point>& locations,
                const cv::Rect& roi,
                const cv::Mat_<uchar>& mask,
                RowHistograms& histogram) const;
    // T == cv::Mat_<U>
    template <typename T>
    static void sumPool(
            const std::vector<WordId>& words,
            const std::vector<cv::Point>& locations,
            const cv::Rect& roi,
            const cv::Mat_<uchar>& mask,
            const unsigned int word_count,
            const unsigned int levels,
            cv::Mat_<T>& histograms);
    // T == cv::Mat_<U>
    template <typename T>
    static void createHistogram(
            const std::vector<WordId>& words,
            const std::vector<cv::Point>& locations,
            const cv::Rect& roi,
            const cv::Mat_<uchar>& mask,
            const unsigned int max_word_count,
            cv::Mat_<T>& histogram);

    /**
     * Get number of histograms for a given level. level must be greater than 0
     * @param levels
     * @return
     */
    static inline int getHistogramCount(const unsigned int levels);

    inline WordId maxWordId() const;
    inline std::size_t wordCount() const;

protected:
    cv::Mat vocabulary_;

    cvflann::SearchParams search_params_;
    cv::flann::GenericIndex<cvflann::L2<float> > index_;

private:
    BowExtractor();

};

//------------------------------------------------------------------------------
inline BowExtractor::WordId BowExtractor::maxWordId() const
{
    //TODO: empty vocabulary? -> underflow !111
    return vocabulary_.rows - 1;
}
inline std::size_t BowExtractor::wordCount() const
{
    return vocabulary_.rows;
}
/*static*/ inline int BowExtractor::getHistogramCount(const unsigned int levels)
{
    return (std::pow(2, 2 * levels) - 1 ) / 3; // same as  ((2 << ((2 * levels) - 1)) - 1 ) / 3 which only works for levels > 0;
}

///*static*/ inline BowExtractor::getLevel(const unsigned int index, const unsigned int vocsize)
//{
//    return ceil(log2( 3 * index + 1)/ 2.f);
//}

template <typename T>
/*static*/ void BowExtractor::sumPool(
        const std::vector<WordId>& words,
        const std::vector<cv::Point>& locations,
        const cv::Rect& roi,
        const cv::Mat_<uchar>& mask,
        const unsigned int word_count,
        const unsigned int levels,
        cv::Mat_<T>& histograms)
{
    if (levels < 1)
    {
        throw "levels > 0 !!111";
    }

    // compute number of histograms
    const int histo_count = getHistogramCount(levels);

    // one histo per row
    histograms.create(histo_count, word_count);
    histograms = 0;

    // compute histograms at lowest scale
    {
        const int box_count = std::pow(2, levels - 1);
        const int box_width = roi.width / box_count;
        const int box_height = roi.height / box_count;
        const int histograms_idx_offset = histo_count - (box_count * box_count);
        for (int c = box_count - 1; c >= 0 ; --c)
        {
            for (int r = box_count - 1; r >= 0; --r)
            {
                const cv::Rect rect(roi.x + c * box_width, roi.y +  r * box_height, box_width, box_height);
//        std::cout << rect.x << " "  << rect.y << " "<< rect.width << " "<< rect.height << " "<<  ((image_dims & rect) == rect) << " " <<  c * box_count + r << " "<<  c * box_count + r + histograms_idx_offset << " "<< std::endl;
                const int histogram_idx = c * box_count + r + histograms_idx_offset;
                cv::Mat_<T> histogram = histograms.row(histogram_idx);
                BowExtractor::createHistogram(words, locations, rect, mask, word_count, histogram);
            }
        }
    }

    // derive histo for upper layers
    for (int level = levels - 2 ; level >= 0; --level)
    {
        const int current_histo_idx_offset = getHistogramCount(level);
        const int lower_histo_idx_offset = getHistogramCount(level + 1);
        const int box_count = std::pow(2, level + 1);

        for (int c = 0; c < box_count ; ++c)
        {
            for (int r = 0; r < box_count ; ++r)
            {
                const int current_level_histo_ix = current_histo_idx_offset + c/2 * box_count/2 + r/2;
                const int lower_histo_idx = lower_histo_idx_offset + c * box_count + r;
                histograms.row(current_level_histo_ix) += histograms.row(lower_histo_idx);
            }
        }
    }
}

template <typename T>
/*static*/ void BowExtractor::createHistogram(
        const std::vector<WordId>& words,
        const std::vector<cv::Point>& locations,
        const cv::Rect& roi,
        const cv::Mat_<uchar>& mask,
        const unsigned int max_word_count,
        cv::Mat_<T>& histogram)
{
    histogram.create(1,max_word_count);
    histogram = 0;
    std::size_t item_count = words.size();

    if (!mask.data)
    {
        for (unsigned int i = 0; i < item_count; ++i)
        {
            const cv::Point& p = locations[i];
            if (roi.contains(p))
            {
                histogram(words[i]) += 1;
            }
        }
    }
    else
    {
        for (unsigned int i = 0; i < item_count; ++i)
        {
            const cv::Point& p = locations[i];
            if (roi.contains(p) && mask(p.y, p.x))
            {
                histogram(words[i]) += 1;
            }
        }
    }
}

} /* namespace features */
} /* namespace vision */
#endif /* FEATURES_BOWDESCRIPTOR_HPP_ */


#include <cstdlib>
#include <vector>
#include <cmath>
#include <Eigen/Core>
#include <moihgp/moihgp.h>
#include <moihgp/matern32ss.h>
#include <moihgp/matern52ss.h>

#ifdef _WIN32
#ifdef LIBRARY_EXPORTS
#define LIBRARY_API __declspec(dllexport)
#else
#define LIBRARY_API __declspec(dllimport)
#endif
#else
#define LIBRARY_API
#endif

typedef moihgp::MOIHGP<moihgp::Matern32StateSpace> GP32;
typedef moihgp::MOIHGP<moihgp::Matern32StateSpace> GP52;

extern "C"
{

LIBRARY_API GP32* gp32_new(double dt, size_t num_output, size_t num_latent, bool threading)
{
    return new GP32(dt, num_output, num_latent, threading);
} // GP32* gp32_new(double dt, size_t num_output, size_t num_latent, bool threading)

LIBRARY_API void gp32_del(GP32* gp)
{
    gp->~MOIHGP();
} // void gp32_del(GP32* gp)

LIBRARY_API void gp32_step1(GP32* gp, double* x, double* y, double* dx, double* xnew, double* yhat, double* dxnew)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dxnew(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    gp->step(_x, _y, _dx, _xnew, _yhat, _dxnew);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                dxnew[idx1 * igp_num_param * dim + idx2 * dim + idx3] = _dxnew[idx1][idx2](idx3);
            }
        }
    }

} // void gp32_step1(GP32* gp, double* x, double* y, double* dx, double* xnew, double* yhat, double* dxnew)

LIBRARY_API void gp32_step2(GP32* gp, double* x, double* y, double* dx, double* xnew, double* dxnew)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    std::vector<std::vector<Eigen::VectorXd>> _dxnew(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    gp->step(_x, _y, _dx, _xnew, _dxnew);

    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                dxnew[idx1 * igp_num_param * dim + idx2 * dim + idx3] = _dxnew[idx1][idx2](idx3);
            }
        }
    }

} // void gp32_step2(GP32* gp, double* x, double* y, double* dx, double* xnew, double* dxnew)

LIBRARY_API void gp32_step3(GP32* gp, double* x, double* y, double* xnew, double* yhat)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    gp->step(_x, _y, _xnew, _yhat);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
    }

} // void gp32_step3(GP32* gp, double* x, double* y, double* xnew, double* yhat)

LIBRARY_API void gp32_step4(GP32* gp, double* x, double* xnew, double* yhat)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);

    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    gp->step(_x, _xnew, _yhat);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
    }

} // void gp32_step4(GP32* gp, double* x, double* xnew, double* yhat)

LIBRARY_API void gp32_update(GP32* gp, double* params)
{
    size_t num_param = gp->getNumParam();
    Eigen::VectorXd _params = Eigen::Map<Eigen::VectorXd>(params, num_param, 1);
    gp->update(_params);
} // void gp32_update(GP32* gp, double* params)

LIBRARY_API double gp32_lik1(GP32* gp, double* x, double* y, double* dx, double* grad)
{
    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t num_param = gp->getNumParam();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    Eigen::VectorXd _grad(num_param);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    double loss = gp->negLogLikelihood(_x, _y, _dx, _grad);

    for (size_t idx = 0; idx < num_param; idx++)
    {
        grad[idx] = _grad(idx);
    }

    return loss;

} // double gp32_lik1(GP32* gp, double* x, double* y, double* dx, Eigen::VectorXd &grad)

LIBRARY_API double gp32_lik2(GP32* gp, double* x, double* y)
{
    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t num_param = gp->getNumParam();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    double loss = gp->negLogLikelihood(_x, _y);

    return loss;

} // double gp32_lik2(GP32* gp, double* x, double* y)

LIBRARY_API void gp32_get_params(GP32* gp, double* params)
{
    size_t num_param = gp->getNumParam();
    Eigen::VectorXd _params = gp->getParams();
    for (size_t idx = 0; idx < num_param; idx++)
    {
        params[idx] = _params(idx);
    }
} // void gp32_get_params(GP32* gp, double* params)

LIBRARY_API size_t gp32_igp_dim(GP32* gp)
{
    return gp->getIGPDim();
}

LIBRARY_API size_t gp32_num_param(GP32* gp)
{
    return gp->getNumParam();
}

LIBRARY_API size_t gp32_num_igp_param(GP32* gp)
{
    return gp->getNumIGPParam();
}

LIBRARY_API GP52* gp52_new(double dt, size_t num_output, size_t num_latent, bool threading)
{
    return new GP52(dt, num_output, num_latent, threading);
} // GP52* gp52_new(double dt, size_t num_output, size_t num_latent, bool threading)

LIBRARY_API void gp52_del(GP52* gp)
{
    gp->~MOIHGP();
} // void gp52_del(GP52* gp)

LIBRARY_API void gp52_step1(GP52* gp, double* x, double* y, double* dx, double* xnew, double* yhat, double* dxnew)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dxnew(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    gp->step(_x, _y, _dx, _xnew, _yhat, _dxnew);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                dxnew[idx1 * igp_num_param * dim + idx2 * dim + idx3] = _dxnew[idx1][idx2](idx3);
            }
        }
    }

} // void gp52_step1(GP52* gp, double* x, double* y, double* dx, double* xnew, double* yhat, double* dxnew)

LIBRARY_API void gp52_step2(GP52* gp, double* x, double* y, double* dx, double* xnew, double* dxnew)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    std::vector<std::vector<Eigen::VectorXd>> _dxnew(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3=0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    gp->step(_x, _y, _dx, _xnew, _dxnew);

    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                dxnew[idx1 * igp_num_param * dim + idx2 * dim + idx3] = _dxnew[idx1][idx2](idx3);
            }
        }
    }

} // void gp52_step2(GP52* gp, double* x, double* y, double* dx, double* xnew, double* dxnew)

LIBRARY_API void gp52_step3(GP52* gp, double* x, double* y, double* xnew, double* yhat)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    gp->step(_x, _y, _xnew, _yhat);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
    }

} // void gp52_step3(GP52* gp, double* x, double* y, double* xnew, double* yhat)

LIBRARY_API void gp52_step4(GP52* gp, double* x, double* xnew, double* yhat)
{

    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    std::vector<Eigen::VectorXd> _xnew(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _yhat(num_output);

    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    gp->step(_x, _xnew, _yhat);

    for (size_t idx = 0; idx < num_output; idx++)
    {
        yhat[idx] = _yhat(idx);
    }
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            xnew[idx1 * dim + idx2] = _xnew[idx1](idx2);
        }
    }

} // void gp52_step4(GP52* gp, double* x, double* xnew, double* yhat)

LIBRARY_API void gp52_update(GP52* gp, double* params)
{
    size_t num_param = gp->getNumParam();
    Eigen::VectorXd _params = Eigen::Map<Eigen::VectorXd>(params, num_param, 1);
    gp->update(_params);
} // void gp52_update(GP52* gp, double* params)

LIBRARY_API double gp52_lik1(GP52* gp, double* x, double* y, double* dx, double* grad)
{
    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t num_param = gp->getNumParam();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);
    std::vector<std::vector<Eigen::VectorXd>> _dx(num_latent, std::vector<Eigen::VectorXd>(igp_num_param, Eigen::VectorXd(dim).setZero()));
    Eigen::VectorXd _grad(num_param);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
        for (size_t idx2 = 0; idx2 < igp_num_param; idx2++)
        {
            for (size_t idx3 = 0; idx3 < dim; idx3++)
            {
                _dx[idx1][idx2](idx3) = dx[idx1 * igp_num_param * dim + idx2 * dim + idx3];
            }
        }
    }

    double loss = gp->negLogLikelihood(_x, _y, _dx, _grad);

    for (size_t idx = 0; idx < num_param; idx++)
    {
        grad[idx] = _grad(idx);
    }

    return loss;

} // double gp52_lik1(GP52* gp, double* x, double* y, double* dx, Eigen::VectorXd &grad)

LIBRARY_API double gp52_lik2(GP52* gp, double* x, double* y)
{
    size_t num_output = gp->getNumOutput();
    size_t num_latent = gp->getNumLatent();
    size_t num_param = gp->getNumParam();
    size_t igp_num_param = gp->getNumIGPParam();
    size_t dim = gp->getIGPDim();

    std::vector<Eigen::VectorXd> _x(num_latent, Eigen::VectorXd(dim).setZero());
    Eigen::VectorXd _y(num_output);

    _y = Eigen::Map<Eigen::VectorXd>(y, num_output, 1);
    for (size_t idx1 = 0; idx1 < num_latent; idx1++)
    {
        for (size_t idx2 = 0; idx2 < dim; idx2++)
        {
            _x[idx1](idx2) = x[idx1 * dim + idx2];
        }
    }

    double loss = gp->negLogLikelihood(_x, _y);

    return loss;

} // double gp52_lik2(GP52* gp, double* x, double* y)

LIBRARY_API void gp52_get_params(GP52* gp, double* params)
{
    size_t num_param = gp->getNumParam();
    Eigen::VectorXd _params = gp->getParams();
    for (size_t idx = 0; idx < num_param; idx++)
    {
        params[idx] = _params(idx);
    }
} // void gp52_get_params(GP52* gp, double* params)

LIBRARY_API size_t gp52_igp_dim(GP52* gp)
{
    return gp->getIGPDim();
}

LIBRARY_API size_t gp52_num_param(GP52* gp)
{
    return gp->getNumParam();
}

LIBRARY_API size_t gp52_num_igp_param(GP52* gp)
{
    return gp->getNumIGPParam();
}

} // extern "C"


#include <chrono>
// #include <functional>
// #include <memory>
// #include <string>
#include <visualization_msgs/msg/marker.hpp>
#include <geometry_msgs/msg/point.hpp>
#include <geometry_msgs/msg/vector3.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>
#include <tf2/convert.h>
#include <tf2/exceptions.h>
#include <tf2_eigen/tf2_eigen.h>
#include <tf2_ros/transform_broadcaster.h>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <tf2_sensor_msgs/tf2_sensor_msgs.h>
#include <vector>
#include <Eigen/Core>

#include <pcl/filters/approximate_voxel_grid.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/registration/ndt.h>
#include <pcl/common/transforms.h>

#include "rclcpp/rclcpp.hpp"
#include "std_msgs/msg/string.hpp"

using namespace std::chrono_literals;
using namespace sensor_msgs::msg;
using namespace nav_msgs::msg;
using geometry_msgs::msg::TransformStamped;

/*
PSEUDOCODE

Subs for:
- /gps/odom
- /map/pcd
- /lidar/filtered
- /tf (special listener)

On PCD map, cache to pcl::PointCloud<pcl::PointXYZ>

Nudge if:
1. New GPS odom received, then
2. Grab next Lidar
3. Transform Lidar to /map and ICP with cached map
3. Set gps_received = false
4. Repeat

*/

class ICPNudger : public rclcpp::Node
{
	public:
		ICPNudger()
		: Node("icp_nudger")
		{
			this->declare_parameter<double>("nudge_period", 1.0); // Time between ICP updates, in seconds
			this->declare_parameter<bool>("publish_tf", false); // If true, ICPNudger will publish result as tf.
																// If enabled, no other node should publish a map tf.
			this->get_parameter<double>("nudge_period", nudge_period_);

			lidar_sub_ = this->create_subscription<PointCloud2>("/lidar", 1, //Only one message, too slow for queue!
				[this](PointCloud2::SharedPtr msg) { lidarCb(msg); }
			);
			map_sub_ = this->create_subscription<PointCloud2>("/map", 1,
				[this](PointCloud2::SharedPtr msg) { pcd_map_ = msg; } // Convert and cache
			);
			gps_sub_ = this->create_subscription<Odometry>("/gps", 1,
				[this](Odometry::SharedPtr msg) { cached_gps_ = msg; } // Convert and cache
			);

			aligned_lidar_pub_ = this->create_publisher<PointCloud2>("/lidar/nudged", 10);

			timer_ = this->create_wall_timer(std::chrono::duration<double>(nudge_period_), std::bind(&ICPNudger::tryNDT, this));

			// TF2
			tf_buffer_ = std::make_unique<tf2_ros::Buffer>(this->get_clock());
			transform_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
		}

	private:
		void lidarCb(PointCloud2::SharedPtr msg) {
			cached_lidar_ = msg;
		}

		void tryNDT() {
			RCLCPP_INFO(this->get_logger(), "Trying NDT...");

			// TF lidar to /map using GPS odom, convert to PCL PCD
			auto transformed_lidar = new PointCloud2();
			geometry_msgs::msg::TransformStamped gps_to_map;

			if(cached_lidar_ == nullptr || cached_gps_ == nullptr) {
				RCLCPP_WARN(this->get_logger(), "No Lidar or GPS data received. Skipping.");
				return;
			}

			auto gps_pose = cached_gps_->pose.pose;
			gps_to_map.transform.translation.x = gps_pose.position.x;
			gps_to_map.transform.translation.y = gps_pose.position.y;
			gps_to_map.transform.translation.z = gps_pose.position.z;
			gps_to_map.transform.rotation = gps_pose.orientation;

			gps_to_map.child_frame_id = "map";
			gps_to_map.header.frame_id = "base_link";
			gps_to_map.header.stamp = this->get_clock()->now();

			auto lidar_to_bl_tf = tf_buffer_->lookupTransform("base_link", cached_lidar_->header.frame_id.c_str(), this->get_clock()->now());

			tf2::doTransform(*cached_lidar_, *cached_lidar_, lidar_to_bl_tf);
			tf2::doTransform(*cached_lidar_, *transformed_lidar, gps_to_map);
			// RCLCPP_INFO(this->get_logger(), "Transform performed");

			pcl::PointCloud<pcl::PointXYZ>::Ptr pcl_lidar(new pcl::PointCloud<pcl::PointXYZ>(5,1)); // Input to ICP, in PCL's pcd format
			pcl::fromROSMsg(*transformed_lidar, *pcl_lidar);
			pcl::PointCloud<pcl::PointXYZ>::Ptr pcl_map(new pcl::PointCloud<pcl::PointXYZ>(5,1));; // Input to ICP, in PCL's pcd format
			pcl::fromROSMsg(*pcd_map_, *pcl_map);

			pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud (new pcl::PointCloud<pcl::PointXYZ>);
			pcl::ApproximateVoxelGrid<pcl::PointXYZ> approximate_voxel_filter;
			approximate_voxel_filter.setLeafSize (0.5, 0.5, 0.5);
			approximate_voxel_filter.setInputCloud (pcl_lidar);
			approximate_voxel_filter.filter (*filtered_cloud);

			// Apply ICP https://pointclouds.org/documentation/tutorials/iterative_closest_point.html
			pcl::NormalDistributionsTransform<pcl::PointXYZ, pcl::PointXYZ> ndt;
			ndt.setTransformationEpsilon(0.01);
			ndt.setStepSize(0.1);
			ndt.setResolution(1.0);
			ndt.setMaximumIterations(35);
			ndt.setInputSource(filtered_cloud);
			ndt.setInputTarget(pcl_map);

			pcl::PointCloud<pcl::PointXYZ> Final;
			ndt.align(Final);

			std::cout << "has converged:" << ndt.hasConverged() << " score: " <<
			ndt.getFitnessScore() << std::endl;
			std::cout << ndt.getMaximumIterations() << std::endl;

			Eigen::Affine3d eigen_affine_transform;
			eigen_affine_transform.matrix() = ndt.getFinalTransformation().cast<double>();
			auto ndt_tf_stamped = tf2::eigenToTransform(eigen_affine_transform);
			auto nudged_lidar = new PointCloud2();
			tf2::doTransform(*transformed_lidar, *nudged_lidar, ndt_tf_stamped);
			tf2::doTransform(*nudged_lidar, *nudged_lidar, gps_to_map);

			aligned_lidar_pub_->publish(*nudged_lidar);

			RCLCPP_INFO_STREAM(this->get_logger(), "GPS TF: "<<
													gps_to_map.transform.translation.x<<", "
													<<gps_to_map.transform.translation.y<<", "
													<<gps_to_map.transform.translation.z<<" - "
													<<gps_to_map.transform.rotation.w<<", "
													<<gps_to_map.transform.rotation.x<<", "
													<<gps_to_map.transform.rotation.y<<", "
													<<gps_to_map.transform.rotation.z);

			RCLCPP_INFO_STREAM(this->get_logger(), "ndt TF: "<<
													ndt_tf_stamped.transform.translation.x<<", "
													<<ndt_tf_stamped.transform.translation.y<<", "
													<<ndt_tf_stamped.transform.translation.z<<" - "
													<<ndt_tf_stamped.transform.rotation.w<<", "
													<<ndt_tf_stamped.transform.rotation.x<<", "
													<<ndt_tf_stamped.transform.rotation.y<<", "
													<<ndt_tf_stamped.transform.rotation.z);

			RCLCPP_INFO_STREAM(this->get_logger(), "Final TF: "<<
													ndt_tf_stamped.transform.translation.x+gps_to_map.transform.translation.x<<", "
													<<ndt_tf_stamped.transform.translation.y+gps_to_map.transform.translation.y<<", "
													<<ndt_tf_stamped.transform.translation.z+gps_to_map.transform.translation.z<<" - "
													<<ndt_tf_stamped.transform.rotation.w<<", "
													<<ndt_tf_stamped.transform.rotation.x<<", "
													<<ndt_tf_stamped.transform.rotation.y<<", "
													<<ndt_tf_stamped.transform.rotation.z);

			bool publish_tf;
			this->get_parameter<bool>("publish_tf", publish_tf);
			if(publish_tf) {
				std::unique_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster = std::make_unique<tf2_ros::TransformBroadcaster>(*this);
				TransformStamped t;

				t.header.frame_id = "map";
				t.header.stamp = this->get_clock()->now();

				t.child_frame_id = "base_link";
				t.transform.translation.x = ndt_tf_stamped.transform.translation.x+gps_to_map.transform.translation.x;
				t.transform.translation.y = ndt_tf_stamped.transform.translation.y+gps_to_map.transform.translation.y;
				t.transform.translation.z = ndt_tf_stamped.transform.translation.z+gps_to_map.transform.translation.z;
				t.transform.rotation = cached_gps_->pose.pose.orientation;
				tf_broadcaster->sendTransform(t);
			}

			// Generate PoseWithCovariance using fitness score, bias, and GPS odom

			// Publish nudged odom as Odometry msg
		}

		std::shared_ptr<tf2_ros::TransformListener> transform_listener_{nullptr};
		std::unique_ptr<tf2_ros::Buffer> tf_buffer_;
		rclcpp::Subscription<PointCloud2>::SharedPtr map_sub_;
		rclcpp::Subscription<PointCloud2>::SharedPtr lidar_sub_;
		rclcpp::Subscription<Odometry>::SharedPtr gps_sub_;
		rclcpp::Publisher<PointCloud2>::SharedPtr aligned_lidar_pub_;
		sensor_msgs::msg::PointCloud2::SharedPtr pcd_map_;
		sensor_msgs::msg::PointCloud2::SharedPtr cached_lidar_;
		Odometry::SharedPtr cached_gps_;
		double nudge_period_;
		rclcpp::TimerBase::SharedPtr timer_{nullptr};
};

int main(int argc, char * argv[])
{
	rclcpp::init(argc, argv);
	rclcpp::spin(std::make_shared<ICPNudger>());
	rclcpp::shutdown();
	return 0;
}


// TnzCore includes
#include "tmeshimage.h"

// tcg includes
#include "tcg/tcg_misc.h"
#include "tcg/tcg_mesh_bgl.h"

using namespace tcg::bgl;

// Boost includes
#include <boost/graph/properties.hpp>
#include <boost/graph/breadth_first_search.hpp>

// STD includes
#include <queue>

#include "ext/plastichandle.h"

#include <cstring>

//***********************************************************************************************
//    Distance building
//***********************************************************************************************

namespace {

typedef tcg::Mesh<TTextureVertex, tcg::Edge, tcg::FaceN<3>> Graph;

//----------------------------------------------------------------------------------

struct DistanceGreater {
  float *m_distances;

public:
  DistanceGreater(float *distances) : m_distances(distances) {}

  bool operator()(int a, int b) { return m_distances[a] > m_distances[b]; }
};

}  // namespace

//==================================================================================

namespace local {

struct BFS_DistanceBuilder {
  typedef boost::graph_traits<Graph>::edge_descriptor edge_descr;

  float *m_distances;  //!< (NOT owned) Distances from selected vertex
  UCHAR *m_colormap;   //!< (NOT owned) Graph BFS colormap

public:
  BFS_DistanceBuilder(float *distances, UCHAR *colormap)
      : m_distances(distances), m_colormap(colormap) {}

  void tree_edge(const edge_descr &e, const Graph &g) {
    int v, v0 = -1, v1 = tcg::bgl::target(e, g);
    double d, dMin = (std::numeric_limits<double>::max)();

    assert(m_colormap[v1] == boost::white_color);

    // Search for the distance-nearest found vertex
    const Graph::vertex_type &vx1 = g.vertex(v1);

    Graph::vertex_type::edges_const_iterator et, eEnd(vx1.edgesEnd());
    for (et = vx1.edgesBegin(); et != eEnd; ++et) {
      v = g.edge(*et).otherVertex(v1);
      if (m_colormap[v] != boost::white_color) {
        d                = tcg::point_ops::dist(g.vertex(v).P(), vx1.P());
        if (d < dMin) v0 = v, dMin = d;
      }
    }

    assert(v0 >= 0);

    // Just add to the distance from that vertex
    m_distances[v1] = m_distances[v0] + dMin;
  }

  // Unused bfs methods

  void initialize_vertex(int v, const Graph &g) {}
  void discover_vertex(int v, const Graph &g) {}
  void examine_vertex(int v, const Graph &g) {}
  void finish_vertex(int v, const Graph &g) {}

  void examine_edge(const edge_descr &e, const Graph &g) {}
  void non_tree_edge(const edge_descr &e, const Graph &g) {}
  void gray_target(const edge_descr &e, const Graph &g) {}
  void black_target(const edge_descr &e, const Graph &g) {}
};

}  // namespace local

//==================================================================================

bool buildDistances(float *distances, const TTextureMesh &mesh,
                    const TPointD &pos, int *faceHint) {
  using namespace local;

  // First off, find the face containing the specified vertex
  int localF = -1, &f = faceHint ? *faceHint : localF;

  if (f < 0 || f >= mesh.facesCount() || !mesh.faceContains(f, pos))
    f = mesh.faceContaining(pos);

  if (f < 0) return false;

  int vCount = mesh.verticesCount();

  // Prepare the interpolation builder
  UCHAR *colorMap = (UCHAR *)calloc(vCount, sizeof(UCHAR));
  BFS_DistanceBuilder visitor(distances, colorMap);

  DistanceGreater gr(distances);
  std::priority_queue<int, std::vector<int>, DistanceGreater> verticesQueue(gr);

  int v0, v1, v2;
  mesh.faceVertices(f, v0, v1, v2);

  // Prepare BFS data
  distances[v0] = distances[v1] = distances[v2] = 0.0f;

  verticesQueue.push(v0), colorMap[v0] = boost::gray_color;
  verticesQueue.push(v1), colorMap[v1] = boost::gray_color;

  // Launch BFS algorithm
  boost::breadth_first_visit((const Graph &)mesh, v2, verticesQueue, visitor,
                             colorMap);

  free(colorMap);
  return true;
}

//----------------------------------------------------------------------------------

void buildSO(double *so, const TTextureMesh &mesh,
             const std::vector<PlasticHandle> &handles, int *faceHints) {
  int v, vCount = mesh.verticesCount();

  // Build the interpolant function data
  const TRectD &bbox = mesh.getBBox();

  const double len = std::max(bbox.getLx(), bbox.getLy()), val = 1e-8;
  const double k = -log(val) / len;

  // Allocate / initialize arrays
  float *distances = (float *)malloc(vCount * sizeof(float));
  double *wSums    = (double *)calloc(vCount, sizeof(double));

  memset(so, 0, vCount * sizeof(double));

  // Iterate handles - for each, add the corresponding interpolant
  int h, hCount = handles.size();
  for (h = 0; h != hCount; ++h) {
    const PlasticHandle &handle = handles[h];

    if (!buildDistances(distances, mesh, handle.m_pos,
                        faceHints ? &faceHints[h] : 0))
      continue;

    for (v = 0; v != vCount; ++v) {
      double w = fabs(distances[v]);

      wSums[v] += w = exp(-k * w) / (1e-3 + w);
      so[v] += w * handle.m_so;
    }
  }

  // Normalize so values by the wSums
  for (v = 0; v != vCount; ++v)
    if (wSums[v] != 0.0) so[v] /= wSums[v];

  // Cleanup
  free(wSums);
  free(distances);
}


#include "FWCore/Utilities/interface/EDMException.h"
#include "PhysicsTools/Utilities/interface/Parameter.h"
#include "PhysicsTools/Utilities/interface/Gaussian.h"
#include "PhysicsTools/Utilities/interface/Numerical.h"
#include "PhysicsTools/Utilities/interface/Exponential.h"
#include "PhysicsTools/Utilities/interface/Polynomial.h"
#include "PhysicsTools/Utilities/interface/Constant.h"
#include "PhysicsTools/Utilities/interface/Operations.h"
#include "PhysicsTools/Utilities/interface/MultiHistoChiSquare.h"
#include "PhysicsTools/Utilities/interface/HistoChiSquare.h"
#include "PhysicsTools/Utilities/interface/HistoPoissonLikelihoodRatio.h"
#include "PhysicsTools/Utilities/interface/RootMinuit.h"
#include "PhysicsTools/Utilities/interface/RootMinuitCommands.h"
#include "PhysicsTools/Utilities/interface/FunctClone.h"
#include "PhysicsTools/Utilities/interface/rootPlot.h"
#include "PhysicsTools/Utilities/interface/Expression.h"
#include "PhysicsTools/Utilities/interface/HistoPdf.h"
#include "FWCore/FWLite/interface/FWLiteEnabler.h"
#include "TROOT.h"
#include "TSystem.h"
#include "TH1.h"
#include "TFile.h"
#include <fstream>
#include <iostream>
#include <algorithm>
#include <exception>
#include <iterator>
#include <string>
#include <vector>
using namespace std;

// A helper function to simplify the main part.
template<class T>
ostream& operator<<(ostream& os, const vector<T>& v) {
  copy(v.begin(), v.end(), ostream_iterator<T>(cout, " "));
  return os;
}

// A function that get histogram and sets contents to 0
// if entries are too small
TH1 * getHisto(TFile * file, const char * name, unsigned int rebin) {
  TObject * h = file->Get(name);
  if(h == 0)
    throw edm::Exception(edm::errors::Configuration)
      << "Can't find object " << name << "\n";
  TH1 * histo = dynamic_cast<TH1*>(h);
  if(histo == 0)
    throw edm::Exception(edm::errors::Configuration)
      << "Object " << name << " is of type " << h->ClassName() << ", not TH1\n";
  histo->Rebin(rebin);
  for(int i = 1; i <= histo->GetNbinsX(); ++i) {
    if(histo->GetBinContent(i) < 0.1) {
      histo->SetBinContent(i, 0.0);
      histo->SetBinError(i, 0.0);
    }
  }
  return histo;
}

struct sig_tag;
struct bkg_tag;

typedef funct::FunctExpression Expr;
typedef fit::HistoChiSquare<funct::FunctExpression> ExprChi2;
typedef fit::HistoPoissonLikelihoodRatio<funct::FunctExpression> ExprPLR;

double fMin, fMax;
 unsigned int rebinMuMuNoIso ,rebinMuMu =1 , rebinMuMu1HLT , rebinMuMu2HLT , rebinMuTk , rebinMuSa;
 // assume that the bin size is 1 GeV!!!
string ext, region;
bool nonIsoTemplateFromMC;

template<typename T>
struct PlotPrefix { };

template<>
struct PlotPrefix<ExprChi2> {
  static string str() { return "chi2"; }
};

template<>
struct PlotPrefix<ExprPLR> {
  static string str() { return "plr"; }
};

template<typename T>
int main_t(const vector<string> & v_file){
  typedef fit::MultiHistoChiSquare<T, T, T, T, T> ChiSquared;
  fit::RootMinuitCommands<ChiSquared> commands("zChi2Fit.txt");

  cout << "minuit command file completed" << endl;

  funct::Constant rebinMuMuNoIsoConst(rebinMuMuNoIso), rebinMuMuConst(rebinMuMu),
    rebinMuMu1HLTConst(rebinMuMu1HLT), rebinMuMu2HLTConst(rebinMuMu2HLT),
    rebinMuTkConst(rebinMuTk), rebinMuSaConst(rebinMuSa);

  for(vector<string>::const_iterator it = v_file.begin(); it != v_file.end(); ++it) {
    TFile * root_file = new TFile(it->c_str(), "read");

    // default when region==all
    //    TH1 * histoZMuMuNoIso = getHisto(root_file, "nonIsolatedZToMuMuPlots/zMass",rebinMuMuNoIso);
    TH1 * histoZMuMuNoIso = getHisto(root_file, "oneNonIsolatedZToMuMuPlots/zMass",rebinMuMuNoIso);
    TH1 * histoZMuMu = getHisto(root_file, "goodZToMuMuPlots/zMass",rebinMuMu);
    TH1 * histoZMuMu1HLT = getHisto(root_file, "goodZToMuMu1HLTPlots/zMass", rebinMuMu1HLT);
    TH1 * histoZMuMu2HLT = getHisto(root_file, "goodZToMuMu2HLTPlots/zMass", rebinMuMu2HLT);
    TH1 * histoZMuTk = getHisto(root_file, "goodZToMuMuOneTrackPlots/zMass", rebinMuTk);
    TH1 * histoZMuSa = getHisto(root_file, "goodZToMuMuOneStandAloneMuonPlots/zMass", rebinMuSa);
    TH1 * histoZMuSaFromMuMu = getHisto(root_file, "zmumuSaMassHistogram/zMass", rebinMuSa);

    TH1 * histoZMuMuNoIsoTemplateFromMC= histoZMuMu;
    if (nonIsoTemplateFromMC) {
      //      histoZMuMuNoIsoTemplateFromMC = getHisto(root_file, "nonIsolatedZToMuMuPlotsMC/zMass",rebinMuMu);
      histoZMuMuNoIsoTemplateFromMC = getHisto(root_file, "oneNonIsolatedZToMuMuPlotsMC/zMass",rebinMuMu);
  }
    if (region=="barrel"){
      histoZMuMuNoIso = getHisto(root_file, "nonIsolatedZToMuMuPlotsBarrel/zMass",rebinMuMuNoIso);
      histoZMuMu = getHisto(root_file, "goodZToMuMuPlotsBarrel/zMass",rebinMuMu);
      histoZMuMu1HLT = getHisto(root_file, "goodZToMuMu1HLTPlotsBarrel/zMass", rebinMuMu1HLT);
      histoZMuMu2HLT = getHisto(root_file, "goodZToMuMu2HLTPlotsBarrel/zMass", rebinMuMu2HLT);
      histoZMuTk = getHisto(root_file, "goodZToMuMuOneTrackPlotsBarrel/zMass", rebinMuTk);
      histoZMuSa = getHisto(root_file, "goodZToMuMuOneStandAloneMuonPlotsBarrel/zMass", rebinMuSa);
      histoZMuSaFromMuMu = getHisto(root_file, "zmumuSaMassHistogramBarrel/zMass", rebinMuSa);
    }

    if (region=="endcap"){
      histoZMuMuNoIso = getHisto(root_file, "nonIsolatedZToMuMuPlotsEndCap/zMass",rebinMuMuNoIso);
      histoZMuMu = getHisto(root_file, "goodZToMuMuPlotsEndCap/zMass",rebinMuMu);
      histoZMuMu1HLT = getHisto(root_file, "goodZToMuMu1HLTPlotsEndCap/zMass", rebinMuMu1HLT);
      histoZMuMu2HLT = getHisto(root_file, "goodZToMuMu2HLTPlotsEndCap/zMass", rebinMuMu2HLT);
      histoZMuTk = getHisto(root_file, "goodZToMuMuOneTrackPlotsEndCap/zMass", rebinMuTk);
      histoZMuSa = getHisto(root_file, "goodZToMuMuOneStandAloneMuonPlotsEndCap/zMass", rebinMuSa);
      histoZMuSaFromMuMu = getHisto(root_file, "zmumuSaMassHistogramEndCap/zMass", rebinMuSa);
    }

    if (region=="barrend"){
      histoZMuMuNoIso = getHisto(root_file, "nonIsolatedZToMuMuPlotsBarrEnd/zMass",rebinMuMuNoIso);
      histoZMuMu = getHisto(root_file, "goodZToMuMuPlotsBarrEnd/zMass",rebinMuMu);
      histoZMuMu1HLT = getHisto(root_file, "goodZToMuMu1HLTPlotsBarrEnd/zMass", rebinMuMu1HLT);
      histoZMuMu2HLT = getHisto(root_file, "goodZToMuMu2HLTPlotsBarrEnd/zMass", rebinMuMu2HLT);
      histoZMuTk = getHisto(root_file, "goodZToMuMuOneTrackPlotsBarrEnd/zMass", rebinMuTk);
      histoZMuSa = getHisto(root_file, "goodZToMuMuOneStandAloneMuonPlotsBarrEnd/zMass", rebinMuSa);
      histoZMuSaFromMuMu = getHisto(root_file, "zmumuSaMassHistogramBarrEnd/zMass", rebinMuSa);
    }

    if (region!="endcap" && region!="barrel" && region!="barrend" && region!="all"  ){
      cout<< "not a valid region selected"<< endl;
      cout << "possible choises are: all, barrel, endcap, barrend " << endl;
      return 0;
    }

    cout << ">>> histogram loaded\n";
    string f_string = *it + "_" + PlotPrefix<T>::str() + "_";
    replace(f_string.begin(), f_string.end(), '.', '_');
    replace(f_string.begin(), f_string.end(), '/', '_');
    string plot_string = f_string + "." + ext;
    cout << ">>> Input files loaded\n";

    const char * kYieldZMuMu = "YieldZMuMu";
    const char * kEfficiencyTk = "EfficiencyTk";
    const char * kEfficiencySa = "EfficiencySa";
    const char * kEfficiencyIso = "EfficiencyIso";
    const char * kEfficiencyHLT = "EfficiencyHLT";
    const char * kYieldBkgZMuTk = "YieldBkgZMuTk";
    const char * kYieldBkgZMuSa = "YieldBkgZMuSa";
    const char * kYieldBkgZMuMuNotIso = "YieldBkgZMuMuNotIso";
    const char * kAlpha = "Alpha";
    const char * kBeta = "Beta";
    const char * kLambda = "Lambda";
    const char * kA0 = "A0";
    const char * kA1 = "A1";
    const char * kA2 = "A2";
    const char * kB0 = "B0";
    const char * kB1 = "B1";
    const char * kB2 = "B2";
    const char * kC0 = "C0";
    const char * kC1 = "C1";
    const char * kC2 = "C2";

    funct::Parameter yieldZMuMu(kYieldZMuMu, commands.par(kYieldZMuMu));
    funct::Parameter effTk(kEfficiencyTk, commands.par(kEfficiencyTk));
    funct::Parameter effSa(kEfficiencySa, commands.par(kEfficiencySa));
    funct::Parameter effIso(kEfficiencyIso, commands.par(kEfficiencyIso));
    funct::Parameter effHLT(kEfficiencyHLT, commands.par(kEfficiencyHLT));
    funct::Parameter yieldBkgZMuTk(kYieldBkgZMuTk, commands.par(kYieldBkgZMuTk));
    funct::Parameter yieldBkgZMuSa(kYieldBkgZMuSa, commands.par(kYieldBkgZMuSa));
    funct::Parameter yieldBkgZMuMuNotIso(kYieldBkgZMuMuNotIso, commands.par(kYieldBkgZMuMuNotIso));
    funct::Parameter alpha(kAlpha, commands.par(kAlpha));
    funct::Parameter beta(kBeta, commands.par(kBeta));
    funct::Parameter lambda(kLambda, commands.par(kLambda));
    funct::Parameter a0(kA0, commands.par(kA0));
    funct::Parameter a1(kA1, commands.par(kA1));
    funct::Parameter a2(kA2, commands.par(kA2));
    funct::Parameter b0(kB0, commands.par(kB0));
    funct::Parameter b1(kB1, commands.par(kB1));
    funct::Parameter b2(kB2, commands.par(kB2));
    funct::Parameter c0(kC0, commands.par(kC0));
    funct::Parameter c1(kC1, commands.par(kC1));
    funct::Parameter c2(kC2, commands.par(kC2));
    funct::Constant cFMin(fMin), cFMax(fMax);

    // count ZMuMu Yield
    double nZMuMu = 0, nZMuMu1HLT = 0, nZMuMu2HLT = 0;
    {
      unsigned int nBins = histoZMuMu->GetNbinsX();
      double xMin = histoZMuMu->GetXaxis()->GetXmin();
      double xMax = histoZMuMu->GetXaxis()->GetXmax();
      double deltaX =(xMax - xMin) / nBins;
      for(unsigned int i = 0; i < nBins; ++i) {
	double x = xMin + (i +.5) * deltaX;
	if(x > fMin && x < fMax){
	  nZMuMu += histoZMuMu->GetBinContent(i+1);
	  nZMuMu1HLT += histoZMuMu1HLT->GetBinContent(i+1);
	  nZMuMu2HLT += histoZMuMu2HLT->GetBinContent(i+1);
	}
      }
    }
    // aggiungi 1HLT 2HLT
    cout << ">>> count of ZMuMu yield in the range [" << fMin << ", " << fMax << "]: " << nZMuMu << endl;
    cout << ">>> count of ZMuMu (1HLT) yield in the range [" << fMin << ", " << fMax << "]: "  << nZMuMu1HLT << endl;
    cout << ">>> count of ZMuMu (2HLT) yield in the range [" << fMin << ", " << fMax << "]: "  << nZMuMu2HLT << endl;
    funct::RootHistoPdf zPdfMuMu(*histoZMuMu, fMin, fMax);
      //assign ZMuMu as pdf
    funct::RootHistoPdf zPdfMuMuNonIso = zPdfMuMu;
    if (nonIsoTemplateFromMC) {
      funct::RootHistoPdf zPdfMuMuNoIsoFromMC(*histoZMuMuNoIsoTemplateFromMC, fMin, fMax);
      zPdfMuMuNonIso = zPdfMuMuNoIsoFromMC;
    }

    funct::RootHistoPdf zPdfMuTk = zPdfMuMu;
    funct::RootHistoPdf zPdfMuMu1HLT = zPdfMuMu;
    funct::RootHistoPdf zPdfMuMu2HLT = zPdfMuMu;
    funct::RootHistoPdf zPdfMuSa(*histoZMuSaFromMuMu, fMin, fMax);
    zPdfMuMuNonIso.rebin(rebinMuMuNoIso/rebinMuMu);
    zPdfMuTk.rebin(rebinMuTk/rebinMuMu);
    zPdfMuMu1HLT.rebin(rebinMuMu1HLT/rebinMuMu);
    zPdfMuMu2HLT.rebin(rebinMuMu2HLT/rebinMuMu);

    funct::Numerical<2> _2;
    funct::Numerical<1> _1;

    //Efficiency term
    Expr zMuMuEff1HLTTerm = _2 * (effTk ^ _2) *  (effSa ^ _2) * (effIso ^ _2) * effHLT * (_1 - effHLT);
    Expr zMuMuEff2HLTTerm = (effTk ^ _2) *  (effSa ^ _2) * (effIso ^ _2) * (effHLT ^ _2) ;
    //    Expr zMuMuNoIsoEffTerm = (effTk ^ _2) * (effSa ^ _2) * (_1 - (effIso ^ _2)) * (_1 - ((_1 - effHLT)^_2));
    // change to both hlt and one not iso
    Expr zMuMuNoIsoEffTerm = _2 * (effTk ^ _2) * (effSa ^ _2) * effIso * (_1 - effIso) * (effHLT^_2);
    Expr zMuTkEffTerm = _2 * (effTk ^ _2) * effSa * (_1 - effSa) * (effIso ^ _2) * effHLT;
    Expr zMuSaEffTerm = _2 * (effSa ^ _2) * effTk * (_1 - effTk) * (effIso ^ _2) * effHLT;

    Expr zMuMu1HLT = rebinMuMu1HLTConst * zMuMuEff1HLTTerm * yieldZMuMu;
    Expr zMuMu2HLT = rebinMuMu2HLTConst * zMuMuEff2HLTTerm * yieldZMuMu;

    Expr zMuTkBkg = yieldBkgZMuTk * funct::Exponential(lambda)* funct::Polynomial<2>(a0, a1, a2);
    Expr zMuTkBkgScaled = rebinMuTkConst * zMuTkBkg;
    Expr zMuTk = rebinMuTkConst * (zMuTkEffTerm * yieldZMuMu * zPdfMuTk + zMuTkBkg);

    Expr zMuMuNoIsoBkg = yieldBkgZMuMuNotIso * funct::Exponential(alpha)* funct::Polynomial<2>(b0, b1, b2);
    Expr zMuMuNoIsoBkgScaled = rebinMuMuNoIsoConst * zMuMuNoIsoBkg;
    Expr zMuMuNoIso = rebinMuMuNoIsoConst * (zMuMuNoIsoEffTerm * yieldZMuMu * zPdfMuMuNonIso + zMuMuNoIsoBkg);

    Expr zMuSaBkg = yieldBkgZMuSa * funct::Exponential(beta)* funct::Polynomial<2>(c0, c1, c2);
    Expr zMuSaBkgScaled = rebinMuSaConst * zMuSaBkg;
    Expr zMuSa = rebinMuSaConst * (zMuSaEffTerm * yieldZMuMu * zPdfMuSa  + zMuSaBkg );

    TH1D histoZCount1HLT("histoZCount1HLT", "", 1, fMin, fMax);
    histoZCount1HLT.Fill(100, nZMuMu1HLT);
    TH1D histoZCount2HLT("histoZCount2HLT", "", 1, fMin, fMax);
    histoZCount2HLT.Fill(100, nZMuMu2HLT);

    ChiSquared chi2(zMuMu1HLT, & histoZCount1HLT,
		    zMuMu2HLT, & histoZCount2HLT,
		    zMuTk, histoZMuTk,
		    zMuSa, histoZMuSa,
		    zMuMuNoIso,histoZMuMuNoIso,
		    fMin, fMax);
    cout << "N. bins: " << chi2.numberOfBins() << endl;

    fit::RootMinuit<ChiSquared> minuit(chi2, true);
    commands.add(minuit, yieldZMuMu);
    commands.add(minuit, effTk);
    commands.add(minuit, effSa);
    commands.add(minuit, effIso);
    commands.add(minuit, effHLT);
    commands.add(minuit, yieldBkgZMuTk);
    commands.add(minuit, yieldBkgZMuSa);
    commands.add(minuit, yieldBkgZMuMuNotIso);
    commands.add(minuit, lambda);
    commands.add(minuit, alpha);
    commands.add(minuit, beta);
    commands.add(minuit, a0);
    commands.add(minuit, a1);
    commands.add(minuit, a2);
    commands.add(minuit, b0);
    commands.add(minuit, b1);
    commands.add(minuit, b2);
    commands.add(minuit, c0);
    commands.add(minuit, c1);
    commands.add(minuit, c2);
    commands.run(minuit);
    const unsigned int nPar = 20;//WARNIG: this must be updated manually for now
    ROOT::Math::SMatrix<double, nPar, nPar, ROOT::Math::MatRepSym<double, nPar> > err;
    minuit.getErrorMatrix(err);

    std::cout << "error matrix:" << std::endl;
    for(unsigned int i = 0; i < nPar; ++i) {
      for(unsigned int j = 0; j < nPar; ++j) {
	  std::cout << err(i, j) << "\t";
      }
      std::cout << std::endl;
    }
    minuit.printFitResults();
    ofstream myfile;
    myfile.open ("fitResult.txt", ios::out | ios::app);
    myfile<<"\n";
    double Y =  minuit.getParameterError("YieldZMuMu");
    double dY = minuit.getParameterError("YieldZMuMu", Y);
    double tk_eff =  minuit.getParameterError("EfficiencyTk");
    double dtk_eff = minuit.getParameterError("EfficiencyTk", tk_eff);
    double sa_eff =  minuit.getParameterError("EfficiencySa");
    double dsa_eff = minuit.getParameterError("EfficiencySa", sa_eff);
    double iso_eff =  minuit.getParameterError("EfficiencyIso");
    double diso_eff = minuit.getParameterError("EfficiencyIso", iso_eff);
    double hlt_eff =  minuit.getParameterError("EfficiencyHLT");
    double dhlt_eff = minuit.getParameterError("EfficiencyHLT",hlt_eff);
    myfile<< Y <<" "<< dY <<" "<< tk_eff <<" "<< dtk_eff <<" "<< sa_eff << " " << dsa_eff << " " << iso_eff <<" " << diso_eff<< " " << hlt_eff << " " << dhlt_eff << " " <<chi2()/(chi2.numberOfBins()- minuit.numberOfFreeParameters());

    myfile.close();

    //Plot
    double s;
    s = 0;
    for(int i = 1; i <= histoZMuMuNoIso->GetNbinsX(); ++i)
      s += histoZMuMuNoIso->GetBinContent(i);
    histoZMuMuNoIso->SetEntries(s);
    s = 0;
    for(int i = 1; i <= histoZMuMu->GetNbinsX(); ++i)
      s += histoZMuMu->GetBinContent(i);
    histoZMuMu->SetEntries(s);
    s = 0;
    for(int i = 1; i <= histoZMuMu1HLT->GetNbinsX(); ++i)
      s += histoZMuMu1HLT->GetBinContent(i);
    histoZMuMu1HLT->SetEntries(s);
    s = 0;
    for(int i = 1; i <= histoZMuMu2HLT->GetNbinsX(); ++i)
      s += histoZMuMu2HLT->GetBinContent(i);
    histoZMuMu2HLT->SetEntries(s);
    s = 0;
    for(int i = 1; i <= histoZMuTk->GetNbinsX(); ++i)
      s += histoZMuTk->GetBinContent(i);
    histoZMuTk->SetEntries(s);
    s = 0;
    for(int i = 1; i <= histoZMuSa->GetNbinsX(); ++i)
      s += histoZMuSa->GetBinContent(i);
    histoZMuSa->SetEntries(s);

    string ZMuMu1HLTPlot = "ZMuMu1HLTFit_" + plot_string;
    root::plot<Expr>(ZMuMu1HLTPlot.c_str(), *histoZMuMu1HLT, zMuMu1HLT, fMin, fMax,
		     effTk, effSa, effIso, effHLT, yieldZMuMu,
		     kOrange-2, 2, kSolid, 100,
		     "Z -> #mu #mu mass", "#mu #mu invariant mass (GeV/c^{2})",
		     "Events");

    string ZMuMu2HLTPlot = "ZMuMu2HLTFit_" + plot_string;
    root::plot<Expr>(ZMuMu2HLTPlot.c_str(), *histoZMuMu2HLT, zMuMu2HLT, fMin, fMax,
		     effTk, effSa, effIso, effHLT, yieldZMuMu,
		     kOrange-2, 2, kSolid, 100,
		     "Z -> #mu #mu mass", "#mu #mu invariant mass (GeV/c^{2})",
		     "Events");

    string ZMuMuNoIsoPlot = "ZMuMuNoIsoFit_X_" + plot_string;
    root::plot<Expr>(ZMuMuNoIsoPlot.c_str(), *histoZMuMuNoIso, zMuMuNoIso, fMin, fMax,
		     effTk, effSa, effIso, effHLT, yieldZMuMu,
		     yieldBkgZMuMuNotIso, alpha, b0, b1, b2,
		     kWhite, 2, kSolid, 100,
		     "Z -> #mu #mu Not Iso mass", "#mu #mu invariant mass (GeV/c^{2})",
		     "Events");
    ZMuMuNoIsoPlot = "ZMuMuNoIsoFit_" + plot_string;
    TF1 funZMuMuNoIso = root::tf1_t<sig_tag, Expr>("ZMuMuNoIsoFunction", zMuMuNoIso, fMin, fMax,
					      effTk, effSa, effIso, effHLT, yieldZMuMu,
					      yieldBkgZMuMuNotIso, alpha, b0, b1, b2);
    funZMuMuNoIso.SetLineColor(kOrange+8);
    funZMuMuNoIso.SetLineWidth(3);
    //funZMuMuNoIso.SetLineStyle(kDashed);

    //funZMuMuNoIso.SetFillColor(kOrange-2);
    //funZMuMuNoIso.SetFillStyle(3325);

    funZMuMuNoIso.SetNpx(10000);
    TF1 funZMuMuNoIsoBkg = root::tf1_t<bkg_tag, Expr>("ZMuMuNoIsoBack", zMuMuNoIsoBkgScaled, fMin, fMax,
						 yieldBkgZMuMuNotIso, alpha, b0, b1, b2);
    funZMuMuNoIsoBkg.SetLineColor(kViolet+3);
    funZMuMuNoIsoBkg.SetLineWidth(2);
    funZMuMuNoIsoBkg.SetLineStyle(kSolid);
    funZMuMuNoIsoBkg.SetFillColor(kViolet-5);
    funZMuMuNoIsoBkg.SetFillStyle(3357);

    funZMuMuNoIsoBkg.SetNpx(10000);
    histoZMuMuNoIso->SetTitle("Z -> #mu #mu Not Iso mass");
    histoZMuMuNoIso->SetXTitle("#mu +  #mu invariant mass (GeV/c^{2})");
    histoZMuMuNoIso->SetYTitle("Events");
    TCanvas *canvas = new TCanvas("canvas");
    histoZMuMuNoIso->Draw("e");
    funZMuMuNoIsoBkg.Draw("same");
    funZMuMuNoIso.Draw("same");
    canvas->SaveAs(ZMuMuNoIsoPlot.c_str());
    canvas->SetLogy();
    string logZMuMuNoIsoPlot = "log_" + ZMuMuNoIsoPlot;
    canvas->SaveAs(logZMuMuNoIsoPlot.c_str());

    double IntSigMMNotIso = ((double) rebinMuMu/  (double) rebinMuMuNoIso) * funZMuMuNoIso.Integral(fMin, fMax);
    double IntSigMMNotIsoBkg = ((double) rebinMuMu/  (double) rebinMuMuNoIso) * funZMuMuNoIsoBkg.Integral(fMin, fMax);
    cout << "*********  ZMuMuNoIsoPlot signal yield from the fit ==> " <<  IntSigMMNotIso << endl;
    cout << "*********  ZMuMuNoIsoPlot background yield from the fit ==> " << IntSigMMNotIsoBkg << endl;

    string ZMuTkPlot = "ZMuTkFit_X_" + plot_string;
    root::plot<Expr>(ZMuTkPlot.c_str(), *histoZMuTk, zMuTk, fMin, fMax,
		     effTk, effSa, effIso, effHLT, yieldZMuMu,
		     yieldBkgZMuTk, lambda, a0, a1, a2,
		     kOrange+3, 2, kSolid, 100,
		     "Z -> #mu + (unmatched) track mass", "#mu #mu invariant mass (GeV/c^{2})",
		     "Events");
    ZMuTkPlot = "ZMuTkFit_" + plot_string;
    TF1 funZMuTk = root::tf1_t<sig_tag, Expr>("ZMuTkFunction", zMuTk, fMin, fMax,
					      effTk, effSa, effIso, effHLT, yieldZMuMu,
					      yieldBkgZMuTk, lambda, a0, a1, a2);
    funZMuTk.SetLineColor(kOrange+8);
    funZMuTk.SetLineWidth(3);
    funZMuTk.SetLineStyle(kSolid);
    //  funZMuTk.SetFillColor(kOrange-2);
    //funZMuTk.SetFillStyle(3325);
    funZMuTk.SetNpx(10000);
    TF1 funZMuTkBkg = root::tf1_t<bkg_tag, Expr>("ZMuTkBack", zMuTkBkgScaled, fMin, fMax,
						 yieldBkgZMuTk, lambda, a0, a1, a2);
    funZMuTkBkg.SetLineColor(kViolet+3);
    funZMuTkBkg.SetLineWidth(2);
    funZMuTkBkg.SetLineStyle(kSolid);
    funZMuTkBkg.SetFillColor(kViolet-5);
    funZMuTkBkg.SetFillStyle(3357);
    funZMuTkBkg.SetNpx(10000);
    histoZMuTk->SetTitle("Z -> #mu + (unmatched) track mass");
    histoZMuTk->SetXTitle("#mu + (unmatched) track invariant mass (GeV/c^{2})");
    histoZMuTk->SetYTitle("Events");
    TCanvas *canvas_ = new TCanvas("canvas_");
    histoZMuTk->Draw("e");
    funZMuTkBkg.Draw("same");
    funZMuTk.Draw("same");
    canvas_->SaveAs(ZMuTkPlot.c_str());
    canvas_->SetLogy();
    string logZMuTkPlot = "log_" + ZMuTkPlot;
    canvas_->SaveAs(logZMuTkPlot.c_str());

    double IntSigMT = ((double) rebinMuMu/  (double) rebinMuTk) * funZMuTk.Integral(fMin, fMax);
    double IntSigMTBkg = ((double) rebinMuMu/  (double) rebinMuTk) * funZMuTkBkg.Integral(fMin, fMax);
    cout << "*********  ZMuMuTkPlot signal yield from the fit ==> " <<  IntSigMT << endl;
    cout << "*********  ZMuMuTkPlot background yield from the fit ==> " << IntSigMTBkg << endl;

    string ZMuSaPlot = "ZMuSaFit_X_" + plot_string;
    root::plot<Expr>(ZMuSaPlot.c_str(), *histoZMuSa, zMuSa, fMin, fMax,
		     effSa, effTk, effIso, yieldZMuMu,
		     yieldBkgZMuSa, beta, c0, c1, c2 ,
		     kOrange+3, 2, kSolid, 100,
		     "Z -> #mu + (unmatched) standalone mass",
		     "#mu + (unmatched) standalone invariant mass (GeV/c^{2})",
		     "Events");

    ZMuSaPlot = "ZMuSaFit_" + plot_string;
    TF1 funZMuSa = root::tf1_t<sig_tag, Expr>("ZMuSaFunction", zMuSa, fMin, fMax,
					      effTk, effSa, effIso, effHLT, yieldZMuMu,
					      yieldBkgZMuSa, beta, c0, c1, c2);
    funZMuSa.SetLineColor(kOrange+8);
    funZMuSa.SetLineWidth(3);
    funZMuSa.SetLineStyle(kSolid);
    // funZMuSa.SetFillColor(kOrange-2);
    // funZMuSa.SetFillStyle(3325);
    funZMuSa.SetNpx(10000);
    TF1 funZMuSaBkg = root::tf1_t<bkg_tag, Expr>("ZMuSaBack", zMuSaBkgScaled, fMin, fMax,
						 yieldBkgZMuSa, beta, c0, c1, c2);
    funZMuSaBkg.SetLineColor(kViolet+3);
    funZMuSaBkg.SetLineWidth(2);
    funZMuSaBkg.SetLineStyle(kSolid);
    funZMuSaBkg.SetFillColor(kViolet-5);
    funZMuSaBkg.SetFillStyle(3357);
    funZMuSaBkg.SetNpx(10000);
    histoZMuSa->SetTitle("Z -> #mu + (unmatched) standalone mass");
    histoZMuSa->SetXTitle("#mu + (unmatched) standalone invariant mass (GeV/c^{2})");
    histoZMuSa->SetYTitle("Events");
    TCanvas *canvas__ = new TCanvas("canvas__");
    histoZMuSa->Draw("e");
    funZMuSaBkg.Draw("same");
    funZMuSa.Draw("same");
    canvas__->SaveAs(ZMuSaPlot.c_str());
    canvas__->SetLogy();
    string logZMuSaPlot = "log_" + ZMuSaPlot;
    canvas__->SaveAs(logZMuSaPlot.c_str());

    double IntSigMS = ((double) rebinMuMu/  (double) rebinMuSa) * funZMuSa.Integral(fMin, fMax);
    double IntSigMSBkg = ((double) rebinMuMu/  (double) rebinMuSa) * funZMuSaBkg.Integral(fMin, fMax);
    cout << "*********  ZMuMuSaPlot signal yield from the fit ==> " <<  IntSigMS << endl;
    cout << "*********  ZMuMuSaPlot background yield from the fit ==> " << IntSigMSBkg << endl;

  }
  return 0;
}

#include <boost/program_options.hpp>
using namespace boost;
namespace po = boost::program_options;

int main(int ac, char *av[]) {
  po::options_description desc("Allowed options");
  desc.add_options()
    ("help,h", "produce help message")
    ("input-file,i", po::value<vector<string> >(), "input file")
    ("min,m", po::value<double>(&fMin)->default_value(60), "minimum value for fit range")
    ("max,M", po::value<double>(&fMax)->default_value(120), "maximum value for fit range")
    ("rebins,R", po::value<vector<unsigned int> >(), "rebins values: rebinMuMu2HLT , rebinMuMu1HLT , rebinMuMuNoIso , rebinMuSa, rebinMuTk")
    ("chi2,c", "perform chi-squared fit")
    ("plr,p", "perform Poisson likelihood-ratio fit")
    ("nonIsoTemplateFromMC,I", po::value<bool>(&nonIsoTemplateFromMC)->default_value(false) , "take the template for nonIso sample from MC")
    ("plot-format,f", po::value<string>(&ext)->default_value("eps"), "output plot format")
    ("detectorRegion,r",po::value<string> (&region)->default_value("all"), "detector region in which muons are detected" );
  po::positional_options_description p;
  p.add("input-file", -1);
  p.add("rebins", -1);

  po::variables_map vm;
  po::store(po::command_line_parser(ac, av).
	    options(desc).positional(p).run(), vm);
  po::notify(vm);

  if (vm.count("help")) {
    cout << "Usage: options_description [options]\n";
    cout << desc;
    return 0;
  }

  if (!vm.count("input-file")) {
    return 1;
  }
  cout << "Input files are: "
       << vm["input-file"].as< vector<string> >() << "\n";
  vector<string> v_file = vm["input-file"].as< vector<string> >();

  if (vm.count("rebins") ) {
    vector<unsigned int> v_rebin = vm["rebins"].as< vector<unsigned int> >();
    if (v_rebin.size()!=5){
      cerr << " please provide 5 numbers in the given order:  rebinMuMu2HLT , rebinMuMu1HLT , rebinMuMuNoIso, rebinMuSa, rebinMuTk \n";
      return 1;
    }
 rebinMuMuNoIso = v_rebin[2], rebinMuMu1HLT = v_rebin[1], rebinMuMu2HLT = v_rebin[0], rebinMuTk = v_rebin[4], rebinMuSa = v_rebin[3];
  }

  bool chi2Fit = vm.count("chi2"), plrFit = vm.count("plr");

  if(!(chi2Fit||plrFit))
    cerr << "Warning: no fit performed. Please, specify either -c or -p options or both" << endl;

  gROOT->SetStyle("Plain");

  int ret = 0;
  try {
    if(plrFit) {
      std::cout << "==================================== " << std::endl;
      std::cout << "=== Poisson Likelihood Ratio fit === " << std::endl;
      std::cout << "==================================== " << std::endl;
      int ret2 = main_t<ExprPLR>(v_file);
      if(ret2 != 0) ret = 1;
    }
    if(chi2Fit) {
      std::cout << "================= " << std::endl;
      std::cout << "=== Chi-2 fit === " << std::endl;
      std::cout << "================= " << std::endl;
      int ret1 = main_t<ExprChi2>(v_file);
      if(ret1 != 0) ret = 1;
    }
  }
  catch(std::exception& e) {
    cerr << "error: " << e.what() << "\n";
    return 1;
  }
  catch(...) {
    cerr << "Exception of unknown type!\n";
  }
  return ret;
}