// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

#include <unordered_map>
#include <cstdint>
#include <string>
#include <iostream>
#include <cstdlib>
#include <time.h>
#include <mutex>
#include <set>
#include <fstream>
#include <thread>
#include <algorithm>
#include <stdlib.h>

#include "timer.h"
#include "rand_int_rng.h"
#include "lda_document.h"
#include "data_block.h"
#include "model_block.h"
#include "lda_engine.hpp"
#include "utils.hpp"
#include "simple_barrier.h"
#include "light_doc_sampler.hpp"

#ifdef _MSC_VER
#include "windows.h"
#elif defined(__APPLE__)
#include <mach/mach_types.h>
#include <mach/thread_act.h>
#include <pthread.h>
#else
#include "sched.h"
#endif

namespace lda {
 LdaEngine::LdaEngine(int numTopic,
 int numVocab,
 float alphaSum,
 float beta,
 int numIter,
 int likelihoodInterval,
 int numThread,
 int mhstep,
 int maxDocToken)
 : K_(numTopic),
 V_(numVocab),
 compute_ll_interval_(likelihoodInterval),
 beta_(beta),
 num_iterations_(numIter),
 mh_step_(mhstep),
 alpha_sum_(alphaSum),
 maxDocToken_(maxDocToken),
 samplers_(nullptr),
 document_buffer_(nullptr)
 {
 if (numThread > 0)
 {
 num_threads_ = numThread;
 }
 else
 {
 unsigned int uNumCPU = std::thread::hardware_concurrency();
 num_threads_ = std::max(1, (int)(uNumCPU - 2));
 }
 printf("using %d thread(s) to do train/test\n", num_threads_);

 bAlphaSumMultiplied = false;
 atomic_stats_ = new LDAEngineAtomics();
 model_block_ = new LDAModelBlock();
 data_block_ = new LDADataBlock(num_threads_);
 process_barrier_ = new SimpleBarrier(num_threads_);
 samplerQueue_ = new CBlockedIntQueue();

 document_buffer_ = new int32_t*[num_threads_];
 for (int i = 0; i < num_threads_; i++)
 document_buffer_[i] = new int32_t[maxDocToken_ * 2 + 1];

 likelihood_in_iter_ = nullptr;

 beta_sum_ = beta_ * V_;
 }

 LdaEngine::LdaEngine(int32_t K, int32_t V, int32_t num_threads, int32_t compute_ll_interval, float beta, int32_t num_iterations, int32_t mh_step, float alpha_sum, int maxDocToken)
 : K_(K),
 V_(V),
 compute_ll_interval_(compute_ll_interval),
 beta_(beta),
 num_iterations_(num_iterations),
 mh_step_(mh_step),
 alpha_sum_(alpha_sum),
 maxDocToken_(maxDocToken),
 samplers_(nullptr),
 document_buffer_(nullptr)
 {
 if (num_threads > 0)
 {
 num_threads_ = num_threads;
 }
 else
 {
 unsigned int uNumCPU = std::thread::hardware_concurrency();
 num_threads_ = std::max(1, (int)(uNumCPU - 2));
 }
 bAlphaSumMultiplied = false;
 process_barrier_ = new SimpleBarrier(num_threads_);
 atomic_stats_ = new LDAEngineAtomics();
 data_block_ = new LDADataBlock(num_threads_);
 model_block_ = new LDAModelBlock();
 samplerQueue_ = new CBlockedIntQueue();

 document_buffer_ = new int32_t*[num_threads_];
 for (int i = 0; i < num_threads_; i++)
 document_buffer_[i] = new int32_t[maxDocToken_ * 2 + 1];

 likelihood_in_iter_ = nullptr;
 beta_sum_ = beta_ * V_;
 }


 LdaEngine::~LdaEngine()
 {
 //delete memory space
 delete process_barrier_;
 process_barrier_ = nullptr;

 delete data_block_;
 data_block_ = nullptr;

 delete atomic_stats_;
 atomic_stats_ = nullptr;

 delete model_block_;
 model_block_ = nullptr;

 delete samplerQueue_;
 samplerQueue_ = nullptr;

 for (int i = 0; i < num_threads_; ++i)
 {
 delete samplers_[i];
 }
 delete[] samplers_;

 if (document_buffer_)
 {
 for (int i = 0; i < num_threads_; ++i)
 {
 delete[]document_buffer_[i];
 document_buffer_[i] = nullptr;
 }
 delete[]document_buffer_;
 document_buffer_ = nullptr;
 }

 if (likelihood_in_iter_)
 {
 delete[] likelihood_in_iter_;
 likelihood_in_iter_ = nullptr;
 }
 }

 bool LdaEngine::InitializeBeforeTrain()
 {
 CTimer tmDebug(true);
 CheckFunction(0, tmDebug, "enter initializeBeforeTrain", false);
 //allocate model memory from the data preloaded
 AllocateModelMemory(data_block_);
 CheckFunction(0, tmDebug, "allocate model memory", false);

 double alloc_start = lda::get_time();
 global_word_topic_table_.resize(V_);
 alias_rng_int_.Init(K_);
 beta_k_v_.resize(K_);
 global_alias_k_v_.resize(V_);

 for (int i = 0; i < V_; ++i)
 {
 global_alias_k_v_[i] = model_block_->get_alias_row(i);
 }
 global_summary_row_.resize(K_);
 CheckFunction(0, tmDebug, "initlaizing global tables used in sampling", false);

 word_range_for_each_thread_.resize(num_threads_ + 1);
 int32_t word_num_each_thread = V_ / num_threads_;
 word_range_for_each_thread_[0] = 0;
 for (int32_t i = 0; i < num_threads_ - 1; ++i)
 {
 word_range_for_each_thread_[i + 1] = word_range_for_each_thread_[i] + word_num_each_thread;
 }
 word_range_for_each_thread_[num_threads_] = V_;

 //setup sampler
 samplers_ = new LightDocSampler*[num_threads_];
 samplerQueue_->clear();

 for (int i = 0; i < num_threads_; ++i)
 {
 samplers_[i] = new LightDocSampler(
 K_,
 V_,
 num_threads_,
 mh_step_,
 beta_,
 alpha_sum_,
 global_word_topic_table_,
 global_summary_row_,
 global_alias_k_v_,
 beta_height_,
 beta_mass_,
 beta_k_v_);

 samplerQueue_->push(i);
 }
 CheckFunction(0, tmDebug, "create samplers", false);
 return true;
 }

 void LdaEngine::InitializeBeforeTest()
 {
 // TODO:
 // Allocating space for word-topic-table and alias table according to the input data of SetModel interface (done)
 // Create multiple thread-specific sampler
 // set word_range_for_each_thread_
 // Adjust the alpha_sum_ parameter for each thread-specific sampler
 CTimer tmDebug(true);
 CheckFunction(0, tmDebug, "enter initializeBeforeTest", false);

 global_word_topic_table_.resize(V_);
 alias_rng_int_.Init(K_);
 beta_k_v_.resize(K_);
 global_alias_k_v_.resize(V_);

 for (int i = 0; i < V_; ++i)
 {
 global_alias_k_v_[i] = model_block_->get_alias_row(i);
 }
 CheckFunction(0, tmDebug, "initlaizing global tables used in sampling", false);

 // Set the word range for each thread
 word_range_for_each_thread_.resize(num_threads_ + 1);
 int32_t word_num_each_thread = V_ / num_threads_;
 word_range_for_each_thread_[0] = 0;
 for (int32_t i = 0; i < num_threads_ - 1; ++i)
 {
 word_range_for_each_thread_[i + 1] = word_range_for_each_thread_[i] + word_num_each_thread;
 }
 word_range_for_each_thread_[num_threads_] = V_;

 //setup sampler
 if (samplers_)
 {
 for (int i = 0; i < num_threads_; ++i)
 {
 delete samplers_[i];
 }
 delete[] samplers_;
 }
 if (document_buffer_)
 {
 for (int i = 0; i < num_threads_; ++i)
 {
 delete[]document_buffer_[i];
 document_buffer_[i] = nullptr;
 }
 delete[]document_buffer_;
 document_buffer_ = nullptr;
 }

 samplers_ = new LightDocSampler*[num_threads_];
 document_buffer_ = new int32_t*[num_threads_];
 samplerQueue_->clear();

 for (int i = 0; i < num_threads_; ++i)
 {
 samplers_[i] = new LightDocSampler(
 K_,
 V_,
 num_threads_,
 mh_step_,
 beta_,
 alpha_sum_,
 global_word_topic_table_,
 global_summary_row_,
 global_alias_k_v_,
 beta_height_,
 beta_mass_,
 beta_k_v_);

 samplers_[i]->AdaptAlphaSum(false);
 document_buffer_[i] = new int32_t[maxDocToken_ * 2 + 1];

 samplerQueue_->push(i);
 }
 CheckFunction(0, tmDebug, "create samplers", false);

 // build alias table
 // build alias table for the dense term, beta_k_v_, which is shared by all the words
 beta_mass_ = 0;
 std::vector<float> proportion(K_);
 for (int k = 0; k < K_; ++k)
 {
 proportion[k] = beta_ / (global_summary_row_[k] + beta_sum_);
 beta_mass_ += proportion[k];
 }
 alias_rng_int_.SetProportionMass(proportion, beta_mass_, beta_k_v_, &beta_height_, samplers_[0]->rng());

 // build alias table for the sparse term
 for (int thread_id = 0; thread_id < num_threads_; ++thread_id)
 {
 LightDocSampler &sampler = *(samplers_[thread_id]);
 sampler.build_alias_table(word_range_for_each_thread_[thread_id], word_range_for_each_thread_[thread_id + 1], thread_id);
 }
 CheckFunction(0, tmDebug, "build alisa table", false);
 }

 void LdaEngine::Train(const char* pTrainOutput)
 {
 std::vector<std::thread> threads(num_threads_);
 atomic_stats_->thread_counter_ = 0;

 for (auto& thr : threads) {
 thr = std::thread(&LdaEngine::Training_Thread, this);
 }

 printf("started training with %d threads\n", num_threads_);
 for (auto& thr : threads) {
 thr.join();
 }

 if (pTrainOutput)
 {
 DumpDocTopicTable(pTrainOutput);
 }
 }

 void LdaEngine::Test(int32_t burnin_iter, float* pLoglikelihood)
 {
 std::vector<std::thread> threads(num_threads_);
 atomic_stats_->thread_counter_ = 0;
 burnin_iterations_ = burnin_iter;

 likelihood_in_iter_ = new float[burnin_iterations_];
 for (int i = 0; i < burnin_iterations_; i++)
 {
 likelihood_in_iter_[i] = 0.0;
 }

 for (auto& thr : threads) {
 thr = std::thread(&LdaEngine::Testing_Thread, this);
 }

 printf("started testing with %d threads\n", num_threads_);

 for (auto& thr : threads) {
 thr.join();
 }

 //get the loglikelihood of each burn in iteration
 for (int i = 0; i < burnin_iterations_; i++)
 {
 pLoglikelihood[i] = likelihood_in_iter_[i]; //just set an arbitary value here for later update
 }
 }

 void LdaEngine::CheckFunction(int thread_id, CTimer &tmDebug, const char* msg, bool waitBarrier)
 {
 }

 void LdaEngine::Training_Thread()
 {
 CTimer tmDebug(true);

 int thread_id = atomic_stats_->thread_counter_++;
 std::vector<std::pair<int, double>> llcontainer;
 // Set core affinity which helps performance improvement
#ifdef _MSC_VER
 DWORD maskLL = 0;
 maskLL |= (1LL << (thread_id));
 DWORD_PTR mask = maskLL;
 SetThreadAffinityMask(GetCurrentThread(), mask);
#elif defined(__APPLE__)
 thread_port_t thread = pthread_mach_thread_np(pthread_self());
 thread_affinity_policy_data_t policy = { thread_id };
 thread_policy_set(thread, THREAD_AFFINITY_POLICY, (thread_policy_t)&policy, 1);
#else
 cpu_set_t set;
 CPU_ZERO(&set);
 CPU_SET(thread_id, &set);
 sched_setaffinity(0, sizeof(cpu_set_t), &set);
#endif

 // Each thread builds a portion of word-topic table. We do this way because each word-topic row 
 // has a thread-specific buffer for rehashing
 process_barrier_->wait();
 LightDocSampler &sampler_ = *(samplers_[thread_id]);
 sampler_.AdaptAlphaSum(true);

 sampler_.build_word_topic_table(thread_id, num_threads_, *model_block_);
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "intialize word_topic_table for sampler - in function train_thread");

 int32_t token_num = 0;
 int32_t doc_start = data_block_->Begin(thread_id);
 int32_t doc_end = data_block_->End(thread_id);

 for (int32_t doc_index = doc_start; doc_index != doc_end; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 int doc_size = doc->size();
 for (int i = 0; i < doc_size; ++i)
 {
 int topic = sampler_.rand_k();
 doc->SetTopic(i, topic);
 }
 int cursor = doc->get_cursor();
 token_num += sampler_.GlobalInit(doc.get());
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "intialize token topic before iterations - in function train_thread");

 for (int i = 0; i < num_threads_; ++i)
 {
 std::vector<word_topic_delta>& wtd_vec = samplers_[i]->get_word_topic_delta(thread_id);
 for (auto& wtd : wtd_vec)
 {
 global_word_topic_table_[wtd.word].inc(wtd.topic, wtd.delta);
 }
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "intialize word topic model before iterations - in function train_thread");

 // use thread-private delta table to get global table
 {
 std::lock_guard<std::mutex> lock(atomic_stats_->global_mutex_);

 std::vector<int64_t> &summary = sampler_.get_delta_summary_row();
 for (int i = 0; i < K_; ++i)
 {
 global_summary_row_[i] += summary[i];
 }
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "global summary & Complete setup train before iterations - in function train_thread");

 for (int iter = 0; iter < num_iterations_; ++iter)
 {
 CheckFunction(thread_id, tmDebug, "----------------------iteration start - in function train_thread---------------------");
 int32_t token_sweeped = 0;
 atomic_stats_->num_tokens_clock_ = 0;
 // build alias table
 // build alias table for the dense term, beta_k_v_, which is shared by all the words
 if (thread_id == 0)
 {
 beta_mass_ = 0;
 std::vector<float> proportion(K_);
 for (int k = 0; k < K_; ++k)
 {
 proportion[k] = beta_ / (global_summary_row_[k] + beta_sum_);
 beta_mass_ += proportion[k];
 }

 alias_rng_int_.SetProportionMass(proportion, beta_mass_, beta_k_v_, &beta_height_, sampler_.rng());
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "built alias table dense - in function train_thread");

 // build alias table for the sparse term
 sampler_.build_alias_table(word_range_for_each_thread_[thread_id], word_range_for_each_thread_[thread_id + 1], thread_id);
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "built alias table sparse - in function train_thread");

 sampler_.EpocInit();
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "EpochInit - in function train_thread");

 // main part of the training - sampling over documents in this iteration
 double iter_start = lda::get_time();
 int32_t doc_start_local = data_block_->Begin(thread_id);
 int32_t doc_end_local = data_block_->End(thread_id);

 for (int32_t doc_index = doc_start_local; doc_index != doc_end_local; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 token_sweeped += sampler_.SampleOneDoc(doc.get());
 }
 atomic_stats_->num_tokens_clock_ += token_sweeped;

 process_barrier_->wait();
 double iter_end = lda::get_time();

 if (thread_id == 0)
 {
 double seconds_this_iter = iter_end - iter_start;

 printf("Iter: %04d", iter);
 std::cout
 << "\tThread = " << thread_id
 << "\tTokens: " << atomic_stats_->num_tokens_clock_
 << "\tTook: " << seconds_this_iter << " sec"
 << "\tThroughput: "
 << static_cast<double>(atomic_stats_->num_tokens_clock_) / (seconds_this_iter) << " token/(thread*sec)"
 << std::endl;
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "train(gibbs sampling) - in function train_thread");

 // syncup global table
 double sync_start = lda::get_time();
 for (int i = 0; i < num_threads_; ++i)
 {
 std::vector<word_topic_delta> & wtd_vec = samplers_[i]->get_word_topic_delta(thread_id);
 for (auto& wtd : wtd_vec)
 {
 global_word_topic_table_[wtd.word].inc(wtd.topic, wtd.delta);
 }
 }

 // use thread-private delta table to update global table
 {
 std::lock_guard<std::mutex> lock(atomic_stats_->global_mutex_);
 std::vector<int64_t> &summary = sampler_.get_delta_summary_row();
 for (int i = 0; i < K_; ++i)
 {
 global_summary_row_[i] += summary[i];
 }
 }
 process_barrier_->wait();
 CheckFunction(thread_id, tmDebug, "syncup global word_topic table - in function train_thread");

 if (compute_ll_interval_ != -1 && (iter % compute_ll_interval_ == 0 || iter == num_iterations_ - 1))
 {
 double ll = EvalLogLikelihood(true, thread_id, iter, sampler_);
 llcontainer.push_back(std::pair<int, double>(iter, ll));
 }

 CheckFunction(thread_id, tmDebug, "----------------------iteration end - in function train_thread---------------------");
 }

 if (thread_id == 0)
 {
 //output the ll once
 for (size_t i = 0; i < llcontainer.size(); i++)
 {
 printf("loglikelihood @iter%04d = %f\n", llcontainer[i].first, llcontainer[i].second);
 }
 }

 process_barrier_->wait();

 snprintf(tmDebug.m_szMessage, 200, "thread_id = %d, training iterations", thread_id);
 tmDebug.InnerTag();
 }

 void LdaEngine::Testing_Thread()
 {
 int thread_id = atomic_stats_->thread_counter_++;

 // Set core affinity which helps performance improvement
#ifdef _MSC_VER
 DWORD maskLL = 0;
 maskLL |= (1LL << (thread_id));
 DWORD_PTR mask = maskLL;
 SetThreadAffinityMask(GetCurrentThread(), mask);
#elif defined(__APPLE__)
 thread_port_t thread = pthread_mach_thread_np(pthread_self());
 thread_affinity_policy_data_t policy = { thread_id };
 thread_policy_set(thread, THREAD_AFFINITY_POLICY, (thread_policy_t)&policy, 1);
#else
 cpu_set_t set;
 CPU_ZERO(&set);
 CPU_SET(thread_id, &set);
 sched_setaffinity(0, sizeof(cpu_set_t), &set);
#endif
 process_barrier_->wait();

 //// Each thread builds a portion of word-topic table. We do this way because each word-topic row 
 //// has a thread-specific buffer for rehashing
 LightDocSampler &sampler_ = *(samplers_[thread_id]);
 sampler_.AdaptAlphaSum(false);

 double init_start = lda::get_time();
 int32_t token_num = 0;
 int32_t doc_start = data_block_->Begin(thread_id);
 int32_t doc_end = data_block_->End(thread_id);

 for (int32_t doc_index = doc_start; doc_index != doc_end; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 int doc_size = doc->size();
 for (int i = 0; i < doc_size; ++i)
 {
 int topic = sampler_.rand_k();
 doc->SetTopic(i, topic);
 }
 int cursor = doc->get_cursor();
 token_num += sampler_.GlobalInit(doc.get());
 }

 process_barrier_->wait();

 // build alias table
 // build alias table for the dense term, beta_k_v_, which is shared by all the words
 if (thread_id == 0)
 {
 beta_mass_ = 0;
 std::vector<float> proportion(K_);
 for (int k = 0; k < K_; ++k)
 {
 proportion[k] = beta_ / (global_summary_row_[k] + beta_sum_);
 beta_mass_ += proportion[k];
 }

 alias_rng_int_.SetProportionMass(proportion, beta_mass_, beta_k_v_, &beta_height_, sampler_.rng());
 }

 // build alias table for the sparse term
 double alias_start = lda::get_time();
 process_barrier_->wait();
 sampler_.build_alias_table(word_range_for_each_thread_[thread_id], word_range_for_each_thread_[thread_id + 1], thread_id);
 process_barrier_->wait();

 // print the log-likelihood before inference
 EvalLogLikelihood(true, thread_id, 0, sampler_);

 double total_start = lda::get_time();
 for (int iter = 0; iter < burnin_iterations_; ++iter)
 {
 double iter_start = lda::get_time();
 int32_t token_sweeped = 0;
 atomic_stats_->num_tokens_clock_ = 0;
 int32_t doc_start_local = data_block_->Begin(thread_id);
 int32_t doc_end_local = data_block_->End(thread_id);

 for (int32_t doc_index = doc_start_local; doc_index != doc_end_local; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 token_sweeped += sampler_.InferOneDoc(doc.get());
 }
 atomic_stats_->num_tokens_clock_ += token_sweeped;

 process_barrier_->wait();
 double iter_end = lda::get_time();

 if (thread_id == 0)
 {
 double seconds_this_iter = iter_end - iter_start;

 printf("Iter: %04d", iter);
 std::cout
 << "\tThread = " << thread_id
 << "\tTokens: " << atomic_stats_->num_tokens_clock_
 << "\tTook: " << seconds_this_iter << " sec"
 << "\tThroughput: "
 << static_cast<double>(atomic_stats_->num_tokens_clock_) / (seconds_this_iter) << " token/(thread*sec)"
 << std::endl;

 }

 process_barrier_->wait();

 if (compute_ll_interval_ != -1 && (iter % compute_ll_interval_ == 0 || iter == burnin_iterations_ - 1))
 {
 EvalLogLikelihood(false, thread_id, iter, sampler_);
 }
 }

 double total_end = lda::get_time();
 printf("thread_id = %d, Total time for burnin iterations : %f sec.\n", thread_id, total_end - total_start);
 }

 void LdaEngine::AllocateDataMemory(int num_document, int64_t corpus_size)
 {
 data_block_->Allocate(num_document, corpus_size);
 }

 void LdaEngine::AllocateModelMemory(const LDADataBlock* data_block)
 {
 model_block_->InitFromDataBlock(data_block, V_, K_);

 global_word_topic_table_.resize(V_);

 for (int i = 0; i < V_; ++i)
 {
 global_word_topic_table_[i] = model_block_->get_row(i, nullptr);
 }
 }

 void LdaEngine::AllocateModelMemory(int num_vocabs, int num_topics, int64_t nonzero_num)
 {
 model_block_->Init(num_vocabs, num_topics, nonzero_num);

 global_word_topic_table_.resize(num_vocabs);

 for (int i = 0; i < num_vocabs; ++i)
 {
 global_word_topic_table_[i] = model_block_->get_row(i, nullptr);
 }
 }

 void LdaEngine::AllocateModelMemory(int num_vocabs, int num_topics, int64_t mem_block_size, int64_t alias_mem_block_size)
 {
 model_block_->Init(num_vocabs, num_topics, mem_block_size, alias_mem_block_size); //memory allocated here

 global_word_topic_table_.resize(num_vocabs);
 global_summary_row_.resize(K_, 0);

 //each value inside the global_word_topic_table_ will be set while call SetWordTopic()
 }

 int LdaEngine::FeedInData(int* term_id, int* term_freq, int32_t term_num, int32_t vocab_size)
 {
 if (V_ == 0) //number vocab could be set in allocating model memory function
 V_ = vocab_size;

 //data_block represent for one doc
 return data_block_->Add(term_id, term_freq, term_num);
 }

 int LdaEngine::FeedInDataDense(int* term_freq, int32_t term_num, int32_t vocab_size)
 {
 if (V_ == 0) //number vocab could be set in allocating model memory function
 V_ = vocab_size;

 //data_block represent for one doc
 return data_block_->AddDense(term_freq, term_num);
 }

 void LdaEngine::TestOneDoc(int* term_id, int* term_freq, int32_t term_num, int* pTopics, int* pProbs, int32_t& numTopicsMax, int32_t numBurnIter, bool reset)
 {
 //numTopicsMax initialy holds the max returned topic number in order to hold the pTopic/pProbs memory in outside function
 //when data return, numTopicsMax should contains the real topic number returned.
 int sampler_id = 0;
 sampler_id = samplerQueue_->pop();

 LightDocSampler &sampler = *(samplers_[sampler_id]);
 int64_t data_length = 1;
 for (int i = 0; i < term_num; ++i)
 {
 for (int j = 0; j < term_freq[i]; ++j)
 {
 data_length += 2;
 }
 }

 assert(data_length <= maxDocToken_ * 2 + 1);

 if (reset)
 {
 // restart the rng seeds, so that we always get consistent result for the same input
 rng_.restart();
 sampler.rng_restart();
 }

 // NOTE: in multi-threaded implementation, the dynamic memory allocation
 // may cause contention at OS heap lock
 int64_t idx = 1;
 for (int i = 0; i < term_num; ++i)
 {
 for (int j = 0; j < term_freq[i]; ++j)
 {
 document_buffer_[sampler_id][idx++] = term_id[i];
 document_buffer_[sampler_id][idx++] = rng_.rand_k(K_);
 }
 }

 std::shared_ptr<LDADocument> doc(new LDADocument(document_buffer_[sampler_id], document_buffer_[sampler_id] + data_length));

 for (int iter = 0; iter < numBurnIter; ++iter)
 {
 sampler.InferOneDoc(doc.get());
 }
 sampler.GetDocTopic(doc.get(), pTopics, pProbs, numTopicsMax);

 samplerQueue_->push(sampler_id);
 }

 void LdaEngine::TestOneDocDense(int* term_freq, int32_t term_num, int* pTopics, int* pProbs, int32_t& numTopicsMax, int32_t numBurnIter, bool reset)
 {
 //numTopicsMax initialy holds the max returned topic number in order to hold the pTopic/pProbs memory in outside function
 //when data return, numTopicsMax should contains the real topic number returned.
 int sampler_id = 0;
 sampler_id = samplerQueue_->pop();

 LightDocSampler &sampler = *(samplers_[sampler_id]);
 int64_t data_length = 1;
 for (int i = 0; i < term_num; ++i)
 {
 for (int j = 0; j < term_freq[i]; ++j)
 {
 data_length += 2;
 }
 }

 assert(data_length <= maxDocToken_ * 2 + 1);

 if (reset)
 {
 // restart the rng seeds, so that we always get consistent result for the same input
 rng_.restart();
 sampler.rng_restart();
 }

 // NOTE: in multi-threaded implementation, the dynamic memory allocation
 // may cause contention at OS heap lock
 int64_t idx = 1;
 for (int i = 0; i < term_num; ++i)
 {
 for (int j = 0; j < term_freq[i]; ++j)
 {
 document_buffer_[sampler_id][idx++] = i;
 document_buffer_[sampler_id][idx++] = rng_.rand_k(K_);
 }
 }

 std::shared_ptr<LDADocument> doc(new LDADocument(document_buffer_[sampler_id], document_buffer_[sampler_id] + data_length));

 for (int iter = 0; iter < numBurnIter; ++iter)
 {
 sampler.InferOneDoc(doc.get());
 }
 sampler.GetDocTopic(doc.get(), pTopics, pProbs, numTopicsMax);

 samplerQueue_->push(sampler_id);
 }

 void LdaEngine::GetDocTopic(int docID, int* pTopic, int* pProb, int32_t& numTopicReturn)
 {
 //get the current topic vector of the document
 int thread_id = 0;
 LightDocSampler &sampler = *(samplers_[thread_id]);

 sampler.GetDocTopic(data_block_->GetOneDoc(docID).get(), pTopic, pProb, numTopicReturn);
 }

 void LdaEngine::SetAlphaSum(float avgDocLength)
 {
 if (!bAlphaSumMultiplied)
 {
 alpha_sum_ = alpha_sum_ * avgDocLength;
 bAlphaSumMultiplied = true;
 }
 printf("alpha_sum was set to %f", alpha_sum_);
 }

 bool LdaEngine::ClearData()
 {
 data_block_->Clear();
 return true;
 }

 bool LdaEngine::ClearModel()
 {
 model_block_->Clear();
 return true;
 }

 //function to support dumping the topic_model model file
 void LdaEngine::GetWordTopic(int32_t wordId, int32_t* pTopic, int32_t* pProb, int32_t& length)
 {
 //cap the topic number here according to inpassed value of length
 int lengthCap = length;

 // NOTE: we MUST check whether the word-topic row is empty before get its value
 if (global_word_topic_table_[wordId].capacity() == 0)
 {
 length = 0;
 return;
 }

 length = 0;
 for (int i = 0; i < K_; ++i)
 {
 if (global_word_topic_table_[wordId][i] > 0)
 {
 pTopic[length] = i;
 pProb[length] = global_word_topic_table_[wordId][i];
 length++;

 if (length >= lengthCap)
 break;
 }
 }
 }

 // Compare by frequencies in descending order.
 bool CompareTerms(const std::pair<int, int> &term1, const std::pair<int, int> &term2)
 {
 // REVIEW: consider changing this to impose a total order, since quicksort is not stable.
 return term2.second < term1.second;
 }

 void LdaEngine::GetTopicSummary(int32_t topicId, int32_t* pWords, float* pProb, int32_t& length)
 {
 std::vector<std::pair<int, int>> allTermsVec;
 int sumFreq = 0;
 for (int i = 0; i < V_; i++) //for all the terms check the topic distribution
 {
 if (global_word_topic_table_[i][topicId] > 0)
 {
 std::pair<int, int> p;
 p.first = i;
 p.second = global_word_topic_table_[i][topicId];
 allTermsVec.push_back(p);
 sumFreq += global_word_topic_table_[i][topicId];
 }
 }

 std::sort(allTermsVec.begin(), allTermsVec.end(), CompareTerms);

 int usedTerm = (int)allTermsVec.size();
 length = std::min(usedTerm, length);
 for (int i = 0; i < length; i++)
 {
 pWords[i] = allTermsVec[i].first;
 pProb[i] = (((float)(allTermsVec[i].second)) + beta_) / (sumFreq + beta_ * V_);
 }
 }

 //function to support loading the topic_model model file
 void LdaEngine::SetWordTopic(int32_t wordId, int32_t* pTopic, int32_t* pProb, int32_t length)
 {
 //NOTE: whether we should really use the "true" here
 model_block_->SetWordInfo(wordId, length, true);
 global_word_topic_table_[wordId] = model_block_->get_row(wordId, nullptr);

 for (int i = 0; i < length; ++i)
 {
 global_word_topic_table_[wordId].inc(pTopic[i], pProb[i]);
 global_summary_row_[pTopic[i]] += pProb[i];
 }
 }

 void LdaEngine::GetModelStat(int64_t &memBlockSize, int64_t &aliasMemBlockSize)
 {
 //NOTE: get the model's value at the end of training stage. try to save these two numbers to disk file
 model_block_->GetModelStat(memBlockSize, aliasMemBlockSize);
 }

 double LdaEngine::EvalLogLikelihood(bool is_train, int thread_id, int iter, LightDocSampler &sampler)
 {
 double doc_ll = 0;
 double word_ll = 0;

 if (thread_id == 0)
 {
 atomic_stats_->doc_ll_ = 0;
 atomic_stats_->word_ll_ = 0;
 }
 process_barrier_->wait();

 int doc_num = 0;
 int32_t doc_start = data_block_->Begin(thread_id);
 int32_t doc_end = data_block_->End(thread_id);
 for (int32_t doc_index = doc_start; doc_index != doc_end; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 doc_ll += sampler.ComputeOneDocLLH(doc.get());
 doc_num++;
 }
 atomic_stats_->doc_ll_ = atomic_stats_->doc_ll_ + doc_ll;
 process_barrier_->wait();

 word_ll = sampler.ComputeWordLLH(word_range_for_each_thread_[thread_id], word_range_for_each_thread_[thread_id + 1]);
 atomic_stats_->word_ll_ = atomic_stats_->word_ll_ + word_ll;
 process_barrier_->wait();

 double total_ll = 0;
 if (thread_id == 0)
 {
 double normalized_ll = sampler.NormalizeWordLLH();

 total_ll = 0;
 total_ll += atomic_stats_->doc_ll_;
 total_ll += atomic_stats_->word_ll_;
 total_ll += normalized_ll;

 if (!is_train)
 {
 likelihood_in_iter_[iter] = (float)total_ll;
 }

 std::cout << "Total likelihood: " << total_ll << "\t";
 std::cout << "..........[Nomralized word ll: " << normalized_ll << "\t"
 << "Word likelihood: " << atomic_stats_->word_ll_ << "\t"
 << "Doc likelihood: " << atomic_stats_->doc_ll_ << "]" << std::endl;
 }
 process_barrier_->wait();

 return total_ll;
 }

 void LdaEngine::DumpDocTopicTable(const std::string& doc_topic_file)
 {
 std::ofstream dt_stream;
 dt_stream.open(doc_topic_file, std::ios::out);
 assert(dt_stream.good());

 int32_t num_documents = data_block_->num_documents();
 int32_t doc_start = 0;
 int32_t doc_end = num_documents;

 lda::light_hash_map doc_topic_counter_(1024);

 for (int32_t doc_index = doc_start; doc_index != doc_end; ++doc_index)
 {
 std::shared_ptr<LDADocument> doc = data_block_->GetOneDoc(doc_index);
 doc_topic_counter_.clear();
 doc->GetDocTopicCounter(doc_topic_counter_);

 dt_stream << doc_index;
 if (doc->size())
 {
 int32_t capacity = doc_topic_counter_.capacity();
 int32_t *key = doc_topic_counter_.key();
 int32_t *value = doc_topic_counter_.value();
 int32_t nonzero_num = 0;

 for (int i = 0; i < capacity; ++i)
 {
 if (key[i] > 0)
 {
 dt_stream << " " << key[i] - 1 << ":" << value[i];
 }
 }
 }
 dt_stream << std::endl;
 }
 dt_stream.close();
 }

 void LdaEngine::DumpFullModel(const std::string& word_topic_dump)
 {
 std::ofstream wt_stream;
 wt_stream.open(word_topic_dump, std::ios::out);
 assert(wt_stream.good());

 for (int w = 0; w < V_; ++w)
 {
 int nonzero_num = global_word_topic_table_[w].nonzero_num();
 if (nonzero_num)
 {
 wt_stream << w;
 for (int t = 0; t < K_; ++t)
 {
 if (global_word_topic_table_[w][t] > 0)
 {
 wt_stream << " " << t << ":" << global_word_topic_table_[w][t];
 }
 }
 wt_stream << std::endl;
 }
 }
 wt_stream.close();

 std::ofstream summary_stream;
 summary_stream.open("summary_row.txt", std::ios::out);
 for (int i = 0; i < K_; ++i)
 {
 summary_stream << global_summary_row_[i] << std::endl;
 }
 summary_stream.close();
 }
} // namespace lda