#include "cnn/nodes.h"
#include "cnn/cnn.h"
#include "cnn/training.h"
#include "cnn/timing.h"
#include "cnn/rnn.h"
#include "cnn/gru.h"
#include "cnn/lstm.h"
#include "cnn/dict.h"
#include "cnn/expr.h"

#include <iostream>
#include <fstream>

#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>

#include <unordered_map>
#include <unordered_set>

using namespace std;
using namespace cnn;

cnn::real pdrop = 0.5;
bool eval = false;
bool verbose = false;
unsigned INPUT_DIM = 64;
unsigned HIDDEN_DIM = 150;

unsigned INPUT_VOCAB_SIZE = 0;

cnn::Dict word_dict;

int kSOS;
int kEOS;
int kUNK;
int kNUM; //新加的用来表示数字

unordered_map<string, vector<cnn::real> > pretrained_embeddings;
vector<cnn::real> averaged_vec;

cnn::real epsilon = 0.001;

template<class T1, class T2>
struct PairCmp
{
    bool operator()(const pair<T1, T2>& a, const pair<T1, T2>& b)
    {
        return a.second > b.second;
    }
};
class Example {
public:
    vector<vector<string> > words; //multiple sents
    vector<vector<int> > word_ids; //sent_ids
    float ylabel;
    float fraction;
    float open;
    float close;
    string date;
};

bool ReadExample(ifstream& fin, Example& example);

/**
* Read Word Embeddings and normalize them
* @param input_file pretrained embeddings file path
* */

void ReadEmbeddings(const string& input_file)
{
    ifstream fin(input_file);
    assert(fin);
    string line;
    bool first = true;
    while(getline(fin, line))
    {
        istringstream sin(line);
        string word;
        sin>>word;
        cnn::real value;
        vector<cnn::real> vecs;
        cnn::real sum = 0.000001;
        while(sin>>value)
        {
            vecs.push_back(value);
            sum += value * value;
        }
        if(first)
        {
            first = false;
            INPUT_DIM = vecs.size();
            averaged_vec.resize(INPUT_DIM, 0.0);
        }
        //normalize to one
        sum = sqrt(sum);
        for(int j = 0; j < INPUT_DIM; j++)
        {
            averaged_vec[j] += vecs[j];
            vecs[j] /= sum;
        }
        pretrained_embeddings[word] = vecs;
        word_dict.Convert(word);
    }
    fin.close();
    cnn::real sum = 0.000001;
    for(int j = 0; j < INPUT_DIM; j++)
    {
        averaged_vec[j] /= pretrained_embeddings.size();
        sum += averaged_vec[j] * averaged_vec[j];
    }
    sum = sqrt(sum);
    for(int j = 0; j < INPUT_DIM; j++)
    {
        averaged_vec[j] /= sum ;
    }
    cerr<<"Pretrained Embeddings: "<< pretrained_embeddings.size()<<endl;
}

void ReadCorpus(const string& input_file,
                vector<Example>& examples,
                cnn::Dict& wordDict,
                int& tlines,
                int& ttoks,
                bool isTrain = true,
                int cut_off = 3)
{
    ifstream fin(input_file);
    if(!fin) {
        cerr<<"Cannot open file: "<< input_file << endl;
        exit(0);
    }

    unordered_map<string, int> word_counts;
    int tokens = 0, types = 0;
    while(true)
    {
        Example example;
        if(!ReadExample(fin,  example)) break;
        tlines++;
        for(int t = 0; t< (int)example.words.size(); t++)
        {
            auto& title = example.words[t];
            if(isTrain)
            {
                for(int j = 0; j < (int)title.size(); j++)
                {
                    auto& word = title[j];
                    auto iter = word_counts.find(word);
                    if( iter == word_counts.end())
                    {
                        word_counts.insert(make_pair(word, 1));
                        types++;
                    }
                    else {
                        iter->second++;
                    }
                    tokens++;
                }
            }
            ttoks += title.size();
        }
        examples.push_back(example);
    }
    fin.close();

    if(isTrain)
    {
        cerr<<"[LOG] Before Pruning there are "<<  tokens << " words and " <<types<< " types."<<endl;
        std::vector<pair<string, float>> pairs;

        unordered_map<string, float> tf_counts;
        unordered_map<string, float> idf_counts;
        int total_titles = 0;
        for(unsigned sent_index = 0; sent_index < examples.size(); sent_index++)
        {
            auto& example = examples[sent_index];
            int num_titles = example.words.size();
            total_titles += num_titles;
            for(int title_index = 0; title_index < num_titles; title_index++)
            {
                unordered_map<string, float> temp_idf_counts;
                for(unsigned word_index = 0; word_index < example.words[title_index].size(); word_index++)
                {
                    const string& word = example.words[title_index][word_index];
                    tf_counts[word]++;
                    temp_idf_counts[word] = 1;
                }
                for(auto& item: temp_idf_counts)
                {
                    idf_counts[item.first]++;
                }
            }
        }
        for(auto& item: tf_counts)
        {
            float tf = item.second;
            float idf = idf_counts[item.first];
            //float tf_idf = (1.0 + std::log(tf)) * std::log( example.words.size()/idf );
            float tf_idf = (1.0 + std::log(tf)) * std::log( total_titles/idf );
            pairs.push_back(make_pair(item.first, tf_idf));
        }
        sort(pairs.begin(), pairs.end(), PairCmp<string, float>());
        /*for(auto p: pairs)
        {
            //cerr<< p.first << "\t" << p.second  << "\t" << tf_counts[p.first] << "\t" << idf_counts[p.first] << "\t" << example.words.size() << endl;
            cerr<< p.first << "\t" << p.second  << "\t" << tf_counts[p.first] << "\t" << idf_counts[p.first] << "\t" << total_titles << endl;
        }*/
        //int rank = pairs.size() * 0.4; //save 40% words
        int rank = pairs.size() * 0.4; //save 40% words
        float survive_value  =  pairs[rank].second;
        int remaining_tokens = 0, remaining_types = 0;
        for(auto iter = pairs.begin(); iter != pairs.end(); iter++)
        {
            if(iter->second > survive_value)
                //if(iter->second >= survive_value)
            {
                word_dict.Convert(iter->first);
                remaining_tokens += iter->second;
                remaining_types++;
            }
        }
        cerr<<"[LOG] After Pruning there are "<<  remaining_tokens << " words and " << remaining_types<< " types."<<endl;
        cerr<<"[LOG] WordDict Size: "<< word_dict.size() << endl;
        word_dict.Freeze();
        word_dict.SetUnk("|unk|");
    }

    for(unsigned sent_index = 0; sent_index < examples.size(); sent_index++)
    {
        auto& example = examples[sent_index];
        int num_titles = example.words.size();
        example.word_ids.resize(num_titles);

        for(int title_index = 0; title_index < num_titles; title_index++)
        {
            example.word_ids[title_index].resize( example.words[title_index].size() );

            for(unsigned word_index = 0; word_index < example.words[title_index].size(); word_index++)
            {
                int wordId = wordDict.Convert( example.words[title_index][word_index] );
                example.word_ids[title_index][word_index] = wordId;
                //cerr<< example.words[title_index][word_index]<<" ";
            }
            //cerr<<endl;
        }
        //cerr<< endl;
    }
    //cerr<<endl;
}

void computeLRP(cnn::real Rf, const vector<cnn::real>& w, const vector<cnn::real>& s, cnn::real b, vector<cnn::real>& Rs)
{
    //implement equation (2) to equation (6) of alex's note
    assert(w.size() == s.size());
    cnn::real sum = b;
    for(unsigned k = 0; k< s.size(); k++)
        sum += w[k] * s[k];
    cnn::real append = epsilon * (sum > 0? 1: -1);
    sum += append;
    Rs.resize(s.size(), 0);
    for(unsigned k = 0; k< s.size(); k++)
        Rs[k] = Rf * w[k] * s[k] / sum;
}
void computeLRP(cnn::real Rf, const vector<cnn::real>& w, const vector<cnn::real>& s, cnn::real b, vector<cnn::real>& Rs, cnn::real& Rb)
{
    //implement equation (2) to equation (6) of alex's note
    assert(w.size() == s.size());
    cnn::real sum = b;
    for(unsigned k = 0; k< s.size(); k++)
        sum += w[k] * s[k];
    cnn::real append = epsilon * (sum > 0? 1: -1);
    sum += append;
    Rs.resize(s.size(), 0);
    for(unsigned k = 0; k< s.size(); k++)
        Rs[k] = Rf * w[k] * s[k] / sum;
    Rb = Rf * b / sum;
}
void computeLRPforSum(float Rf, const vector<cnn::real>& s, vector<cnn::real>& Rs)
{
    Rs.resize(s.size(), 0);
    cnn::real sum = 0;
    for(unsigned k = 0; k< s.size(); k++)
        sum += s[k];
    cnn::real append = epsilon * (sum > 0? 1: -1);
    sum += append;
    for(unsigned k = 0; k< s.size(); k++)
        Rs[k] = Rf  * s[k] / sum;
}

struct TitleEncoder {
    LookupParameters* p_word;  //

    vector<Parameters*> p_output_w;
    vector<Parameters*> p_output_b;

    vector<Parameters*> p_res_w;

    Parameters* p_fed_o;
    Parameters* p_fed_o_b;

    int layers;
    explicit TitleEncoder(Model& model): layers(1)
    {
        p_word = model.add_lookup_parameters(INPUT_VOCAB_SIZE, {INPUT_DIM});

        p_output_w.resize(layers);
        p_output_b.resize(layers);
        p_res_w.resize(layers);
        for(int i = 0; i< layers; i++)
        {
            p_output_w[i] = model.add_parameters({HIDDEN_DIM, INPUT_DIM});
            p_output_b[i] = model.add_parameters({HIDDEN_DIM});
            p_res_w[i] = model.add_parameters({INPUT_DIM, HIDDEN_DIM});
        }

        p_fed_o = model.add_parameters({1, INPUT_DIM});
        p_fed_o_b = model.add_parameters({1}, 0.01);

        for( size_t  word_index = 0; word_index <  INPUT_VOCAB_SIZE; word_index++)
        {
            auto word_iter = pretrained_embeddings.find( word_dict.Convert(word_index) );
            if(word_iter != pretrained_embeddings.end())
            {
                ((LookupParameters*)p_word)->Initialize(word_index, word_iter->second);
            }
        }
        ((LookupParameters*)p_word)->Initialize(kUNK, averaged_vec);
        ((LookupParameters*)p_word)->Initialize(kNUM, averaged_vec);
    }

    Expression buildLRPGraph(ComputationGraph&cg, const Example& example)
    {
        assert(eval);
        vector<Expression> title_exps;
        vector<vector<Expression> > title_words;
        vector<Expression> title_cols;

        for(size_t title_index = 0; title_index < example.words.size(); title_index++)
        {
            auto& insent = example.word_ids[title_index];

            const unsigned slen = insent.size();

            vector<Expression> i_words;

            for (unsigned t = 0; t < insent.size(); ++t)
            {
                Expression i_word_t;

                i_word_t = lookup(cg, p_word, insent[t]);
                if (!eval) {
                    i_word_t = dropout(i_word_t, pdrop);
                }
                if(insent[t] != kUNK)//this way better
                {
                    i_words.push_back(i_word_t);
                }
            }
            if(i_words.size() > 1 )
            {
                title_words.push_back(i_words);
                //关键是这个地方
                Expression current_col = tanh(kmh_ngram(concatenate_cols(i_words), 2));
                title_exps.push_back(sum_cols(current_col) );
                title_cols.push_back(current_col);
            }
            else if(i_words.size()> 0) {
                //这个地方用得到每个words
                title_exps.push_back(average(i_words));
                title_words.push_back(i_words);
            }
        }
        if(!eval)
        {
            for(size_t title_index = 0; title_index < title_exps.size(); title_index++)
            {
                title_exps[title_index] = block_dropout(title_exps[title_index], pdrop);
            }
        }
        //这个地方用得到每个title_exps的
        Expression example_rep = average(title_exps);
        vector<Expression> param_tW(layers), param_tb(layers), param_resW(layers);
        for(int i = 0; i< layers; i++)
        {
            param_tW[i] = parameter(cg, p_output_w[i]);
            param_tb[i] = parameter(cg, p_output_b[i]);
            param_resW[i] = parameter(cg, p_res_w[i]);
        }
        //Dropout可以不用管
        if(!eval)
            example_rep = dropout(example_rep, pdrop);
        vector<Expression> example_reps, example_rep_temps;
        example_reps.push_back(example_rep);

        for(int i = 0; i < layers; i++)
        {
            Expression example_rep_temp = tanh(param_tW[i] * example_rep + param_tb[i]);
            example_rep_temps.push_back(example_rep_temp);

            example_rep = param_resW[i] * example_rep_temp + example_rep;
            example_reps.push_back(example_rep);
        }
        Expression i_output = example_rep;
        //得到example_rep的相关度也没有问题了
        if(verbose)
        {
            cg.incremental_forward();
            cerr<<(*i_output.value())<<endl;
        }

        Expression param_oW = parameter(cg, p_fed_o);
        Expression param_ob = parameter(cg, p_fed_o_b);
        Expression output_h = param_oW * i_output + param_ob;
        //===============LRP==============================
        //Step 1: Get Rf
        auto feature_vector  = as_scalar(cg.incremental_forward());
        cnn::real final_r_f = feature_vector; //预测涨的相关性, 如果final_r_f 大于0， 那么就预测涨
        //cerr<<"final_f: " << final_r_f  << "\t" << feature_vector[1]<< endl;

        //Step 2: Get R_i_output
        auto param_ob_value = as_scalar(param_ob.value());
        auto param_oW_values = *(param_oW.value());
        auto i_output_values = as_vector(i_output.value());
        /*
        cnn::real sum = b_values[1];
        for(int i = 0; i< HIDDEN_DIM; i++)
            sum += param_oW_values(1, i) * i_output_values[i];
        cerr<<"comp: " << sum << "\t" << feature_vector[1] << endl; //几乎相等了
        */
        vector<cnn::real> w(INPUT_DIM);
        for(unsigned i = 0; i< INPUT_DIM; i++)
            w[i] = param_oW_values(0,i);
        vector<cnn::real> R_i_output;
        computeLRP(final_r_f, w, i_output_values, param_ob_value, R_i_output);
        /*copy(R_i_output.begin(), R_i_output.end(), ostream_iterator<cnn::real>(cerr, " "));
        cerr<< endl;*/
        //Step 2.5; Get R_example_rep final
        auto example_rep_final_values = as_vector(example_reps.back().value());

        auto example_rep_input_values = as_vector(example_reps[0].value());
        auto example_rep_temp_values = as_vector(example_rep_temps.back().value());
        auto param_resW_values = *(param_resW[0].value());

        vector<cnn::real> R_example_rep(INPUT_DIM, 0);
        vector<cnn::real> R_example_rep_tanh(HIDDEN_DIM, 0);
        for(unsigned i = 0; i< INPUT_DIM; i++)
        {
            vector<cnn::real> temp_R_example_rep;
            vector<cnn::real> sw(HIDDEN_DIM);
            for(unsigned j = 0; j< HIDDEN_DIM; j++)
                sw[j] = param_resW_values(i,j);
            //also need to compute b
            computeLRP(R_i_output[i], sw, example_rep_temp_values, example_rep_input_values[i], temp_R_example_rep, R_example_rep[i]);
            for(unsigned j = 0; j< HIDDEN_DIM; j++)
                R_example_rep_tanh[j] += temp_R_example_rep[j];
        }
        //Step 3: Accumulate R_example_rep input
        auto param_tW_values = *(param_tW[0].value());
        auto param_tb_values = as_vector(param_tb[0].value());
        for(unsigned i = 0; i< HIDDEN_DIM; i++)
        {
            vector<cnn::real> temp_R_example_rep;
            vector<cnn::real> sw(INPUT_DIM);
            for(unsigned j = 0; j< INPUT_DIM; j++)
                sw[j] = param_tW_values(i,j);
            computeLRP(R_example_rep_tanh[i], sw, example_rep_input_values, param_tb_values[i], temp_R_example_rep);
            for(unsigned j = 0; j< INPUT_DIM; j++)
                R_example_rep[j] += temp_R_example_rep[j];
        }
        /*copy(R_example_rep.begin(), R_example_rep.end(), ostream_iterator<cnn::real>(cerr, " "));
        cerr<< endl;*/
        //Step 4: Get R_title_exps, using Equations (12) to Equations (16)
        //Expression example_rep = average(title_exps);
        vector<vector<cnn::real> > title_exps_values(title_exps.size()), R_title_exps(title_exps.size());
        for(unsigned i = 0; i< title_exps.size(); i++)
        {
            title_exps_values[i] = as_vector(title_exps[i].value());
            R_title_exps[i].resize(INPUT_DIM, 0.0);
        }
        for(unsigned i = 0; i< INPUT_DIM; i++)
        {
            vector<cnn::real> s(title_exps.size());
            for(unsigned j = 0; j< title_exps.size(); j++)
                s[j] = title_exps_values[j][i];
            vector<cnn::real> temp_R_s(title_exps.size());
            computeLRPforSum(R_example_rep[i], s, temp_R_s);
            for(unsigned j = 0; j< title_exps.size(); j++)
                R_title_exps[j][i] += temp_R_s[j];
        }
        /*for(unsigned i = 0; i< title_exps.size(); i++) //这个地方把标题的每一维度加起来可以实现得到每个标题的相关性
        {
            copy(R_title_exps[i].begin(), R_title_exps[i].end(), ostream_iterator<cnn::real>(cerr, " "));
            cerr<< endl;
        }*/
        //Step 5: Get R_i_words
        size_t real_title_index = 0, title_col_index = 0;
        unordered_map<unsigned, vector<cnn::real> > R_words;
        for(size_t title_index = 0; title_index < example.words.size(); title_index++)
        {
            auto& insent = example.word_ids[title_index];

            vector<unsigned> selected_words; //选中的词

            for (unsigned t = 0; t < insent.size(); ++t)
            {
                if(insent[t] != kUNK)
                    selected_words.push_back(insent[t]);
            }
            if(selected_words.size() > 1 )
            {
                vector<Expression>& i_words  = title_words[real_title_index];
                Expression& current_col = title_cols[title_col_index];
                //title_exps.push_back( sum_cols(tanh(kmh_ngram(concatenate_cols(i_words), 2))) );
                //title_exps -> sum_cols
                //title_exps[real_title_index][i] = sum_j(current_col[i,j])
                auto current_col_values = *(current_col.value());
                assert(current_col_values.cols() == i_words.size() -1);
                vector<vector<cnn::real> > R_current_col(INPUT_DIM);
                unsigned num_cols = current_col_values.cols(); //列数
                for(size_t row_index = 0; row_index < INPUT_DIM; row_index++)
                {
                    vector<cnn::real> s(num_cols);
                    for(unsigned col_index = 0; col_index < num_cols ; col_index++)
                        s[col_index] = current_col_values(row_index, col_index);
                    vector<cnn::real> temp_R_s(num_cols);
                    computeLRPforSum(R_title_exps[real_title_index][row_index], s, temp_R_s);
                    R_current_col[row_index] = temp_R_s;
                }
                //sum->kmh_ngram->i_words
                vector<vector<cnn::real> > i_words_values(i_words.size()), R_i_words(i_words.size());
                for(unsigned i = 0; i< i_words.size(); i++)
                {
                    i_words_values[i] = as_vector(i_words[i].value());
                    R_i_words[i].resize(INPUT_DIM, 0.0);
                }
                unsigned ngram = 2;
                for(size_t row_index = 0; row_index < INPUT_DIM; row_index++)
                {
                    for(size_t col_index = 0; col_index < num_cols-ngram+1; col_index++)
                    {
                        vector<cnn::real> s(ngram);
                        for(unsigned ngram_index = 0; ngram_index < ngram; ngram_index++)
                            s[ngram_index] = i_words_values[col_index + ngram_index][row_index];
                        vector<cnn::real> temp_R_s(ngram);
                        computeLRPforSum(R_current_col[row_index][col_index], s, temp_R_s);
                        for(unsigned ngram_index = 0; ngram_index < ngram; ngram_index++)
                            R_i_words[col_index+ngram_index][row_index] += temp_R_s[ngram_index];
                    }
                }
                //加到每个词的身上
                for(unsigned i = 0; i< R_i_words.size(); i++)
                {
                    unsigned word = selected_words[i];
                    auto iter = R_words.find(word);
                    if(iter  == R_words.end())
                    {
                        R_words[word] = R_i_words[i];
                    }
                    else {
                        for(unsigned j = 0; j<INPUT_DIM; j++)
                            iter->second[j] += R_i_words[i][j];
                    }
                }

                real_title_index ++;
                title_col_index++;
            }
            else if(selected_words.size()> 0) {
                //title_exps.push_back(average(i_words));
                vector<Expression>& i_words  = title_words[real_title_index];
                vector<vector<cnn::real> > i_words_values(i_words.size()), R_i_words(i_words.size());
                for(unsigned i = 0; i< i_words.size(); i++)
                {
                    i_words_values[i] = as_vector(i_words[i].value());
                    R_i_words[i].resize(INPUT_DIM, 0.0);
                }
                for(unsigned i = 0; i< INPUT_DIM; i++)
                {
                    vector<cnn::real> s(i_words.size());
                    for(unsigned j = 0; j< i_words.size(); j++)
                        s[j] = i_words_values[j][i];
                    vector<cnn::real> temp_R_s(i_words.size());
                    computeLRPforSum(R_title_exps[real_title_index][i], s, temp_R_s);
                    for(unsigned j = 0; j< i_words.size(); j++)
                        R_i_words[j][i] += temp_R_s[j];
                }
                //加到每个词的身上
                for(unsigned i = 0; i< R_i_words.size(); i++)
                {
                    unsigned word = selected_words[i];
                    auto iter = R_words.find(word);
                    if(iter  == R_words.end())
                    {
                        R_words[word] = R_i_words[i];
                    }
                    else {
                        for(unsigned j = 0; j<INPUT_DIM; j++)
                            iter->second[j] += R_i_words[i][j];
                    }
                }
                real_title_index ++;
            }
        }
        //6, Get sum of R_words for each word according to equation 17
        vector<pair<unsigned, cnn::real> > R_sum_words;
        for(auto& item: R_words)
        {
            auto& scores = item.second;
            cnn::real sum_score = 0;
            for(auto& score: scores)
                sum_score += score;
            R_sum_words.push_back(make_pair(item.first, sum_score));
        }
        sort(R_sum_words.begin(), R_sum_words.end(), PairCmp<unsigned, cnn::real>());
        for(auto& item: R_sum_words)
        {
            cerr<< word_dict.Convert(item.first) << "\t" << item.second << endl;
        }
        cerr<< "===" << endl;
        return output_h;
    }
};

bool ReadExample(ifstream& fin, Example& example)
{
    string line;
    float open_price, close_price, fraction;
    int num_titles;
    string date;
    if(!getline(fin, line)) return false;
    istringstream sin(line);
    sin>>fraction>>num_titles>>open_price>>close_price>>date;
    //cout<<label_str<<", "<<num_titles<<endl;
    example.ylabel = fraction * 100;
    example.fraction = fraction;
    example.open = open_price;
    example.close = close_price;
    example.date = date;
    example.words.resize(num_titles);

    for(int i = 0; i< num_titles; i++)
    {
        if(!getline(fin, line))
            return false;
        istringstream wsin(line);
        string word;
        while(wsin>>word)
        {
            if(word == "|||") break;
            example.words[i].push_back(word);
        }
    }
    return true;
    //copy(data.begin(), data.end(), ostream_iterator<cnn::real>(cout, " "));
    //cout<<endl;
}

int main(int argc, char** argv) {
    cerr<<"Command: ";
    for(int i = 0; i< argc; i++)
        cerr<<argv[i]<<" ";
    cerr<<endl;

    cnn::Initialize(argc, argv, 4);

    if (argc != 5 && argc != 6) {
        cerr << "Usage: " << argv[0] << " code vec train dev [model.params]\n";
        return 1;
    }
    kSOS = word_dict.Convert("<s>");
    kEOS = word_dict.Convert("</s>");
    kUNK = word_dict.Convert("|unk|");
    kNUM = word_dict.Convert("<num>"); //number

    string stock_code = argv[1];
    vector<Example> training_examples, dev_examples;

    int tlc = 0;
    int ttoks = 0;
    cerr << "Reading Pretrained Word Vectors ... " << argv[2] << endl;
    ReadEmbeddings(argv[2]);

    cerr << "Reading training data from " << argv[3] << "...\n";

    ReadCorpus(argv[3], training_examples, word_dict, tlc, ttoks, true, 1);
    cerr << tlc << " lines, " << ttoks << " tokens, " << word_dict.size() << " types\n";
    int pos_count = 0;
    for(auto& example: training_examples)
        pos_count+= (example.ylabel >= 0);
    cerr<< "Total Training Examples: " << training_examples.size() << endl;
    cerr<< "Positive Training Examples: " << pos_count << " " << pos_count/(float)training_examples.size() << endl;
    cerr<< "Negative Training Examples: " << training_examples.size() - pos_count << " "<< 1.0 - pos_count /(float)training_examples.size() << endl;
    INPUT_VOCAB_SIZE = word_dict.size();

    int dlc = 0;
    int dtoks = 0;
    cerr << "Reading dev data from " << argv[4] << "...\n";
    ReadCorpus(argv[4], dev_examples, word_dict, dlc, dtoks, false, 0);
    cerr << dlc << " lines, " << dtoks << " tokens\n";
    int pos_dev_count = 0;
    for(auto& example: dev_examples)
        pos_dev_count+= (example.ylabel >=0);
    cerr<< "Total Dev Examples: " << dev_examples.size() << endl;
    cerr<< "Positive DEV Examples: " << pos_dev_count << " " << pos_dev_count/(float)dev_examples.size() << endl;
    cerr<< "Negative DEV Examples: " << dev_examples.size() - pos_dev_count << " "<< 1.0 - pos_dev_count /(float)dev_examples.size() << endl;

    ostringstream os;
    os << stock_code<<".hier.avg"
       << '_' << INPUT_DIM
       << '_' << HIDDEN_DIM
       << '_' << pdrop << "_pid." << getpid();

    const string fname = os.str();
    cerr << "Parameters will be written to: " << fname << endl;

    Model model;
    bool use_momentum = true;
    Trainer* sgd = nullptr;
    if (use_momentum)
    {
        //sgd = new MomentumSGDTrainer(&model);
        sgd = new MomentumSGDTrainer(&model, 1e-4, 0.01);
        sgd->eta_decay = 0.05;
        //sgd = new AdagradTrainer(&model);
    }
    else
        sgd = new SimpleSGDTrainer(&model, 1e-4);
    //sgd = new SimpleSGDTrainer(&model);
    //sgd->clipping_enabled = false;

    TitleEncoder titleEncoder(model);

    if(argc == 6)
    {
        string fname = argv[5];
        ifstream in(fname);
        boost::archive::text_iarchive ia(in);
        ia >> model;

        cnn::real dloss = 0;
        eval = true;
        //verbose = true;
        int dev_index = 0;
        for (auto& sent : dev_examples) {
            dev_index ++;
            auto y = sent.ylabel;
            ComputationGraph cg;
            //cerr<< "begin " << dev_index << endl;
            Expression y_pred =  titleEncoder.buildLRPGraph(cg, sent);
            //cerr << "DEVLINE: " << dtags << endl;
            dloss += as_scalar(cg.incremental_forward());
        }
        return 0;
    }

    delete sgd;
}


/* rosrun kamu_robotu_map_analyzer map_analyzer /
home/alperen/MY_WORKSPACES/oko_slam/ros_ws/saved_maps/real_results/cartographer_maps/4Objects128pt3 */

// ROS
#include <ros/ros.h>
#include <std_msgs/String.h>
// OpenCV
#include <opencv2/opencv.hpp>
//#include <opencv2/highgui/highgui.hpp>

// File system and regular C++ headers
#include <string>
#include <iomanip>
#include <sstream>
#include <vector>

#include <boost/filesystem.hpp>
#include <dirent.h>
#include <sys/stat.h>
// GetCurrentDir
#include <unistd.h>

#include <yaml-cpp/yaml.h>

#define GetCurrentDir getcwd
// Function Prototypes
std::string GetCurrentWorkingDir(void);

int main(int argc, char** argv)
{
    ros::init(argc, argv, "map_analyzer");
    ros::NodeHandle nh("~");

    // Get parameter from user
    std::string file_dir = argv[1];
    std::string yaml_file = file_dir + ".yaml";
    std::string map_file = file_dir + ".pgm";

    // Read Original Map Image
    cv::Mat input_map_image, input_map_image_gs;
    input_map_image = cv::imread(map_file);

    // Read yaml files
    YAML::Node config = YAML::LoadFile(yaml_file);
    float resolution = config["resolution"].as<float>();
    std::vector<float> origin_yaml = config["origin"].as<std::vector<float>>();
    int offset_coordinate_axis = input_map_image.size().height;
    int robot_initial_point_x = abs( origin_yaml[0]/resolution );
    int robot_initial_point_y = offset_coordinate_axis - abs( origin_yaml[1]/resolution );

    std::cout << "robot initial pose x: " << robot_initial_point_x << std::endl;
    std::cout << "robot initial pose y: " << robot_initial_point_y << std::endl;

    // Draw a Small Circle to Specify Initial Starting Point of Robot
    cv::circle( input_map_image, cv::Point(robot_initial_point_x,robot_initial_point_y), 1.0, cv::Scalar(0,0,255), -1, 8, 0 );
    // Draw Coordinate Axis whose origin is robot initial position
    cv::line( input_map_image, cv::Point(robot_initial_point_x,robot_initial_point_y) , cv::Point(robot_initial_point_x+10,robot_initial_point_y), cv::Scalar(0,0,255), 1,8); //X axis
    cv::line( input_map_image, cv::Point(robot_initial_point_x,robot_initial_point_y) , cv::Point(robot_initial_point_x,robot_initial_point_y-10), cv::Scalar(0,255,0), 1,8); //Y axis

//    cv::namedWindow( "Original Map", CV_WINDOW_NORMAL);
//    cv::resizeWindow("Original Map", 600, 800);
//    cv::imshow("Original Map",input_map_image);
//    cv::waitKey(0);

    // Convert original map image to grayscale
    cv::cvtColor( input_map_image, input_map_image_gs, CV_BGR2GRAY );
    ////// Smooth the image to filter noise and find edges using Canny
    cv::Mat map_im_blurred, map_im_edges;
    cv::GaussianBlur(input_map_image_gs, map_im_blurred, cv::Size(3,3),0,0);
    cv::Canny(map_im_blurred, map_im_edges, 10,30);
//    cv::namedWindow( "Map Edges", CV_WINDOW_NORMAL);
//    cv::resizeWindow("Map Edges", 600, 800);
//    cv::imshow("Map Edges",map_im_edges);
//    cv::waitKey(0);

    ////// Find Contours of the map with largest area and get a rectangle around
    // findContours uses edges to detect contours which possible belong to
    // closed structures like objects
    std::vector<std::vector<cv::Point> > map_im_contours;
    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(map_im_edges, map_im_contours, hierarchy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);
    double largest_area = 0.0;
    int largest_contour_index = 0;
    double second_largest_area = 0.0;
    int second_largest_contour_index = 0;

    for( int i = 0; i< map_im_contours.size(); i++ )
    {
        double a = cv::contourArea( map_im_contours[i],false);
        //std::cout << "Area: " << a << std::endl;
        if(a>=largest_area && (a-largest_area) >= 500)
        {
            second_largest_area = largest_area;
            second_largest_contour_index = largest_contour_index;
            largest_area = a;
            largest_contour_index=i;
        }
        else if (a>second_largest_area && a<largest_area)
        {
            second_largest_area = a;
            second_largest_contour_index = i;
        }
//    // Debug
//    cv::drawContours(input_map_image,map_im_contours,i
//    ,cv::Scalar(255,0,0),0,8,hierarchy);
//    cv::namedWindow( "Map Edges2", CV_WINDOW_NORMAL);
//    cv::resizeWindow("Map Edges2", 600, 800);
//    cv::imshow("Map Edges2",input_map_image);
//    cv::waitKey(0);
//    cv::drawContours(input_map_image,map_im_contours,i
//    ,cv::Scalar(0,0,0),0,8,hierarchy);
    }

    //convex hulls
    std::vector<std::vector<cv::Point> >hull(map_im_contours.size());
    std::vector<std::vector<int> > hullsI(map_im_contours.size());
    std::vector<std::vector<cv::Vec4i>> defects(map_im_contours.size());
    for (int i = 0; i < map_im_contours.size(); i++)
    {
        cv::convexHull(map_im_contours[i], hull[i], false);
        cv::convexHull(map_im_contours[i], hullsI[i], false);
        if(hullsI[i].size() > 3 )
        {
            cv::convexityDefects(map_im_contours[i], hullsI[i], defects[i]);
        }
    }
    //REQUIRED contour is detected,then convex hell is found and also convexity defects are found and stored in defects

    if (largest_area>100){
        std::cout << "Area is found" << std::endl;
        cv::drawContours(input_map_image, hull, second_largest_contour_index, cv::Scalar(0, 255, 255), 2, 8, std::vector<cv::Vec4i>(), 0, cv::Point());

        /// Draw convexityDefects
        for(int j=0; j<defects[second_largest_contour_index].size(); ++j)
        {
            const cv::Vec4i& v = defects[second_largest_contour_index][j];
            float depth = v[3] / 256;
            if (depth > 10) //  filter defects by depth
            {
                int startidx = v[0];
                cv::Point ptStart(map_im_contours[second_largest_contour_index][startidx]);
                int endidx = v[1];
                cv::Point ptEnd(map_im_contours[second_largest_contour_index][endidx]);
                int faridx = v[2];
                cv::Point ptFar(map_im_contours[second_largest_contour_index][faridx]);

                cv::line(input_map_image, ptStart, ptEnd, cv::Scalar(0, 255, 0), 0.2);
                cv::line(input_map_image, ptStart, ptFar, cv::Scalar(0, 255, 0), 0.2);
                cv::line(input_map_image, ptEnd, ptFar, cv::Scalar(0, 255, 0), 0.2);
                cv::circle(input_map_image, ptFar, 1, cv::Scalar(0, 255, 0), 1);
            }
        }
    }

//    cv::namedWindow( "Convex Defects", CV_WINDOW_NORMAL);
//    cv::resizeWindow("Convex Defects", 600, 800);
//    cv::imshow( "Convex Defects", input_map_image );
//    cv::waitKey(0);

    cv::Rect bounding_rect=cv::boundingRect(map_im_contours[second_largest_contour_index]);
    cv::rectangle(input_map_image, bounding_rect,  cv::Scalar(255,0,0),1, 8,0);

    // Crop the original map image to eliminate redundant parts
    cv::Mat map_im_cropped;
    cv::Mat map_im_cropped_gs;
    map_im_cropped = input_map_image(bounding_rect);

//    cv::namedWindow( "Cropped Map Im", CV_WINDOW_NORMAL);
//    cv::resizeWindow("Cropped Map Im", 600, 800);
//    cv::imshow( "Cropped Map Im", map_im_cropped);
//    cv::waitKey(0);

    // Count edge numbers of contours to detect objects
    cv::cvtColor( map_im_cropped, map_im_cropped_gs, CV_BGR2GRAY );
    std::vector<std::vector<cv::Point> > cropped_contours;
    std::vector<cv::Vec4i> cropped_hierarchy;
    cv::findContours(map_im_cropped_gs, cropped_contours, cropped_hierarchy, CV_RETR_LIST, CV_CHAIN_APPROX_SIMPLE);

    std::vector<cv::Point2f> triangle;
    cv::Point2f center, vtx[4];
    float radius = 0;

    for( int i = 0; i< cropped_contours.size(); i++ )
    {
        std::vector<cv::Point> approx_cont;
        cv::approxPolyDP(cropped_contours[i], approx_cont, 0.1*cv::arcLength(cropped_contours[i],true),true);
        std::cout << "Edge Number:" << approx_cont.size() << std::endl;
        if ( !cv::isContourConvex(approx_cont) )
        {
            continue;
        }
        else if (approx_cont.size() == 3 && (cv::contourArea( cropped_contours[i],false) < largest_area/10) && (cv::contourArea( cropped_contours[i],false) > (largest_area/1000)))
        {
            //cv::drawContours(map_im_cropped,cropped_contours,i,cv::Scalar(0,255,255),0,8,cropped_hierarchy);
            cv::minEnclosingTriangle(cropped_contours[i], triangle);
        }
        else if ( (approx_cont.size() == 4 || approx_cont.size() == 5) && (cv::contourArea( cropped_contours[i],false) < (largest_area/10)) && (cv::contourArea( cropped_contours[i],false) > (largest_area/500)) )
        {
            //cv::drawContours(map_im_cropped,cropped_contours,i,cv::Scalar(255,0,255),0,8,cropped_hierarchy);
            //cv::RotatedRect box = minAreaRect(cropped_contours[i]);
            //box.points(vtx);
            cv::Rect bound_rect=cv::boundingRect(cropped_contours[i]);
            cv::rectangle(map_im_cropped, bound_rect,  cv::Scalar(255,0,0),1, 8,0);
        }
        else if (approx_cont.size() >= 6 && (cv::contourArea( cropped_contours[i],false) < (largest_area/10)) && (cv::contourArea( cropped_contours[i],false) > (largest_area/1000)))
        {
            //cv::drawContours(map_im_cropped,cropped_contours,i,cv::Scalar(255,0,255),0,8,cropped_hierarchy);
            cv::minEnclosingCircle(cv::Mat(cropped_contours[i]), center, radius);
        }
    }

    // Draw Detected Objects

    for(int i = 0; i < 4; i++ )
    {
        cv::line(map_im_cropped, vtx[i], vtx[(i+1)%4], cv::Scalar(255, 0, 0), 1);
    }
    cv::circle(map_im_cropped, center, cvRound(radius), cv::Scalar(0, 255, 0), 1);

    for( int i = 0; i < 3; i++ )
    {
        cv::line(map_im_cropped, triangle[i], triangle[(i+1)%3], cv::Scalar(0, 0, 255), 1);
    }

    //cv::namedWindow( "Detected Objects", CV_WINDOW_NORMAL);
    //cv::resizeWindow("Detected Objects", 600, 800);
    //cv::imshow("Detected Objects",map_im_cropped);
    //cv::waitKey(0);

    std::string output_image = file_dir + "_processed.png";
    cv::imwrite(output_image, map_im_cropped);

    return 0;
}


/**
 * @file
 * @author University of Warwick
 * @version 1.0
 *
 * @section LICENSE
 *
 * @section DESCRIPTION
 *
 * Tests for the concrete methods of the DistributionConfigSourceNormalJSON class.
 */

#define BOOST_TEST_MODULE DistributionConfigSourceNormalJSON

// Boost setup
#include <boost/test/unit_test.hpp>
#include <boost/test/output_test_stream.hpp>
#include <stdexcept>

#include "DistributionConfigSourceNormalJSON.h"
#include "Error.h"

// Namespaces
namespace utf = boost::unit_test;

using namespace cupcfd::distributions;

// === Constructor ===
// Test 1: Parameter Constructor
BOOST_AUTO_TEST_CASE(constructor_test1, * utf::tolerance(0.00001))
{
	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);
}

// Test 2: Configuration Copy Constructor
BOOST_AUTO_TEST_CASE(constructor_test2, * utf::tolerance(0.00001))
{

}

// === operator= ===
// Test 1: Test correct copy
BOOST_AUTO_TEST_CASE(operator_equals_test1, * utf::tolerance(0.00001))
{

}

// === clone ===
BOOST_AUTO_TEST_CASE(clone_test1, * utf::tolerance(0.00001))
{

}

// === getLBound ===
// Test 1: Get the correct value
BOOST_AUTO_TEST_CASE(getLBound_test1, * utf::tolerance(0.00001))
{
	cupcfd::error::eCodes status;
	double lBound;

	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);

	status = jsonConfigSource.getLBound(&lBound);
	BOOST_CHECK_EQUAL(status, cupcfd::error::E_SUCCESS);
	BOOST_TEST(lBound == 2.3);
}

// === getUBound ===
BOOST_AUTO_TEST_CASE(getUBound_test1, * utf::tolerance(0.00001))
{
	cupcfd::error::eCodes status;
	double uBound;

	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);

	status = jsonConfigSource.getUBound(&uBound);
	BOOST_CHECK_EQUAL(status, cupcfd::error::E_SUCCESS);
	BOOST_TEST(uBound == 4.6);
}

// === getMean ===
BOOST_AUTO_TEST_CASE(getMean_test1, * utf::tolerance(0.00001))
{
	cupcfd::error::eCodes status;
	double mean;

	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);

	status = jsonConfigSource.getMean(&mean);
	BOOST_CHECK_EQUAL(status, cupcfd::error::E_SUCCESS);
	BOOST_TEST(mean == 3.7);
}

// === getStDev ===
BOOST_AUTO_TEST_CASE(getStDev_test1, * utf::tolerance(0.00001))
{
	cupcfd::error::eCodes status;
	double stDev;

	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);

	status = jsonConfigSource.getStDev(&stDev);
	BOOST_CHECK_EQUAL(status, cupcfd::error::E_SUCCESS);
	BOOST_TEST(stDev == 0.3);
}

// === buildDistributionConfig ===
BOOST_AUTO_TEST_CASE(buildDistribution_test1, * utf::tolerance(0.00001))
{
	cupcfd::error::eCodes status;

	std::string topLevels[1] = {"Level1"};
	DistributionConfigSourceNormalJSON<int,double> jsonConfigSource("../tests/distributions/data/DistributionConfigSourceNormalTest.json", topLevels, 1);

	DistributionConfig<int,double> * distConfig;
	status = jsonConfigSource.buildDistributionConfig(&distConfig);
	BOOST_CHECK_EQUAL(status, cupcfd::error::E_SUCCESS);
}


#include <boost/graph/wavefront.hpp>


/* Copyright (c) 2012, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See LICENSE.txt or
 * http://www.opensource.org/licenses/mit-license.php */

#pragma once

#include "hdp_base.hpp"
#include "hdp_var_base.hpp"

#include "random.hpp"
#include "baseMeasure.hpp"
#include "probabilityHelpers.hpp"

#include <stddef.h>
#include <stdint.h>
#include <typeinfo>

#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/beta.hpp>
#include <armadillo>

using namespace std;
using namespace arma;

/*
 * this one assumes that the number of words per document is
 * smaller than the dictionary size
 *
 * http://en.wikipedia.org/wiki/Virtual_inheritance
 */
template <class U>
class HDP_var: public HDP<U>, public virtual HDP_var_base
{
  public:

    HDP_var(const BaseMeasure<U>& base, double alpha, double omega)
      : HDP_var_base(0,0,0), HDP<U>(base, alpha, omega)
    {};

    ~HDP_var()
    {};

    // interface mainly for python
    uint32_t addDoc(const Mat<U>& x_i)
    {
      uint32_t x_ind = HDP<U>::addDoc(x_i);
      // TODO: potentially slow
      // add the index of the added x_i
      mInd2Proc.resize(mInd2Proc.n_elem+1);
      mInd2Proc[mInd2Proc.n_elem-1] = x_ind;
      return x_ind;
    };

    /*
     * Initializes the corpus level parameters mA and mLambda according
     * to Blei's Stochastic Variational paper
     *
     * @param D is the assumed number of documents for init
     */
    void initCorpusParams(uint32_t Nw, uint32_t K, uint32_t T, uint32_t D)
    {

      mT = T;
      mK = K;
      mNw = Nw;

      mA.ones(K,2);
      mA.col(1) *= HDP<U>::mOmega;

      // initialize lambda
      HDP<U>::mLambda.init(HDP<U>::mH0,K); // initialize the priors of the topics from the base measure.

//      HDP<U>::mLambda.zeros(K,Nw);
//      GammaRnd gammaRnd(1.0, 1.0);
//      for (uint32_t k=0; k<K; ++k){
//        for (uint32_t w=0; w<Nw; ++w) HDP<U>::mLambda(k,w) = gammaRnd.draw();
//        HDP<U>::mLambda.row(k) += ((Dir*)(&mH0))->mAlphas;
//      }
    };

    /* From: Online Variational Inference for the HDP
     * method for "one shot" computation without storing data in this class
     *  Nw: number of different words
     *  kappa=0.9: forgetting rate
     *  uint32_t T=10; // truncation on document level
     *  uint32_t K=100; // truncation on corpus level
     *  S = batch size
     */
    void densityEst(const vector<Mat<U> >& x, uint32_t Nw,
        double kappa, uint32_t K, uint32_t T, uint32_t S)
    {
      cout<<"densityEstimate with: K="<<K<<"; T="<<T<<"; kappa="<<kappa<<"; Nw="<<Nw<<"; S="<<S<<endl;

      HDP<U>::mX = x;
      uint32_t D=HDP<U>::mX.size();
      cout<<"D="<<D<<endl;
      cout<<"mX[0].shape= "<<HDP<U>::mX[0].n_rows<<"x"<<HDP<U>::mX[0].n_cols<<endl;

      mInd2Proc = linspace<Row<uint32_t> >(0,D-1,D);

      mT = T;
      mK = K;
      mNw = Nw;

      initCorpusParams(mNw,mK,mT,D);
      cout<<"Init of corpus params done"<<endl;
      Row<uint32_t> ind = updateEst_batch(mInd2Proc,mZeta,mPhi,mGamma,mA,HDP<U>::mLambda,mPerp,HDP<U>::mOmega,kappa,S,true);
//      cout<<"mPhi -> D="<<mPhi.size()<<endl;
//      cout<<"mPhi -> D="<<HDP_var_base::mPhi.size()<<endl;
//      cout<<"mPerp="<<mPerp.t()<<endl;

      Mat<double> pi(D,T);
      Mat<double> sigPi(D,T+1);
      Mat<uint32_t> c(D,T);
      getDocTopics(pi,sigPi,c);
//      cout<<"c:"<<c<<endl;

      mInd2Proc.set_size(0); // all processed

    };

    /*
     * compute density estimate based on data previously fed into the class using addDoc
     */
    bool densityEst(uint32_t Nw, double kappa, uint32_t K, uint32_t T, uint32_t S)
    {
      if(HDP<U>::mX.size() > 0)
      {
        densityEst(HDP<U>::mX,Nw,kappa,K,T,S);
        //TODO: return p_d(x)
        return true;
      }else{
        return false;
      }
    };

    /*
     * updated the estimate using newly added docs in mX
     * "newly added": the ones that are indicated in mInd2Proc
     */
    bool updateEst_batch(double kappa, uint32_t S)
    {
      uint32_t Db=mInd2Proc.n_elem;
      if (Db >0){
        cout<<"updatedEstimate with: K="<<mK<<"; T="<<mT<<"; kappa="<<kappa<<"; Nw="<<mNw<<"; S="<<S<<endl;
        vector<Mat<double> > zeta; // will get resized accordingly inside updateEst_batch
        vector<Mat<double> > phi;
        vector<Mat<double> > gamma;
        Col<double> perp;

        Row<uint32_t> ind = updateEst_batch(mInd2Proc,zeta,phi,gamma,mA,HDP<U>::mLambda,perp,HDP<U>::mOmega,kappa,S);

        mZeta.resize(mZeta.size()+Db);
        mPhi.resize(mPhi.size()+Db);
        mGamma.resize(mGamma.size()+Db);
        mPerp.resize(mPerp.n_elem+Db);
        mPerp.rows(mPerp.n_elem-Db,mPerp.n_elem-1) = perp;
        for (uint32_t i=0; i<ind.n_elem; ++i){
          mZeta[ind[i]] = zeta[i];
          mPhi[ind[i]] = phi[i];
          mGamma[ind[i]] = gamma[i];
        }

        mInd2Proc.set_size(0); // all processed

        return true;
      }else{
        cout<<"add more documents before starting to process"<<endl;
      }
      return false;
    };

    /*
     * after an initial densitiy estimate has been made using densityEst()
     * can use this to update the estimate with information from additional x
     */
    bool  updateEst(const Mat<U>& x, double kappa)
    {
      if (HDP<U>::mX.size() > 0 && HDP<U>::mX.size() == mPhi.size()) { // this should indicate that there exists an estimate already
        uint32_t N = x.n_cols;
        uint32_t T = mT;
        uint32_t K = mK;
        HDP<U>::mX.push_back(x);
        mZeta.push_back(Mat<double>(T,K));
        mPhi.push_back(Mat<double>(N,T));
        //    mZeta.set_size(T,K);
        //    mPhi.set_size(N,T);
        mGamma.push_back(Mat<double>(T,2));
        uint32_t d = HDP<U>::mX.size()-1;
        mPerp.resize(d+1);

        if(updateEst(HDP<U>::mX[d],mZeta[d],mPhi[d],mGamma[d],mA,HDP<U>::mLambda,HDP<U>::mOmega,d,kappa))
        {
          mPerp[d] = 0.0;
          for (uint32_t i=0; i < HDP<U>::mX_ho.size(); ++i)
          {
            Row<double> logP=logP_w(HDP<U>::mX[d],mPhi[d], mZeta[d], mGamma[d], HDP<U>::mLambda);
            mPerp[d] += HDP<U>::perplexity(HDP<U>::mX_ho[i], logP);
            //mPerp[d] += perplexity(mX_ho[i], mZeta[d], mPhi[d], mGamma[d], HDP<U>::mLambda);
          }
          mPerp[d] /= double(HDP<U>::mX_ho.size());
          cout<<"Perplexity="<<mPerp[d]<<endl;
          return true;
        }else{
          return false;
        }
      }else{
        return false;
      }
    };

    /*
     *
     */
    bool updateEst(const Mat<U>& x, Mat<double>& zeta, Mat<double>& phi, Mat<double>& gamma, Mat<double>& a, DistriContainer<U>& lambda, double omega, uint32_t d, double kappa)
    {
      uint32_t D = d+1; // assume that doc d is appended to the end
//      uint32_t Nw = lambda.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      Mat<double> eLogBeta(mK,x.n_cols); //TODO
      Col<double> digam_lamb_sum(mK);

      compElogBeta(eLogBeta,  lambda, x);

      Col<double> eLogSig_a(K);
      compElogSig(eLogSig_a, a);

      //    cout<<"---------------- Document "<<d<<" N="<<N<<" -------------------"<<endl;
      initZeta(zeta,eLogBeta);
      initPhi(phi,zeta,eLogBeta);

      if(!is_finite(zeta))
      {
        cout<<"updateEst::a="<<size(a)<<endl<<a;
        cout<<"updateEst::lambda="<<lambda[0]->asRow();
        cout<<"updateEst::x="<<size(x)<<endl<<x;
        cout<<"zeta_init="<<zeta<<endl;
        cout<<"phi_init="<<phi<<endl;
        cout<<"eLogBeta="<<eLogBeta<<endl;
        exit(0);
      }

      // ------------------------ doc level updates --------------------
      bool converged = false;
      Col<double> eLogSig_gam(T);
      Mat<double> gamma_prev(T,2);
      gamma_prev.ones();

      uint32_t o=0;
      while(!converged){
        //      cout<<"-------------- Iterating local params #"<<o<<" -------------------------"<<endl;
        updateGamma(gamma,phi);

        if (!is_finite(gamma)){
          cout<<"gamma="<<gamma;
          cout<<"phi="<<phi;
          exit(1);
        }

        compElogSig(eLogSig_gam,gamma); // precompute

        updateZeta(zeta,phi,eLogSig_a,eLogBeta);
        updatePhi(phi,zeta,eLogSig_gam,eLogBeta);

        converged = (accu(gamma_prev != gamma))==0 || o>60 ;
        gamma_prev = gamma;
        ++o;

        if(!is_finite(zeta))
        {
          cout<<"o="<<o<<endl;
          cout<<"zeta="<<zeta.row(0)<<" |.|="<<sum(exp(zeta.row(0)))<<endl;
          cout<<"phi="<<phi.row(0)<<" |.|="<<sum(exp(phi.row(0)))<<endl;
          exit(0);
        }
      }

      //    cout<<" --------------------- natural gradients --------------------------- "<<endl;
      //    cout<<"\tD="<<D<<" omega="<<omega<<endl;
      DistriContainer<U> d_lambda(HDP<U>::mH0,K);
      Mat<double> d_a(K,2);
      computeNaturalGradients(d_lambda, d_a, zeta, phi, omega, D, x);

      cout<<"update::d_lambda:"<<d_lambda.toMat().rows(0,5);
      cout<<"update::lambda:"<<lambda.toMat().rows(0,5);

      //    cout<<" ------------------- global parameter updates: ---------------"<<endl;
      double ro = exp(-kappa*log(1+double(d+1)));
      //    cout<<"\tro="<<ro<<endl;
      for (uint32_t k=0; k<mK; ++k)
        lambda[k]->fromRow( (1.0-ro)*lambda[k]->asRow() + ro* d_lambda[k]->asRow());

      //lambda = (1.0-ro)*lambda + ro*d_lambda;
      a = (1.0-ro)*a+ ro*d_a;
      return true;
    };

    /*
     * Updates the estimate using mini batches
     * @param ind_x indices of docs to process within docs mX. !these are assumed to be in order!
     * @param sameIndAsX == true -> zeta,phi,gamma have same indices as x (ind_x). This typically happens for the initial batch update. Setting this to true eliminates the need of reordering the results afterwords to match the indices of the docs x.
     * @return the randomly shuffled indices to show how the data was processed -> this allows association of zetas, phis and gammas with docs in mX
     */
    Row<uint32_t> updateEst_batch(const Row<uint32_t>& ind_x, vector<Mat<double> >& zeta, vector<Mat<double> >& phi, vector<Mat<double> >& gamma, Mat<double>& a, DistriContainer<U>& lambda, Col<double>& perp, double omega, double kappa, uint32_t S, bool sameIndAsX=false)
    {
      uint32_t d_0 = min(ind_x); // thats the doc number that we start with -> needed for ro computation; assumes that all indices in mX prior to d_0 have already been processed.
      uint32_t D= max(ind_x)+1; // D is the maximal index of docs that we are processing +1

      Row<uint32_t> ind = shuffle(ind_x,1);
//        cout<<"ind_x: "<<ind_x.cols(0,S)<<endl;
//        cout<<"ind  : "<<ind.cols(0,S)<<endl;

      zeta.resize(ind.n_elem,Mat<double>(mT,mK));
      phi.resize(ind.n_elem);
      gamma.resize(ind.n_elem,Mat<double>(mT,2));
      perp.zeros(ind.n_elem);

      for (uint32_t dd=0; dd<ind.n_elem; dd += S)
      {
        DistriContainer<U> db_lambda(HDP<U>::mH0,mK);

        Mat<double> db_a(mK,2);
        db_a.zeros();

        Col<double> eLogSig_a(mK);
        compElogSig(eLogSig_a, a);

#pragma omp parallel for schedule(dynamic)
        for (uint32_t db=dd; db<min(dd+S,ind.n_elem); db++)
        {
          uint32_t d=ind[db];
          const Mat<U>& x_d = HDP<U>::mX[d];
          uint32_t dout=sameIndAsX?d:db;
          uint32_t N=x_d.n_cols;
          //      cout<<"---------------- Document "<<d<<" N="<<N<<" -------------------"<<endl;

          cout<<"-- db="<<db<<" d="<<d<<" N="<<N<<endl;

          Mat<double> eLogBeta(mK,x_d.n_cols);
          Col<double> digam_lamb_sum(mK);
          compElogBeta(eLogBeta, lambda, x_d);

          //Mat<double> zeta(T,K);
          phi[dout].resize(N,mT);
          initZeta(zeta[dout],eLogBeta);
          initPhi(phi[dout],zeta[dout],eLogBeta);

//            cout<<"zeta_init="<<zeta[dout]<<endl;
//            cout<<"phi_init="<<phi[dout]<<endl;
//            cout<<"eLogBeta="<<eLogBeta<<endl;
//            if(!is_finite(zeta[dout]))
//              exit(0);

          //cout<<" ------------------------ doc level updates --------------------"<<endl;
          //Mat<double> gamma(T,2);
          Col<double> eLogSig_gam(mT);
          Mat<double> gamma_prev(mT,2);
          gamma_prev.ones();
          gamma_prev.col(1) += HDP<U>::mAlpha;
          bool converged = false;
          uint32_t o=0;
          while(!converged){
//            cout<<"-------------- Iterating local params #"<<o<<" -------------------------"<<endl;
            updateGamma(gamma[dout],phi[dout]);

            if (!is_finite(gamma[dout])){
              cout<<"gamma="<<gamma[dout];
              cout<<"phi="<<phi[dout];
              exit(1);
            }

            compElogSig(eLogSig_gam,gamma[dout]); // precompute

            updateZeta(zeta[dout],phi[dout],eLogSig_a,eLogBeta);
            updatePhi(phi[dout],zeta[dout],eLogSig_gam,eLogBeta);

            converged = (accu(gamma_prev != gamma[dout]))==0 || o>30 ;
            gamma_prev = gamma[dout];
            ++o;

//            cout<<"o="<<o<<endl;
//            cout<<"zeta="<<zeta[dout].row(0)<<" |.|="<<sum(zeta[dout].row(0))<<endl;
//            cout<<"phi="<<phi[dout].row(0)<<" |.|="<<sum(phi[dout].row(0))<<endl;
//            if(!is_finite(zeta[dout]))
//              exit(0);
          }

          DistriContainer<U> d_lambda(HDP<U>::mH0,mK); // batch updates
          Mat<double> d_a(mK,2);
//          cout<<" --------------------- natural gradients dout="<< dout<<" dd="<< dd<<" --------------------------- "<<endl;
          computeNaturalGradients(d_lambda, d_a, zeta[dout], phi[dout], HDP<U>::mOmega, D, x_d);
#pragma omp critical
          {
            cout<<"update_batch::d_lambda:"<<endl<<d_lambda.toMat().rows(0,5);
            for (uint32_t k=0; k<d_lambda.size(); ++k){
              db_lambda[k]->fromRow(db_lambda[k]->asRow() + d_lambda[k]->asRow());
            }
            db_a += d_a;
          }
        }
        //for (uint32_t k=0; k<K; ++k)
        //  cout<<"delta lambda_"<<k<<" min="<<min(d_lambda.row(k))<<" max="<< max(d_lambda.row(k))<<" #greater 0.1="<<sum(d_lambda.row(k)>0.1)<<endl;
        // ----------------------- update global params -----------------------
        uint32_t t=dd+d_0; // d_0 is the timestep of the the first index to process; dd is the index in the current batch
        uint32_t bS = min(S,ind.n_elem-dd); // necessary for the last batch, which migth not form a complete batch
        //TODO: what is the time dd? d_0 needed?
        double ro = exp(-kappa*log(1+double(t)+double(bS)/2.0)); // as "time" use the middle of the batch
//        cout<<" -- global parameter updates t="<<t<<" bS="<<bS<<" ro="<<ro<<endl;
//        cout<<"d_a="<<db_a<<endl;

//        cout<<"dLambda"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<db_lambda[k]->asRow();
//
//        cout<<"Before"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<lambda[k]->asRow();

        cout<<"update_batch::db_lambda:"<<endl<<db_lambda.toMat().rows(0,5);
        for (uint32_t k=0; k<db_lambda.size(); ++k)
          lambda[k]->fromRow((1.0-ro)*lambda[k]->asRow() + (ro/S)*db_lambda[k]->asRow()); //TODO: doies this make sense for NIW prior???
        cout<<"update_batch::lambda(after):"<<endl<<lambda.toMat().rows(0,5);

//        cout<<"After"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<lambda[k]->asRow();

        //lambda = (1.0-ro)*lambda + (ro/S)*db_lambda;
        a = (1.0-ro)*a + (ro/S)*db_a;
        cout<<"update_batch::lambda:"<<lambda.toMat().rows(0,5);

        perp[dd+bS/2] = 0.0;
        if (HDP<U>::mX_te.size() > 0) {
          cout<<"computing "<<HDP<U>::mX_te.size()<<" perplexities"<<endl;
#pragma omp parallel for schedule(dynamic)
          for (uint32_t i=0; i < HDP<U>::mX_te.size(); ++i)
          {
            //cout<<"mX_te: "<< mX_te[i].n_rows << "x"<< mX_te[i].n_cols<<endl;
            //cout<<"mX_ho: "<< mX_ho[i].n_rows << "x"<< mX_ho[i].n_cols<<endl;
            //TODO: these subfunctions work on the member variables!!! dont do that...
            double perp_i =  perplexity(HDP<U>::mX_te[i],HDP<U>::mX_ho[i],dd+bS/2+1,ro); //perplexity(mX_ho[i], mZeta[d], mPhi[d], mGamma[d], lambda);
            //cout<<"perp_"<<i<<"="<<perp_i<<endl;
#pragma omp critical
            {
              perp[dd+bS/2] += perp_i;
            }
          }
          perp[dd+bS/2] /= double(HDP<U>::mX_te.size());
          cout<<"Perplexity="<<perp[dd+bS/2]<<endl;
        }
      }
      cout<<"perp="<<perp.t()<<endl;
      return ind;
    };

    // compute the perplexity of a given document split into x_test (to find a topic model for the doc) and x_ho (to evaluate the perplexity)
    double perplexity(const Mat<U>& x_te, const Mat<U>& x_ho, uint32_t d, double kappa=0.75)
    {
      if (HDP<U>::mX.size() > 0 && HDP<U>::mX.size() == mPhi.size()) { // this should indicate that there exists a estimate already

        uint32_t N = x_te.n_cols;
        uint32_t T = mT;
        uint32_t K = mK;

        Mat<double> zeta(T,K);
        Mat<double> phi(N,T);
        Mat<double> gamma(T,2);
        //uint32_t d = mX.size()-1;

        Mat<double> a(mA);// DONE: make deep copy here!

        DistriContainer<U> lambda(HDP<U>::mLambda);
        double omega = HDP<U>::mOmega;

        cout<<"perplexity::lambda:"<<lambda.toMat().rows(0,5);
        cout<<"perplexity::mLambda:"<<HDP<U>::mLambda.toMat().rows(0,5);

        cout<<"updating copied model with x"<<endl;
        updateEst(x_te,zeta,phi,gamma,a,lambda,omega,d,kappa);
//        cout<<" lambda.shape="<<lambda.size()<<endl;
        cout<<"computing perplexity under updated model"<<endl;
        //TODO: compute probabilities then use that to compute perplexity

        cout<<"x_te: "<<size(x_te);
        Mixture<U> mix = docMixture(phi, zeta, gamma, lambda);
        return HDP<U>::perplexity(x_ho, mix);
        //return perplexity(x_ho, zeta, phi, gamma, lambda);
      }else{
        return 1.0/0.0;
      }
    };

    Mixture<U> docMixture(uint32_t d) const {
      return docMixture(mPhi[d],mZeta[d],mGamma[d],HDP<U>::mLambda);
    };
    Mixture<U> docMixture(const Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& gamma, const DistriContainer<U>& lambda) const
    {
      Col<double> pi;
      Col<double> sigPi;
      Col<uint32_t> c;
      getDocTopics(pi,sigPi,c,gamma,zeta);
      //cout<<"getDocTopics done"<<endl;
      //cout<<"c="<<c.t()<<size(c);
      Col<uint32_t> z(mNw);
      getWordTopics(z, phi);
      //cout<<"getWordTopics done"<<endl;
      //cout<<"z="<<z.t()<<size(z);

      DistriContainer<U> beta;
      HDP<U>::getCorpTopics(beta,lambda);
//      cout<<"getCorpTopics done"<<endl;
      //cout<<"beta:\t"<<size(beta);

      Col<uint32_t> c_u = unique(c);
      Row<double> ps(c_u.n_elem); // proportions in  the mixture
      DistriContainer<U> beta_d(c_u.n_elem);
      for (uint32_t i=0; i< c_u.n_elem; ++i){
        beta_d[i] = beta[c_u(i)]->getCopy();
        ps[i] = sum(sigPi.elem(find(c == c_u(i) )));
      }
//      cout<<"Mixture done"<<endl;
      return Mixture<U>(beta_d,ps);
    };

    /* Probability distribution over the words in document d
     *
     * TODO: so is that here not some MAP or ML estimate?!
     */
    Row<double> logP_w(uint32_t d) const {
//      cout<<"mX.size="<<mX.size()<<endl;
      return logP_w(HDP<U>::mX[d],mPhi[d],mZeta[d],mGamma[d],HDP<U>::mLambda);
    };
    /*
     * log probability using the samples x (not using sufficient statistics -> x is just a list of words)
     */
    Row<double> logP_w(const Mat<U>& x, const Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& gamma, const DistriContainer<U>& lambda) const
    {
      Row<double> p(x.n_cols);
      p.zeros();
//      cout<<"x:\t"<<size(x);
//      cout<<"phi:\t"<<size(phi);
//      cout<<"zeta:\t"<<size(zeta);
//      cout<<"gamma:\t"<<size(gamma);
//      cout<<"lambda:\t"<<size(lambda);

      Col<double> pi;
      Col<double> sigPi;
      Col<uint32_t> c;
      getDocTopics(pi,sigPi,c,gamma,zeta);
      //cout<<"getDocTopics done"<<endl;
      //cout<<"c="<<c.t()<<size(c);
      Col<uint32_t> z(mNw);
      getWordTopics(z, phi);
      //cout<<"getWordTopics done"<<endl;
      //cout<<"z="<<z.t()<<size(z);

      DistriContainer<U> beta;
      HDP<U>::getCorpTopics(beta,lambda);
      //cout<<"getCorpTopics done"<<endl;
      //cout<<"beta:\t"<<size(beta);

      for (uint32_t i=0; i<x.n_cols; ++i){
        //cout<<"z_"<<i<<"="<<z[i]<<endl;
        p[i] = beta[c[z[i]]]->logP(x.col(i));// logCat(x[i], beta.row( c[ z[i] ]));
        //cout<<"p_"<<x[i]<<"="<<p[x[i]]<<endl;
      }
//      for (uint32_t w=0; w<mNw; ++w)
//        p[w] = p[w]==0.0?-1e10:p[w];

      //cout<<"p="<<p<<endl;
      return p;
    };

  protected:

    Row<uint32_t> mInd2Proc; // indices of docs that have not been processed

  private:

    /*
     * precompute necessary digamma function values, because these are slowing the whole algorithm down
     * all the update methods for zeta and phi need these values very often! I can precumpute these once after updating the global parameters (and hence lambda)
     */
    void compElogBeta(Mat<double>& eLogBeta, const DistriContainer<U>& lambda, const Mat<U>& x_d) const
    {
      eLogBeta.zeros(mK,x_d.n_cols);

      Col<double> digam_lamb_sum(mK);
    //  Mat<uint32_t> x_u = unique(x_d); // cannot do this trick anymore since x_d are continuous for the NIW base measure
//      cout<<"x_d="<<x_d<<endl;
//      cout<<"x_u="<<x_u<<endl;

      for (uint32_t i = 0; i < x_d.n_cols ; i++) {
        for (uint32_t k = 0; k < mK; k++) {
           eLogBeta(k,i) = lambda[k]->Elog(x_d.col(i)); // E[log beta] computation in paper
        }
      }
//      for (uint32_t k = 0; k < mK; k++) {
//        digam_lamb_sum(k) = digamma(sum(lambda.row(k)));
//      }
//      for (uint32_t i = 0; i < x_u.n_elem ; i++) {
//        for (uint32_t k = 0; k < mK; k++) {
//           eLogBeta(k,x_u(i)) = digamma(lambda(k,x_u(i))) - digam_lamb_sum(k);
//        }
//      }
    }

    void compElogSig(Col<double>& eLogSig, const Mat<double>& a) const
    {
      for (uint32_t k=0; k<a.n_rows; ++k){
        eLogSig(k) = digamma(a(k,0)) - digamma(a(k,0) + a(k,1));
        for (uint32_t l=0; l<k; ++l)
          eLogSig(k) += digamma(a(l,1)) - digamma(a(l,0) + a(l,1));
      }
    }

    void initZeta(Mat<double>& zeta, const Mat<double>& eLogBeta)
    {
      uint32_t N = eLogBeta.n_cols; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;
      //cerr<<"\tinit zeta"<<endl;
      for (uint32_t i=0; i<T; ++i) {
        for (uint32_t k=0; k<K; ++k) {
          zeta(i,k)=0.0;
          for (uint32_t n=0; n<N; ++n) {
            //if(i==0 && k==0) cout<<zeta(i,k)<<" -> ";
            zeta(i,k) += eLogBeta(k,n); //ElogBeta(lambda, k, x_d(n));
          }
        }
        normalizeLogDistribution(zeta.row(i));
        //cout<<" normalized="<<zeta(0,0)<<endl;
      }
      //cerr<<"zeta>"<<endl<<zeta<<"<zeta"<<endl;
      //cerr<<"normalization check:"<<endl<<sum(zeta,1).t()<<endl; // sum over rows
    };

    void initPhi(Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& eLogBeta)
    {
      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;
      //cout<<"\tinit phi"<<endl;
      for (uint32_t n=0; n<N; ++n){
        for (uint32_t i=0; i<T; ++i) {
          phi(n,i)=0.0;
          for (uint32_t k=0; k<K; ++k) {
            phi(n,i)+=zeta(i,k)* eLogBeta(k,n); // ElogBeta(lambda, k, x_d(n));
          }
        }
        normalizeLogDistribution(phi.row(n));
      }
      //cerr<<"phi>"<<endl<<phi<<"<phi"<<endl;
    };

    void updateGamma(Mat<double>& gamma, const Mat<double>& phi)
    {
      uint32_t N = phi.n_rows;
      uint32_t T = phi.n_cols;

      gamma.ones();
      gamma.col(1) *= HDP<U>::mAlpha;
      for (uint32_t i=0; i<T; ++i)
      {
        for (uint32_t n=0; n<N; ++n){
          gamma(i,0) += phi(n,i);
          for (uint32_t j=i+1; j<T; ++j) {
            gamma(i,1) += phi(n,j);
          }
        }
      }
      //cout<<gamma.t()<<endl;
    };

    void updateZeta(Mat<double>& zeta, const Mat<double>& phi, const Col<double>& eLogSig_a, const Mat<double>& eLogBeta)
    {
//      assert(x_d.n_rows == 1);

      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      for (uint32_t i=0; i<T; ++i){
        //zeta(i,k)=0.0;
        for (uint32_t k=0; k<K; ++k) {
          zeta(i,k) = eLogSig_a(k); //ElogSigma(a,k);
          //cout<<zeta(i,k)<<endl;
          for (uint32_t n=0; n<N; ++n){
            zeta(i,k) += phi(n,i)* eLogBeta(k,n); //ElogBeta(lambda,k,x_d(n));
          }
        }
        normalizeLogDistribution(zeta.row(i));
      }
    }

    void updatePhi(Mat<double>& phi, const Mat<double>& zeta, const Col<double>& eLogSig_gam, const Mat<double>& eLogBeta)
    {
//      assert(x_d.n_rows == 1);

      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      for (uint32_t n=0; n<N; ++n){
        //phi(n,i)=0.0;
        for (uint32_t i=0; i<T; ++i) {
          phi(n,i) = eLogSig_gam(i); //ElogSigma(gamma,i);
          for (uint32_t k=0; k<K; ++k) {
            phi(n,i) += zeta(i,k)* eLogBeta(k,n); //ElogBeta(lambda,k,x_d(n)) ;
          }
        }
        normalizeLogDistribution(phi.row(n));
      }
    }

    void computeNaturalGradients(DistriContainer<U>& d_lambda, Mat<double>& d_a, const Mat<double>& zeta, const Mat<double>&  phi, double omega, uint32_t D, const Mat<U>& x_d)
    {
//      uint32_t N = x_d.n_cols;
//      uint32_t Nw = d_lambda.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

//      d_lambda.init(HDP<U>::mH0,mK);

//      d_lambda.zeros();
      d_a.zeros();
      for (uint32_t k=0; k<K; ++k)
      { // for all K corpus level topics
//        cout<<zeta.col(k).t()<<endl;
//        cout<<phi;
        d_lambda[k]->posteriorHDP_var(zeta.col(k),phi,D,x_d);

//        for (uint32_t i=0; i<T; ++i)
//        {
//          Row<double> _lambda(Nw); _lambda.zeros();
//          for (uint32_t n=0; n<N; ++n){
//            _lambda(x_d(n)) += phi(n,i);
//          }
//          d_lambda.row(k) += zeta(i,k) * _lambda;
//        }
//        d_lambda.row(k) = D*d_lambda.row(k);
//        //cout<<"lambda-nu="<<d_lambda[k].t()<<endl;
//        d_lambda.row(k) += ((Dir*)(&mH0))->mAlphas;
//        //cout<<"lambda="<<d_lambda[k].t()<<endl;

        for (uint32_t i=0; i<T; ++i)
        {
          d_a(k,0) += zeta(i,k);
          for (uint32_t l=k+1; l<K; ++l) {
            d_a(k,1) += zeta(i,l);
          }
        }
        d_a(k,0) = D*d_a(k,0)+1.0;
        d_a(k,1) = D*d_a(k,1)+omega;
      }
    }

    //bool normalizeLogDistribution(Row<double>& r)
    bool normalizeLogDistribution(arma::subview_row<double> r) const
    {
      // known as the log sum exp trick!
//      cout<<" r="<<r<<endl;
      double maxR = as_scalar(max(r));
//      cout<<"  maxR="<<maxR<<endl;
//      cout<<"  exp(r-maxR)="<<exp(r-maxR)<<endl;

      r -= maxR + log(sum(exp(r-maxR)));
      r=exp(r);

//      cout<<" r="<<r<<endl;
      return true;

//      double minR = as_scalar(min(r));
//      cout<<" minR="<<minR<<endl;
//      if(minR > -100.0) {
//        cout<<" logDenom="<<sum(exp(r),1)<<endl;
//        double denom = as_scalar(sum(exp(r),1));
//        cout<<" logDenom="<<denom<<endl;
//        r -= log(denom); // avoid division by 0
//        cout<<" r - logDenom="<<r<<endl;
//        r = exp(r);
//        r /= sum(r);
//        cout<<" exp(r - logDenom)="<<r<<endl;
//        return true;
//      }else{ // cannot compute this -> set the smallest r to 1.0 and the rest to 0
//        double maxR = as_scalar(max(r));
//        cout<<"maxR="<<maxR<<" <-" <<arma::max(r) <<endl;
//        uint32_t kMax=as_scalar(find(r==maxR,1));
//        cout<<"maxR="<<maxR<<" kMax="<<kMax<<" <-" <<arma::max(r) <<endl;
//        r.zeros();
//        r(kMax)=1.0;
//        cout<<" r ="<<r<<endl;
//        return false;
//      }
    }

};


/*
    Copyright (C) 2008 Wei Dong <wdong@princeton.edu>. All Rights Reserved.

    This file is part of LSHKIT.

    LSHKIT is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    LSHKIT is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with LSHKIT.  If not, see <http://www.gnu.org/licenses/>.
*/

/**
  * \file lsh-run.cpp
  * \brief Example of using LSH Forest index for L2 distance.
  *
  * This program is an example of using LSH Forest index.  The LSH used is
  * LSB<GaussianLsh>, so the trees in the forest are binary trees.
  *
  * The current implementation of LSH Forest doesn't allow read/write
  * index files.
  *
\verbatim
Allowed options:
  -h [ --help ]                   produce help message.
  -c [ -- ] arg (=20)             # points to scan from each tree
  -L [ -- ] arg (=1)              number of trees
  -D [ -- ] arg (=10)             maximal depth of tree
  -W [ -- ] arg (=1)              hash function window size
  -Q [ -- ] arg (=100)            number of queries to use
  -K [ -- ] arg (=50)             number of nearest neighbors to retrieve
  -R [ -- ] arg (=3.40282347e+38) R-NN distance range
  -D [ --data ] arg               dataset path
  -B [ --benchmark ] arg
\endverbatim
  *
  */

#include <boost/program_options.hpp>
#include <boost/progress.hpp>
#include <boost/format.hpp>
#include <boost/timer.hpp>
#include <lshkit.h>

using namespace std;
using namespace lshkit;
namespace po = boost::program_options;

int main (int argc, char *argv[])
{
    string data_file;
    string benchmark;

    float R, W;
    unsigned c, L, H;
    unsigned Q, K;
    bool do_benchmark = true;
    // bool use_index = false; // load the index from a file

    boost::timer timer;

    po::options_description desc("Allowed options");
    desc.add_options()
        ("help,h", "produce help message.")
        (",c", po::value<unsigned>(&c)->default_value(20), "# points to scan from each tree")
        (",L", po::value<unsigned>(&L)->default_value(1), "number of trees")
        (",H", po::value<unsigned>(&H)->default_value(10), "maximal depth of tree")
        (",W", po::value<float>(&W)->default_value(1.0), "hash function window size")
        (",Q", po::value<unsigned>(&Q)->default_value(100), "number of queries to use")
        (",K", po::value<unsigned>(&K)->default_value(50), "number of nearest neighbors to retrieve")
        (",R", po::value<float>(&R)->default_value(numeric_limits<float>::max()), "R-NN distance range")
        ("data,D", po::value<string>(&data_file), "dataset path")
        ("benchmark,B", po::value<string>(&benchmark), "benchmark path")
        // ("index", po::value<string>(&index_file), "index file")
        ;

    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);

    if (vm.count("help") || (vm.count("data") < 1))
    {
        cout << desc;
        return 0;
    }

    if ((Q == 0) || (vm.count("benchmark") == 0)) {
        do_benchmark = false;
    }

    /*
    if (vm.count("index") == 1) {
        use_index = true;
    }
    */

    cout << "LOADING DATA..." << endl;
    timer.restart();
    FloatMatrix data(data_file);
    cout << boost::format("LOAD TIME: %1%s.") % timer.elapsed() << endl;

    //typedef Tail<RepeatHash<CauchyLsh> > MyLsh;

    typedef LSB<GaussianLsh> MyLsh;
    typedef ForestIndex<MyLsh, unsigned> Index;

    FloatMatrix::Accessor accessor(data);
    metric::l2<float> l2(data.getDim());
    Index index;

    // bool index_loaded = false;

    /*
    if (use_index) {
        ifstream is(index_file.c_str(), ios_base::binary);
        if (is) {
            is.exceptions(ios_base::eofbit | ios_base::failbit | ios_base::badbit);
            cout << "LOADING INDEX..." << endl;
            timer.restart();
            index.load(is);
            verify(is);
            cout << boost::format("LOAD TIME: %1%s.") % timer.elapsed() << endl;
            index_loaded = true;
        }
    }

    if (!index_loaded) {
        // We define a short name for the MPLSH index.
        float min = numeric_limits<float>::max();
        float max = -numeric_limits<float>::max();

        for (unsigned i = 0; i < data.getSize(); ++i) {
            for (unsigned j = 0; j < data.getDim(); ++j) {
                if (data[i][j] > max) max = data[i][j];
                if (data[i][j] < min) min = data[i][j];
            }
        }

        */

        Index::Parameter param;

        // Setup the parameters.  Note that L is not provided here.
        param.W = W;
        param.dim = data.getDim();
        DefaultRng rng;

        index.init(param, rng, L, H);
        // The accessor.

        // Initialize the index structure.  Note L is passed here.
        cout << "CONSTRUCTING INDEX..." << endl;

        timer.restart();
        {
            boost::progress_display progress(data.getSize());
            for (unsigned i = 0; i < data.getSize(); ++i)
            {
                // Insert an item to the hash table.
                // Note that only the key is passed in here.
                // MPLSH will get the feature from the accessor.
                index.insert(i, accessor);
                ++progress;
            }
        }
        cout << boost::format("CONSTRUCTION TIME: %1%s.") % timer.elapsed() << endl;

        /*
        if (use_index) {
            timer.restart();
            cout << "SAVING INDEX..." << endl;
            {
                ofstream os(index_file.c_str(), ios_base::binary);
                os.exceptions(ios_base::eofbit | ios_base::failbit | ios_base::badbit);
                index.save(os);
                verify(os);
            }
            cout << boost::format("SAVING TIME: %1%s") % timer.elapsed() << endl;
        }
    }
    */

    if (do_benchmark) {

        Benchmark<> bench;
        cout << "LOADING BENCHMARK..." << endl;
        bench.load(benchmark);
        bench.resize(Q, K);
        cout << "DONE." << endl;

        for (unsigned i = 0; i < Q; ++i)
        {
            for (unsigned j = 0; j < K; ++j)
            {
                assert(bench.getAnswer(i)[j].key < data.getSize());
            }
        }

        cout << "RUNNING QUERIES..." << endl;

        Stat recall;
        Stat cost;

        timer.restart();
        {
            TopkScanner<FloatMatrix::Accessor, metric::l2<float> > query(accessor, l2, K, R);
            boost::progress_display progress(Q);
            for (unsigned i = 0; i < Q; ++i)
            {
                query.reset(data[bench.getQuery(i)]);
                index.query(data[bench.getQuery(i)], c * L, query);
                recall << bench.getAnswer(i).recall(query.topk());
                cost << double(query.cnt())/double(data.getSize());
                ++progress;
            }
        }
        cout << boost::format("QUERY TIME: %1%s.") % timer.elapsed() << endl;

        cout << "[RECALL] " << recall.getAvg() << " +/- " << recall.getStd() << endl;
        cout << "[COST] " << cost.getAvg() << " +/- " << cost.getStd() << endl;

    }

    return 0;
}


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

#define BOOST_TEST_MAIN
#define BOOST_TEST_MODULE CudaMatrix_test

// Standard includes
#include <string>

// Third party includes
#include <boost/test/unit_test.hpp>

// Local VOTCA includes
#include "votca/xtp/cudapipeline.h"
#include "votca/xtp/eigen.h"

using namespace votca::xtp;

BOOST_AUTO_TEST_SUITE(CudaMatrix_test)

BOOST_AUTO_TEST_CASE(create_cudamatrix) {
  // Call the class to handle GPU resources
  CudaPipeline cp(0);

  // Call matrix multiplication GPU
  Eigen::MatrixXd B = Eigen::MatrixXd::Random(5, 7);

  // Copy matrix back and for to the GPU
  CudaMatrix cumatrix{B, cp.get_stream()};
  Eigen::MatrixXd tmp = cumatrix;

  // Expected results
  BOOST_TEST(B.isApprox(tmp));
}

BOOST_AUTO_TEST_CASE(create_cudamatrixblock) {
  // Call the class to handle GPU resources
  CudaPipeline cp(0);

  // Call matrix multiplication GPU
  Eigen::MatrixXd X = Eigen::MatrixXd::Random(10, 8);

  // Copy matrix back and for to the GPU
  CudaMatrix cumatrix{X.block(2, 4, 3, 1), cp.get_stream()};
  Eigen::MatrixXd tmp = cumatrix;

  // Expected results
  BOOST_TEST(X.block(2, 4, 3, 1).isApprox(tmp));
}

BOOST_AUTO_TEST_CASE(setZero) {
  // Call the class to handle GPU resources
  CudaPipeline cp(0);

  // Call matrix multiplication GPU
  Eigen::MatrixXd X = Eigen::MatrixXd::Random(10, 8);

  // Copy matrix back and for to the GPU
  CudaMatrix cumatrix{X, cp.get_stream()};
  cumatrix.setZero();
  Eigen::MatrixXd tmp = cumatrix;

  // Expected results
  BOOST_TEST(tmp.isApproxToConstant(0, 1e-14));
}

BOOST_AUTO_TEST_SUITE_END()