#include <cstdlib>
#include <algorithm>
#include <sstream>
#include <iostream>
#include <vector>
#include <limits>
#include <cmath>
#include <chrono>
#include <ctime>

#include <unordered_map>
#include <unordered_set>

#include <execinfo.h>
#include <unistd.h>
#include <signal.h>

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

#include "dynet/training.h"
#include "dynet/dynet.h"
#include "dynet/expr.h"
#include "dynet/nodes.h"
#include "dynet/lstm.h"
#include "dynet/rnn.h"
#include "dynet/io.h"
#include "c2.h"

cpyp::Corpus corpus;
volatile bool requested_stop = false;
unsigned LAYERS = 2;
unsigned INPUT_DIM = 40;
unsigned HIDDEN_DIM = 60;
unsigned ACTION_DIM = 36;
unsigned PRETRAINED_DIM = 50;
unsigned LSTM_INPUT_DIM = 60;
unsigned POS_DIM = 10;
unsigned REL_DIM = 8;

unsigned LSTM_CHAR_OUTPUT_DIM = 100; //Miguel
bool USE_SPELLING = false;

float DROPOUT = 0.0f;

bool USE_POS = false;

constexpr const char* ROOT_SYMBOL = "ROOT";
unsigned kROOT_SYMBOL = 0;
unsigned ACTION_SIZE = 0;
unsigned VOCAB_SIZE = 0;
unsigned POS_SIZE = 0;

unsigned CHAR_SIZE = 255; //size of ascii chars... Miguel

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

vector<unsigned> possible_actions;
unordered_map<unsigned, vector<float>> pretrained;

void InitCommandLine(int argc, char** argv, po::variables_map* conf) {
  po::options_description opts("Configuration options");
  opts.add_options()
        ("training_data,T", po::value<string>(), "List of Transitions - Training corpus")
        ("dev_data,d", po::value<string>(), "Development corpus")
        ("test_data,p", po::value<string>(), "Test corpus")
        ("dropout,D", po::value<float>(), "Dropout rate")
        ("unk_strategy,o", po::value<unsigned>()->default_value(1), "Unknown word strategy: 1 = singletons become UNK with probability unk_prob")
        ("unk_prob,u", po::value<double>()->default_value(0.2), "Probably with which to replace singletons with UNK in training data")
        ("model,m", po::value<string>(), "Load saved model from this file")
        ("use_pos_tags,P", "make POS tags visible to parser")
        ("beam_size,b", po::value<unsigned>()->default_value(1), "beam size")
        ("layers", po::value<unsigned>()->default_value(2), "number of LSTM layers")
        ("action_dim", po::value<unsigned>()->default_value(16), "action embedding size")
        ("input_dim", po::value<unsigned>()->default_value(32), "input embedding size")
        ("hidden_dim", po::value<unsigned>()->default_value(64), "hidden dimension")
        ("pretrained_dim", po::value<unsigned>()->default_value(50), "pretrained input dimension")
        ("pos_dim", po::value<unsigned>()->default_value(12), "POS dimension")
        ("rel_dim", po::value<unsigned>()->default_value(10), "relation dimension")
        ("lstm_input_dim", po::value<unsigned>()->default_value(60), "LSTM input dimension")
        ("train,t", "Should training be run?")
        ("words,w", po::value<string>(), "Pretrained word embeddings")
        ("use_spelling,S", "Use spelling model") //Miguel. Spelling model
        ("help,h", "Help");
  po::options_description dcmdline_options;
  dcmdline_options.add(opts);
  po::store(parse_command_line(argc, argv, dcmdline_options), *conf);
  if (conf->count("help")) {
    cerr << dcmdline_options << endl;
    exit(1);
  }
  if (conf->count("training_data") == 0) {
    cerr << "Please specify --traing_data (-T): this is required to determine the vocabulary mapping, even if the parser is used in prediction mode.\n";
    exit(1);
  }
}

struct ParserBuilder {

  LSTMBuilder stack_lstm; // (layers, input, hidden, trainer)
  LSTMBuilder output_lstm; // (layers, input, hidden, trainer)
  LSTMBuilder buffer_lstm;
  LSTMBuilder action_lstm;

  LSTMBuilder ent_lstm_fwd;
  LSTMBuilder ent_lstm_rev;

  LookupParameter p_w; // word embeddings
  LookupParameter p_t; // pretrained word embeddings (not updated)
  LookupParameter p_a; // input action embeddings
  LookupParameter p_r; // relation embeddings
  LookupParameter p_p; // pos tag embeddings
  Parameter p_pbias; // parser state bias
  Parameter p_A; // action lstm to parser state
  Parameter p_B; // buffer lstm to parser state
  Parameter p_O; // output lstm to parser state

  Parameter p_S; // stack lstm to parser state
  Parameter p_H; // head matrix for composition function
  Parameter p_D; // dependency matrix for composition function
  Parameter p_R; // relation matrix for composition function
  Parameter p_w2l; // word to LSTM input
  Parameter p_p2l; // POS to LSTM input
  Parameter p_t2l; // pretrained word embeddings to LSTM input
  Parameter p_ib; // LSTM input bias
  Parameter p_cbias; // composition function bias
  Parameter p_p2a;   // parser state to action
  Parameter p_action_start;  // action bias
  Parameter p_abias;  // action bias
  Parameter p_buffer_guard;  // end of buffer
  Parameter p_stack_guard;  // end of stack
  Parameter p_output_guard;  // end of output buffer

  Parameter p_start_of_word;//Miguel -->dummy <s> symbol
  Parameter p_end_of_word; //Miguel --> dummy </s> symbol
  LookupParameter char_emb; //Miguel-> mapping of characters to vectors

  LSTMBuilder fw_char_lstm; // Miguel
  LSTMBuilder bw_char_lstm; //Miguel

  Parameter p_cW;

  explicit ParserBuilder(ParameterCollection & model, const unordered_map<unsigned, vector<float>>& pretrained) :
      stack_lstm(LAYERS, LSTM_INPUT_DIM, HIDDEN_DIM, model),
      output_lstm(LAYERS, LSTM_INPUT_DIM, HIDDEN_DIM, model),
      buffer_lstm(LAYERS, LSTM_INPUT_DIM, HIDDEN_DIM, model),
      action_lstm(LAYERS, ACTION_DIM, HIDDEN_DIM, model),
      ent_lstm_fwd(LAYERS, LSTM_INPUT_DIM, LSTM_INPUT_DIM, model),
      ent_lstm_rev(LAYERS, LSTM_INPUT_DIM, LSTM_INPUT_DIM, model),
      p_w(model.add_lookup_parameters(VOCAB_SIZE, {INPUT_DIM})),
      p_a(model.add_lookup_parameters(ACTION_SIZE, {ACTION_DIM})),
      p_r(model.add_lookup_parameters(ACTION_SIZE, {REL_DIM})),
      p_pbias(model.add_parameters({HIDDEN_DIM})),
      p_A(model.add_parameters({HIDDEN_DIM, HIDDEN_DIM})),
      p_B(model.add_parameters({HIDDEN_DIM, HIDDEN_DIM})),
      p_O(model.add_parameters({HIDDEN_DIM, HIDDEN_DIM})),
      p_S(model.add_parameters({HIDDEN_DIM, HIDDEN_DIM})),
      p_H(model.add_parameters({LSTM_INPUT_DIM, LSTM_INPUT_DIM})),
      p_D(model.add_parameters({LSTM_INPUT_DIM, LSTM_INPUT_DIM})),
      p_R(model.add_parameters({LSTM_INPUT_DIM, REL_DIM})),
      p_w2l(model.add_parameters({LSTM_INPUT_DIM, INPUT_DIM})),
      p_ib(model.add_parameters({LSTM_INPUT_DIM})),
      p_cbias(model.add_parameters({LSTM_INPUT_DIM})),
      p_p2a(model.add_parameters({ACTION_SIZE, HIDDEN_DIM})),
      p_action_start(model.add_parameters({ACTION_DIM})),
      p_abias(model.add_parameters({ACTION_SIZE})),

      p_buffer_guard(model.add_parameters({LSTM_INPUT_DIM})),
      p_stack_guard(model.add_parameters({LSTM_INPUT_DIM})),
      p_output_guard(model.add_parameters({LSTM_INPUT_DIM})),

      p_start_of_word(model.add_parameters({LSTM_INPUT_DIM})), //Miguel
      p_end_of_word(model.add_parameters({LSTM_INPUT_DIM})), //Miguel

      char_emb(model.add_lookup_parameters(CHAR_SIZE, {INPUT_DIM})),//Miguel

//      fw_char_lstm(LAYERS, LSTM_CHAR_OUTPUT_DIM, LSTM_INPUT_DIM, model), //Miguel
//      bw_char_lstm(LAYERS, LSTM_CHAR_OUTPUT_DIM, LSTM_INPUT_DIM,  model), //Miguel

      fw_char_lstm(LAYERS, LSTM_INPUT_DIM, LSTM_CHAR_OUTPUT_DIM/2, model), //Miguel
      bw_char_lstm(LAYERS, LSTM_INPUT_DIM, LSTM_CHAR_OUTPUT_DIM/2, model), /*Miguel*/
      p_cW(model.add_parameters({LSTM_INPUT_DIM, LSTM_INPUT_DIM * 2})) { //ner. {
    if (USE_POS) {
      p_p = model.add_lookup_parameters(POS_SIZE, {POS_DIM});
      p_p2l = model.add_parameters({LSTM_INPUT_DIM, POS_DIM});
    }
    if (pretrained.size() > 0) {
      p_t = model.add_lookup_parameters(VOCAB_SIZE, {PRETRAINED_DIM});
      for (auto it : pretrained)
        p_t.initialize(it.first, it.second);
      p_t2l = model.add_parameters({LSTM_INPUT_DIM, PRETRAINED_DIM});
    } else {
      p_t = nullptr;
      p_t2l = nullptr;
    }
  }

static bool IsActionForbidden(const string& a, unsigned bsize, unsigned ssize, vector<int> stacki) {

  bool is_shift = (a[0] == 'S');  //MIGUEL
  bool is_reduce = (a[0] == 'R');
  bool is_output = (a[0] == 'O');
//  std::cout<<"is red:"<<is_reduce<<"\n";
  if (is_shift && bsize == 1) return true;
  if (is_reduce && ssize ==1) return true;
  if (is_output && bsize ==1) return true;

  return false;
}

/*static bool IsActionForbidden(const string& a, unsigned bsize, unsigned ssize) {
  bool is_shift = (a[0] == 'S');
  bool is_reduce = !is_shift;
  if (is_shift && bsize == 1) return true;
  if (is_reduce && ssize < 3) return true;
  if (bsize == 2 && // ROOT is the only thing remaining on buffer
      ssize > 2 && // there is more than a single element on the stack
      is_shift) return true;
  // only attach left to ROOT
  if (bsize == 1 && ssize == 3 && a[0] == 'R') return true;
  return false;
}*/

map<int,string> compute_ents(unsigned sent_len, const vector<unsigned>& actions, const vector<string>& setOfActions, map<int,string>* pr = nullptr) {
  map<int,string> r;
  map<int,string>& rels = (pr ? *pr : r);
  for(unsigned i=0;i<sent_len;i++) { rels[i]="ERROR"; }
  vector<int> bufferi(sent_len + 1, 0), stacki(1, -999), outputi(1, -999);
  for (unsigned i = 0; i < sent_len; ++i)
    bufferi[sent_len - i] = i;
  bufferi[0] = -999;

  for (auto action: actions) { // loop over transitions for sentence
    const string& actionString=setOfActions[action];
   // std::cout<<"int"<<action<<"-"<<"actionString"<<actionString<<"\n";
    const char ac = actionString[0];
    if (ac =='S') {  // SHIFT
      assert(bufferi.size() > 1); // dummy symbol means > 1 (not >= 1)
      stacki.push_back(bufferi.back());
      bufferi.pop_back();
    }

    else if (ac=='R') { // REDUCE
      assert(stacki.size() > 1); // dummy symbol means > 2 (not >= 2)
      while(stacki.size()>1) {
	rels[stacki.back()] = actionString;
        outputi.push_back(stacki.back());
        stacki.pop_back();
      }
    }
    else if (ac =='O') {
       assert(bufferi.size() > 1); // dummy symbol means > 1 (not >= 1)
       outputi.push_back(bufferi.back());
       rels[bufferi.back()] = "0";
       bufferi.pop_back();

    }
  }
  assert(bufferi.size() == 1);
  //assert(stacki.size() == 2);
  return rels;
}

// given the first character of a UTF8 block, find out how wide it is
// see http://en.wikipedia.org/wiki/UTF-8 for more info
inline unsigned int UTF8Len(unsigned char x) {
  if (x < 0x80) return 1;
  else if ((x >> 5) == 0x06) return 2;
  else if ((x >> 4) == 0x0e) return 3;
  else if ((x >> 3) == 0x1e) return 4;
  else if ((x >> 2) == 0x3e) return 5;
  else if ((x >> 1) == 0x7e) return 6;
  else return 0;
}

// *** if correct_actions is empty, this runs greedy decoding ***
// returns parse actions for input sentence (in training just returns the reference)
// OOV handling: raw_sent will have the actual words
//               sent will have words replaced by appropriate UNK tokens
// this lets us use pretrained embeddings, when available, for words that were OOV in the
// parser training data
pair<vector<unsigned>, Expression> log_prob_parser(ComputationGraph* hg,
                                    const vector<unsigned>& raw_sent,  // raw sentence
                                    const vector<unsigned>& sent,  // sent with oovs replaced
                                    const vector<unsigned>& sentPos,
                                    const vector<unsigned>& correct_actions,
                                    const vector<string>& setOfActions,
                                    const map<unsigned, std::string>& intToWords,
                                    bool is_evaluation,
                                    double *right) {
  //for (unsigned i = 0; i < sent.size(); ++i) cerr << ' ' << intToWords.find(sent[i])->second;
  //cerr << endl;
    vector<unsigned> results;
    const bool build_training_graph = correct_actions.size() > 0;
    //std::cout<<"****************"<<"\n";
    bool apply_dropout = (DROPOUT && !is_evaluation);

    stack_lstm.new_graph(*hg);
    buffer_lstm.new_graph(*hg);
    output_lstm.new_graph(*hg);
    action_lstm.new_graph(*hg);

    ent_lstm_fwd.new_graph(*hg);
    ent_lstm_rev.new_graph(*hg);

  if (apply_dropout) {
      stack_lstm.set_dropout(DROPOUT);
      action_lstm.set_dropout(DROPOUT);
      buffer_lstm.set_dropout(DROPOUT);
      ent_lstm_fwd.set_dropout(DROPOUT);
      ent_lstm_rev.set_dropout(DROPOUT);
    } else {
      stack_lstm.disable_dropout();
      action_lstm.disable_dropout();
      buffer_lstm.disable_dropout();
      ent_lstm_fwd.disable_dropout();
      ent_lstm_rev.disable_dropout();
    }

    stack_lstm.start_new_sequence();
    buffer_lstm.start_new_sequence();
    output_lstm.start_new_sequence();
    action_lstm.start_new_sequence();
    // variables in the computation graph representing the parameters
    Expression pbias = parameter(*hg, p_pbias);
    // Expression H = parameter(*hg, p_H);
    // Expression D = parameter(*hg, p_D);
    Expression R = parameter(*hg, p_R);
    Expression cbias = parameter(*hg, p_cbias);
    Expression S = parameter(*hg, p_S);
    Expression B = parameter(*hg, p_B);
    Expression O = parameter(*hg, p_O);

    Expression A = parameter(*hg, p_A);
    Expression ib = parameter(*hg, p_ib);
    Expression w2l = parameter(*hg, p_w2l);
    Expression p2l;
    if (USE_POS)
      p2l = parameter(*hg, p_p2l);
    Expression t2l;
    if (p_t2l.p)
      t2l = parameter(*hg, p_t2l);
    Expression p2a = parameter(*hg, p_p2a);
    Expression abias = parameter(*hg, p_abias);
    Expression action_start = parameter(*hg, p_action_start);

    action_lstm.add_input(action_start);

    Expression cW = parameter(*hg, p_cW);

    vector<Expression> buffer(sent.size() + 1);  // variables representing word embeddings (possibly including POS info)
    vector<int> bufferi(sent.size() + 1);  // position of the words in the sentence
    // precompute buffer representation from left to right

    Expression word_end = parameter(*hg, p_end_of_word); //Miguel
    Expression word_start = parameter(*hg, p_start_of_word); //Miguel

    if (USE_SPELLING){
       fw_char_lstm.new_graph(*hg);
        //    fw_char_lstm.add_parameter_edges(hg);

       bw_char_lstm.new_graph(*hg);
       //    bw_char_lstm.add_parameter_edges(hg);
    }

    for (unsigned i = 0; i < sent.size(); ++i) {
      assert(sent[i] < VOCAB_SIZE);
      //Expression w = lookup(*hg, p_w, sent[i]);

      unsigned wi=sent[i];
      std::string ww=intToWords.at(wi);
      Expression w;
      /**********SPELLING MODEL*****************/
      if (USE_SPELLING) {
        //std::cout<<"using spelling"<<"\n";
        if (ww.length()==4  && ww[0]=='R' && ww[1]=='O' && ww[2]=='O' && ww[3]=='T'){
          w=lookup(*hg, p_w, sent[i]); //we do not need a LSTM encoding for the root word, so we put it directly-.
        }
        else {

            fw_char_lstm.start_new_sequence();
            //cerr<<"start_new_sequence done"<<"\n";

            fw_char_lstm.add_input(word_start);
            //cerr<<"added start of word symbol"<<"\n";
            /*for (unsigned j=0;j<w.length();j++){

                //cerr<<j<<":"<<w[j]<<"\n";
                Expression cj=lookup(*hg, char_emb, w[j]);
                fw_char_lstm.add_input(cj, hg);

               //std::cout<<"Inputdim:"<<LSTM_INPUT_DIM<<"\n";
               //hg->incremental_forward();

            }*/
	    std::vector<int> strevbuffer;
            for (unsigned j=0;j<ww.length();j+=UTF8Len(ww[j])){

                //cerr<<j<<":"<<w[j]<<"\n";
                std::string wj;
                for (unsigned h=j;h<j+UTF8Len(ww[j]);h++) wj+=ww[h];
                //std::cout<<"fw"<<wj<<"\n";
                int wjint=corpus.charsToInt[wj];
		//std::cout<<"fw:"<<wjint<<"\n";
		strevbuffer.push_back(wjint);
                Expression cj=lookup(*hg, char_emb, wjint);
                fw_char_lstm.add_input(cj);

               //std::cout<<"Inputdim:"<<LSTM_INPUT_DIM<<"\n";
               //hg->incremental_forward();

            }
            fw_char_lstm.add_input(word_end);
            //cerr<<"added end of word symbol"<<"\n";

            Expression fw_i=fw_char_lstm.back();

            //cerr<<"fw_char_lstm.back() done"<<"\n";

            bw_char_lstm.start_new_sequence();
            //cerr<<"bw start new sequence done"<<"\n";

            bw_char_lstm.add_input(word_end);
	    //for (unsigned j=w.length()-1;j>=0;j--){
            /*for (unsigned j=w.length();j-->0;){
               //cerr<<j<<":"<<w[j]<<"\n";
               Expression cj=lookup(*hg, char_emb, w[j]);
               bw_char_lstm.add_input(cj);
            }*/

	    while(!strevbuffer.empty()) {
		int wjint=strevbuffer.back();
		//std::cout<<"bw:"<<wjint<<"\n";
		Expression cj=lookup(*hg, char_emb, wjint);
                bw_char_lstm.add_input(cj);
		strevbuffer.pop_back();
	    }

            /*for (unsigned j=w.length()-1;j>0;j=j-UTF8Len(w[j])) {

                //cerr<<j<<":"<<w[j]<<"\n";
                std::string wj;
                for (unsigned h=j;h<j+UTF8Len(w[j]);h++) wj+=w[h];
                std::cout<<"bw"<<wj<<"\n";
                int wjint=corpus.charsToInt[wj];
                Expression cj=lookup(*hg, char_emb, wjint);
                bw_char_lstm.add_input(cj);

               //std::cout<<"Inputdim:"<<LSTM_INPUT_DIM<<"\n";
               //hg->incremental_forward();

            }*/
            bw_char_lstm.add_input(word_start);
            //cerr<<"start symbol in bw seq"<<"\n";

            Expression bw_i=bw_char_lstm.back();

            vector<Expression> tt = {fw_i, bw_i};
            w=concatenate(tt); //and this goes into the buffer...
            //cerr<<"fw and bw done"<<"\n";
         }

	}
      /**************************************************/
      //cerr<<"concatenate?"<<"\n";

      /***************NO SPELLING*************************************/

      // Expression w = lookup(*hg, p_w, sent[i]);
      else { //NO SPELLING
          //Don't use SPELLING
          //std::cout<<"don't use spelling"<<"\n";
          w=lookup(*hg, p_w, sent[i]);
      }

      Expression i_i;
      if (USE_POS) {
        Expression p = lookup(*hg, p_p, sentPos[i]);
        i_i = affine_transform({ib, w2l, w, p2l, p});
      } else {
        i_i = affine_transform({ib, w2l, w});
      }
      if (p_t.p && pretrained.count(raw_sent[i])) {
        Expression t = const_lookup(*hg, p_t, raw_sent[i]);
        i_i = affine_transform({i_i, t2l, t});
      }
      buffer[sent.size() - i] = rectify(i_i);
      bufferi[sent.size() - i] = i;
    }
    // dummy symbol to represent the empty buffer
    buffer[0] = parameter(*hg, p_buffer_guard);
    bufferi[0] = -999;
    for (auto& b : buffer)
      buffer_lstm.add_input(b);

    vector<Expression> stack;  // variables representing subtree embeddings
    vector<int> stacki; // position of words in the sentence of head of subtree
    stack.push_back(parameter(*hg, p_stack_guard));
    stacki.push_back(-999); // not used for anything

    vector<Expression> output;  // variables representing subtree embeddings
    vector<int> outputi;
    output.push_back(parameter(*hg, p_output_guard));
    outputi.push_back(-999); // not used for anything
    // drive dummy symbol on stack through LSTM
    stack_lstm.add_input(stack.back());
    output_lstm.add_input(output.back());
    vector<Expression> log_probs;
    string rootword;
    unsigned action_count = 0;  // incremented at each prediction
    while(buffer.size() > 1 || stack.size()>1) {

      // get list of possible actions for the current parser state
      vector<unsigned> current_valid_actions;
      for (auto a: possible_actions) {
        if (IsActionForbidden(setOfActions[a], buffer.size(), stack.size(), stacki))
          continue;
        current_valid_actions.push_back(a);
      }

      // p_t = pbias + S * slstm + B * blstm + A * almst
      Expression p_t = affine_transform({pbias, O, output_lstm.back(), S, stack_lstm.back(), B, buffer_lstm.back(), A, action_lstm.back()});
      Expression nlp_t = rectify(p_t);
      // r_t = abias + p2a * nlp
      Expression r_t = affine_transform({abias, p2a, nlp_t});

      // adist = log_softmax(r_t, current_valid_actions)
      Expression adiste = log_softmax(r_t, current_valid_actions);
      vector<float> adist = as_vector(hg->incremental_forward(adiste));
      double best_score = adist[current_valid_actions[0]];
      unsigned best_a = current_valid_actions[0];
      for (unsigned i = 1; i < current_valid_actions.size(); ++i) {
        if (adist[current_valid_actions[i]] > best_score) {
          best_score = adist[current_valid_actions[i]];
          best_a = current_valid_actions[i];
        }
      }
      unsigned action = best_a;
      if (build_training_graph) {  // if we have reference actions (for training) use the reference action
        action = correct_actions[action_count];
        if (best_a == action) { (*right)++; }
      }
      ++action_count;
      // action_log_prob = pick(adist, action)
      log_probs.push_back(pick(adiste, action));
      results.push_back(action);
      //std::cout<<"action:"<<action<<"\n";

      // add current action to action LSTM
      Expression actione = lookup(*hg, p_a, action);
      action_lstm.add_input(actione);

      // get relation embedding from action (TODO: convert to relation from action?)
      Expression relation = lookup(*hg, p_r, action);

      // do action
      const string& actionString=setOfActions[action];
      //cerr << "A=" << actionString << " Bsize=" << buffer.size() << " Ssize=" << stack.size() << endl;
      const char ac = actionString[0];

      //std::cout<<"buffer-size:"<<bufferi.size()<<"\n";
      //std::cout<<"stack-size:"<<stacki.size()<<"\n";
      //std::cout<<"output-size:"<<outputi.size()<<"\n";
      //std::cout<<ac<<"\n";
      if (ac =='S') {  // SHIFT
         assert(buffer.size() > 1); // dummy symbol means > 1 (not >= 1)
          stack.push_back(buffer.back());
          stack_lstm.add_input(buffer.back());
          stacki.push_back(bufferi.back());
          buffer.pop_back();
          buffer_lstm.rewind_one_step();
          bufferi.pop_back();
       }
       else if (ac=='R') { // REDUCE
	   Expression previous;
           Expression comp;
           vector<Expression> entities(stacki.size());
           ent_lstm_fwd.start_new_sequence();
           ent_lstm_rev.start_new_sequence();
           for (unsigned i = 0; i < stacki.size(); ++i) {
              ent_lstm_fwd.add_input(stack[i]);
              ent_lstm_rev.add_input(stack[stacki.size() - i - 1]);
           }
           while(stacki.size()>1) {
              outputi.push_back(stacki.back());
              stack_lstm.rewind_one_step();
              stack.pop_back();
              stacki.pop_back();
            //COMPOSITION FUNCTION!! ??
            //Expression composed = affine_transform({cbias, H, head, D, dep, R, relation});
            //Expression nlcomposed = tanh(composed);
           }
           Expression efwd = ent_lstm_fwd.back();
           Expression erev = ent_lstm_rev.back();
           if (apply_dropout) {
             efwd = dropout(efwd, DROPOUT);
             erev = dropout(erev, DROPOUT);
           }
           Expression c = concatenate({efwd, erev});
           //Expression c = concatenate({ent_lstm_fwd.back(), ent_lstm_rev.back()});
           Expression composed = rectify(affine_transform({cbias, cW, c, R, relation}));
	   output.push_back(composed);
           output_lstm.add_input(composed);

        }
        else if (ac =='O') {
           assert(bufferi.size() > 1); // dummy symbol means > 1 (not >= 1)
           outputi.push_back(bufferi.back());
           output.push_back(buffer.back());
           output_lstm.add_input(buffer.back());
           buffer.pop_back();
           bufferi.pop_back();
           buffer_lstm.rewind_one_step();

        }

    }
    assert(stack.size() == 1); // guard symbol, root
    assert(stacki.size() == 1);
    assert(buffer.size() == 1); // guard symbol
    assert(bufferi.size() == 1);
    Expression tot_neglogprob = -sum(log_probs);
    assert(tot_neglogprob.pg != nullptr);
    return std::make_pair(results, tot_neglogprob);
  }

};

void signal_callback_handler(int /* signum */) {
  if (requested_stop) {
    cerr << "\nReceived SIGINT again, quitting.\n";
    _exit(1);
  }
  cerr << "\nReceived SIGINT terminating optimization early...\n";
  requested_stop = true;
}

unsigned compute_correct(const map<int,string>& ref, const map<int,string>& hyp, unsigned len) {
  unsigned res = 0;
  for (unsigned i = 0; i < len; ++i) {
    auto ri = ref.find(i);
    auto hi = hyp.find(i);
    assert(ri != ref.end());
    assert(hi != hyp.end());
    if (ri->second.compare(hi->second)==0) ++res;
  }
  return res;
}

void output_conll(const vector<unsigned>& sentence, const vector<unsigned>& pos,
                  const vector<string>& sentenceUnkStrings,
                  const map<unsigned, string>& intToWords,
                  const map<unsigned, string>& intToPos,
                  const map<int,string>& rel_ref,
                  const map<int,string>& rel_hyp) {
  for (unsigned i = 0; i < (sentence.size()); ++i) {
    // auto index = i + 1;
    assert(i < sentenceUnkStrings.size() &&
           ((sentence[i] == corpus.get_or_add_word(cpyp::Corpus::UNK) &&
             sentenceUnkStrings[i].size() > 0) ||
            (sentence[i] != corpus.get_or_add_word(cpyp::Corpus::UNK) &&
             sentenceUnkStrings[i].size() == 0 &&
             intToWords.find(sentence[i]) != intToWords.end())));
    string wit = (sentenceUnkStrings[i].size() > 0)?
      sentenceUnkStrings[i] : intToWords.find(sentence[i])->second;
    auto pit = intToPos.find(pos[i]);
    //assert(hyp.find(i) != hyp.end());
    //auto hyp_head = hyp.find(i)->second + 1;
    //if (hyp_head == (int)sentence.size()) hyp_head = 0;
    auto hyp_rel_it = rel_hyp.find(i);
    auto ref_rel_it = rel_ref.find(i);
    assert(hyp_rel_it != rel_hyp.end());
    auto hyp_rel = hyp_rel_it->second;
    auto ref_rel = ref_rel_it->second;
    size_t first_char_in_rel = hyp_rel.find('(') + 1;
    size_t last_char_in_rel = hyp_rel.rfind(')') - 1;
    size_t first_char_in_ref = ref_rel.find('(') + 1;
    size_t last_char_in_ref = ref_rel.rfind(')') - 1;

    hyp_rel = hyp_rel.substr(first_char_in_rel, last_char_in_rel - first_char_in_rel + 1);
    ref_rel = ref_rel.substr(first_char_in_ref, last_char_in_ref - first_char_in_ref + 1);
    if (hyp_rel.compare("0")!=0){
        hyp_rel="I-"+hyp_rel;
    }
    else hyp_rel="O";
    if (ref_rel.compare("0")!=0){
        ref_rel="I-"+ref_rel;
    }
    else ref_rel="O";

    //cout << index << '\t'       // 1. ID
    cout << wit << ' '        // 2. FORM
    //      << "_" << '\t'         // 3. LEMMA
    //     << "_" << '\t'         // 4. CPOSTAG
         << pit->second << ' ' // 5. POSTAG
         << "_" << ' '         // 6. tree. _ empty?
   //      << hyp_head << '\t'    // 7. HEAD
         << ref_rel << ' '     // 8. DEPREL
         << hyp_rel << endl;     // 8. DEPREL
//         << "_" << '\t'         // 9. PHEAD
 //        << "_" << endl;        // 10. PDEPREL
  }
  cout << endl;
}

int main(int argc, char** argv) {
  dynet::initialize(argc, argv);

  cerr << "COMMAND:";
  for (unsigned i = 0; i < static_cast<unsigned>(argc); ++i) cerr << ' ' << argv[i];
  cerr << endl;
  unsigned status_every_i_iterations = 100;

  po::variables_map conf;
  InitCommandLine(argc, argv, &conf);
  USE_POS = conf.count("use_pos_tags");
  if (conf.count("dropout"))
    DROPOUT = conf["dropout"].as<float>();

  USE_SPELLING=conf.count("use_spelling"); //Miguel
  corpus.USE_SPELLING=USE_SPELLING;

  LAYERS = conf["layers"].as<unsigned>();
  INPUT_DIM = conf["input_dim"].as<unsigned>();
  PRETRAINED_DIM = conf["pretrained_dim"].as<unsigned>();
  HIDDEN_DIM = conf["hidden_dim"].as<unsigned>();
  ACTION_DIM = conf["action_dim"].as<unsigned>();
  LSTM_INPUT_DIM = conf["lstm_input_dim"].as<unsigned>();
  POS_DIM = conf["pos_dim"].as<unsigned>();
  REL_DIM = conf["rel_dim"].as<unsigned>();
  // const unsigned beam_size = conf["beam_size"].as<unsigned>();
  const unsigned unk_strategy = conf["unk_strategy"].as<unsigned>();
  cerr << "Unknown word strategy: ";
  if (unk_strategy == 1) {
    cerr << "STOCHASTIC REPLACEMENT\n";
  } else {
    abort();
  }
  const double unk_prob = conf["unk_prob"].as<double>();
  assert(unk_prob >= 0.); assert(unk_prob <= 1.);
  ostringstream os;
  os << "parser_" << (USE_POS ? "pos" : "nopos")
     << '_' << LAYERS
     << '_' << INPUT_DIM
     << '_' << HIDDEN_DIM
     << '_' << ACTION_DIM
     << '_' << LSTM_INPUT_DIM
     << '_' << POS_DIM
     << '_' << REL_DIM
     << "-pid" << getpid() << ".params";
  // int best_correct_heads = 0;
  double best_f1_score=-1.0;
  const string fname = os.str();
  cerr << "Writing parameters to file: " << fname << endl;
  bool softlinkCreated = false;
  corpus.load_correct_actions(conf["training_data"].as<string>());
  const unsigned kUNK = corpus.get_or_add_word(cpyp::Corpus::UNK);
  kROOT_SYMBOL = corpus.get_or_add_word(ROOT_SYMBOL);

  if (conf.count("words")) {
    pretrained[kUNK] = vector<float>(PRETRAINED_DIM, 0);
    cerr << "Loading from " << conf["words"].as<string>() << " with" << PRETRAINED_DIM << " dimensions\n";
    ifstream in(conf["words"].as<string>().c_str());
    string line;
    getline(in, line);
    vector<float> v(PRETRAINED_DIM, 0);
    string word;
    while (getline(in, line)) {
      istringstream lin(line);
      lin >> word;
      for (unsigned i = 0; i < PRETRAINED_DIM; ++i) lin >> v[i];
      unsigned id = corpus.get_or_add_word(word);
      pretrained[id] = v;
    }
  }

  set<unsigned> training_vocab; // words available in the training corpus
  set<unsigned> singletons;
  {  // compute the singletons in the parser's training data
    map<unsigned, unsigned> counts;
    for (auto sent : corpus.sentences)
      for (auto word : sent.second) { training_vocab.insert(word); counts[word]++; }
    for (auto wc : counts)
      if (wc.second == 1) singletons.insert(wc.first);
  }

  // correct ner label set (not action)
  set<string> label;
  map<string,int> label2id;
  for (unsigned i = 0; i < corpus.actions.size(); ++i) {
      if (corpus.actions[i][0] == 'R') {
          label2id[corpus.actions[i]] = label.size();

          cerr << corpus.actions[i] << " :: "<< label.size() << endl;
          label.insert(corpus.actions[i]);
      }
  }

  cerr << "Number of words: " << corpus.nwords << endl;
  VOCAB_SIZE = corpus.nwords + 1;

  cerr << "Number of UTF8 chars: " << corpus.maxChars << endl;
  if (corpus.maxChars>255) CHAR_SIZE=corpus.maxChars;

  ACTION_SIZE = corpus.nactions + 1;
  //POS_SIZE = corpus.npos + 1;
  POS_SIZE = corpus.npos + 10;
  possible_actions.resize(corpus.nactions);
  for (unsigned i = 0; i < corpus.nactions; ++i)
    possible_actions[i] = i;

  ParameterCollection model;
  ParserBuilder parser(model, pretrained);
  if (conf.count("model")) {
    TextFileLoader loader(conf["model"].as<string>());
    loader.populate(model);
  }

  // OOV words will be replaced by UNK tokens
  corpus.load_correct_actionsDev(conf["dev_data"].as<string>());
  if (USE_SPELLING) VOCAB_SIZE = corpus.nwords + 1;
  //TRAINING
  if (conf.count("train")) {
    signal(SIGINT, signal_callback_handler);
    SimpleSGDTrainer sgd(model);
    //MomentumSGDTrainer sgd(&model);
    float eta_decay = 0.08;
    //sgd.eta_decay = 0.05;
    cerr << "Training started."<<"\n";
    vector<unsigned> order(corpus.nsentences);
    for (unsigned i = 0; i < corpus.nsentences; ++i)
      order[i] = i;
    double tot_seen = 0;
    status_every_i_iterations = min(status_every_i_iterations, corpus.nsentences);
    unsigned si = corpus.nsentences;
    cerr << "NUMBER OF TRAINING SENTENCES: " << corpus.nsentences << endl;
    unsigned trs = 0;
    double right = 0;
    double llh = 0;
    bool first = true;
    int iter = -1;
    while(!requested_stop) {
      ++iter;
      for (unsigned sii = 0; sii < status_every_i_iterations; ++sii) {
           if (si == corpus.nsentences) {
             si = 0;
             if (first) { first = false; } else { sgd.learning_rate *= 1 - eta_decay; }
             cerr << "**SHUFFLE\n";
             random_shuffle(order.begin(), order.end());
           }
           tot_seen += 1;
           const vector<unsigned>& sentence=corpus.sentences[order[si]];
           vector<unsigned> tsentence=sentence;
           if (unk_strategy == 1) {
             for (auto& w : tsentence)
               if (singletons.count(w) && dynet::rand01() < unk_prob) w = kUNK;
           }
	   const vector<unsigned>& sentencePos=corpus.sentencesPos[order[si]];
	   const vector<unsigned>& actions=corpus.correct_act_sent[order[si]];
           ComputationGraph hg;
           auto pred_loss = parser.log_prob_parser(&hg,sentence,tsentence,sentencePos,actions,corpus.actions,corpus.intToWords,false,&right);
           double lp = as_scalar(hg.incremental_forward(pred_loss.second));
           if (lp < 0) {
             cerr << "Log prob < 0 on sentence " << order[si] << ": lp=" << lp << endl;
             assert(lp >= 0.0);
           }
           hg.backward(pred_loss.second);
           sgd.update();
           llh += lp;
           ++si;
           trs += actions.size();
      }
      sgd.status();
      cerr << "update #" << iter << " (epoch " << (tot_seen / corpus.nsentences) << ")\tllh: "<< llh<<" ppl: " << exp(llh / trs) << " err: " << (trs - right) / trs << endl;
      llh = trs = right = 0;

      static int logc = 0;
      ++logc;
      if (logc % 25 == 1) { // report on dev set
        unsigned dev_size = corpus.nsentencesDev;
        // dev_size = 100;
        double llh = 0;
        double trs = 0;
        double right = 0;
        double correct_heads = 0;
        double total_heads = 0;

        // TODO fix it to variable length label
        int confusion[4][4]= {{0}};
        int class_tp[4] = {0};
        int class_fp[4] = {0};
        int class_fn[4] = {0};

        auto t_start = std::chrono::high_resolution_clock::now();
        for (unsigned sii = 0; sii < dev_size; ++sii) {
           const vector<unsigned>& sentence=corpus.sentencesDev[sii];
    	   const vector<unsigned>& sentencePos=corpus.sentencesPosDev[sii];
	       const vector<unsigned>& actions=corpus.correct_act_sentDev[sii];
           vector<unsigned> tsentence=sentence;
        if (!USE_SPELLING) {
                 for (auto& w : tsentence)
                     if (training_vocab.count(w) == 0) w = kUNK;
         }

        ComputationGraph hg;
        auto pred_loss = parser.log_prob_parser(&hg,sentence,tsentence,sentencePos,vector<unsigned>(),corpus.actions,corpus.intToWords,true,&right);
        vector<unsigned> & pred = pred_loss.first;
        double lp = 0;
        //vector<unsigned> pred = parser.log_prob_parser_beam(&hg,sentence,sentencePos,corpus.actions,beam_size,&lp);
        llh -= lp;
        trs += actions.size();
        //for (int i=0;i<actions.size();i++) std::cout<<"a->"<<actions[i]<<"\n";
        //for (int i=0;i<actions.size();i++) std::cout<<"p->"<<pred[i]<<"\n";
        map<int,string> ref = parser.compute_ents(sentence.size(), actions, corpus.actions);
        map<int,string> hyp = parser.compute_ents(sentence.size(), pred, corpus.actions);

        //update confusion matrix
        for (unsigned i = 0; i < sentence.size()-1; ++i) {
          auto ri = ref.find(i);
          auto hi = hyp.find(i);
          assert(ri != ref.end());
          assert(hi != hyp.end());

          if (ri->second != "0"){
              int pr = label2id[hi->second];
              int tr = label2id[ri->second];

              confusion[pr][tr]++;
          }
        }

        correct_heads += compute_correct(ref, hyp, sentence.size() - 1);
        total_heads += sentence.size() - 1;
        }

        // compute tp, fp and fn for each class
        cerr << "confusion matrix" << endl;
        for (unsigned i = 0; i < label.size(); ++i){
            for (unsigned j = 0; j < label.size(); ++j){
                cerr << i << " " << j << " " << confusion[i][j] << endl;
                if (i == j){
                    class_tp[i] = confusion[i][j];
                } else {
                    class_fp[i] += confusion[i][j];
                    class_fn[i] += confusion[j][i];
                }
            }
        }

        // compute precision, recall, f1
        double global_f1_score = 0;
        for (unsigned i = 0; i < label.size(); ++i){
            double precision = (float)class_tp[i] / (class_tp[i] + class_fp[i]);
            double recall = (float)class_tp[i] / (class_tp[i] + class_fn[i]);
            double f1_score = 2 * (float)(precision * recall) / (precision + recall);
            double pr=precision+recall;
            if (pr==0) f1_score=0.0;
            global_f1_score += f1_score / label.size(); // Macro averaged F1
            if (std::isnan(global_f1_score)) global_f1_score=0.0;

            cerr << i << "class Precision " << precision << " Recall " << recall << " F1 " << f1_score << endl;
        }
        cerr << "F1 >> " << global_f1_score << endl;

        auto t_end = std::chrono::high_resolution_clock::now();
        cerr << "  **dev (iter=" << iter << " epoch=" << (tot_seen / corpus.nsentences) << ")\tllh=" << llh << " ppl: " << exp(llh / trs) << " err: " << (trs - right) / trs << " f1: " << global_f1_score << "\t[" << dev_size << " sents in " << std::chrono::duration<double, std::milli>(t_end-t_start).count() << " ms]" << endl;
        if (global_f1_score > best_f1_score) {
          best_f1_score = global_f1_score;
          TextFileSaver saver(fname);
          saver.save(model);
          // Create a soft link to the most recent model in order to make it
          // easier to refer to it in a shell script.
          if (!softlinkCreated) {
            string softlink = " latest_model";
            if (system((string("rm -f ") + softlink).c_str()) == 0 &&
                system((string("ln -s ") + fname + softlink).c_str()) == 0) {
              cerr << "Created " << softlink << " as a soft link to " << fname
                   << " for convenience." << endl;
            }
            softlinkCreated = true;
          }
        }
      }
    }
  } // should do training?
  if (true) { // do test evaluation
    double llh = 0;
    double trs = 0;
    double right = 0;
    double correct_heads = 0;
    double total_heads = 0;
    // double f1score=0.0;

    // TODO fix it to variable length label
    int confusion[4][4]= {{0}};
    int class_tp[4] = {0};
    int class_fp[4] = {0};
    int class_fn[4] = {0};

    auto t_start = std::chrono::high_resolution_clock::now();
    unsigned corpus_size = corpus.nsentencesDev;
    for (unsigned sii = 0; sii < corpus_size; ++sii) {
      const vector<unsigned>& sentence=corpus.sentencesDev[sii];
      const vector<unsigned>& sentencePos=corpus.sentencesPosDev[sii];
      const vector<string>& sentenceUnkStr=corpus.sentencesStrDev[sii];
      const vector<unsigned>& actions=corpus.correct_act_sentDev[sii];
      vector<unsigned> tsentence=sentence;
      if (!USE_SPELLING) {
        for (auto& w : tsentence)
	  if (training_vocab.count(w) == 0) w = kUNK;
      }
      ComputationGraph cg;
      double lp = 0;
      auto pred_loss = parser.log_prob_parser(&cg,sentence,tsentence,sentencePos,vector<unsigned>(),corpus.actions,corpus.intToWords,true,&right);
      vector<unsigned> & pred = pred_loss.first;
      // if (beam_size == 1)
      //   pred = parser.log_prob_parser(&cg,sentence,tsentence,sentencePos,vector<unsigned>(),corpus.actions,corpus.intToWords,true,&right);
      //else
      //  pred = parser.log_prob_parser_beam(&cg,sentence,tsentence,sentencePos,corpus.actions,beam_size,&lp);
      llh -= lp;
      trs += actions.size();
      map<int, string> rel_ref, rel_hyp;
      map<int,string> ref = parser.compute_ents(sentence.size(), actions, corpus.actions, &rel_ref);
      map<int,string> hyp = parser.compute_ents(sentence.size(), pred, corpus.actions, &rel_hyp);
      output_conll(sentence, sentencePos, sentenceUnkStr, corpus.intToWords, corpus.intToPos, ref, hyp);

      //update confusion matrix
      for (unsigned i = 0; i < sentence.size()-1; ++i) {
        auto ri = ref.find(i);
        auto hi = hyp.find(i);
        assert(ri != ref.end());
        assert(hi != hyp.end());

        if (ri->second != "0"){
            int pr = label2id[hi->second];
            int tr = label2id[ri->second];

            confusion[pr][tr]++;
        }
      }

      correct_heads += compute_correct(ref, hyp, sentence.size() - 1);
      total_heads += sentence.size() - 1;
    }

    // compute tp, fp and fn for each class
    cerr << "confusion matrix" << endl;
    for (unsigned i = 0; i < label.size(); ++i){
        for (unsigned j = 0; j < label.size(); ++j){
            cerr << i << " " << j << " " << confusion[i][j] << endl;
            if (i == j){
                class_tp[i] = confusion[i][j];
            } else {
                class_fp[i] += confusion[i][j];
                class_fn[i] += confusion[j][i];
            }
        }
    }

    // compute precision, recall, f1
    double global_f1_score = 0;
    for (unsigned i = 0; i < label.size(); ++i){
        double precision = (float)class_tp[i] / (class_tp[i] + class_fp[i]);
        double recall = (float)class_tp[i] / (class_tp[i] + class_fn[i]);
        double f1_score = 2 * (float)(precision * recall) / (precision + recall);
        double pr=precision+recall;
        if (pr==0) f1_score=0.0;
        global_f1_score += f1_score / label.size(); // Macro averaged F1
        if (std::isnan(global_f1_score)) global_f1_score=0.0;

        cerr << i << "class Precision " << precision << " Recall " << recall << " F1 " << f1_score << endl;
    }
    cerr << "F1 >> " << global_f1_score << endl;

    auto t_end = std::chrono::high_resolution_clock::now();
    cerr << "TEST llh=" << llh << " ppl: " << exp(llh / trs) << " err: " << (trs - right) / trs << " f1: " << global_f1_score << "\t[" << corpus_size << " sents in " << std::chrono::duration<double, std::milli>(t_end-t_start).count() << " ms]" << endl;
  }
  for (unsigned i = 0; i < corpus.actions.size(); ++i) {
    //cerr << corpus.actions[i] << '\t' << parser.p_r->values[i].transpose() << endl;
    //cerr << corpus.actions[i] << '\t' << parser.p_p2a->values.col(i).transpose() << endl;
  }
}


/************************************************************************************
*Title:facemorph
*Author:Yiming Xia
*Library File: Eigen  OpenCV Dlib
*Tips:
*1.You can run this in Visual Studio 2015 or newer.Library file should be configured before compiling
*  Dlib and Eigen is included in this package,OpenCV is not.
*2.You should change the address of the shape_predictor_68_face_landmarks.dat
*   which is in the facedetect function.(It is included in the package)
*3.You should also change the adress of the two images with the people whose mouth is open
*  to get the .avi successful.
*4.Release is faster than Debug.
*************************************************************************************/

#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <fstream>
#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing/render_face_detections.h>
#include <dlib/image_processing.h>
#include <dlib/gui_widgets.h>
#include <dlib/image_io.h>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <dlib/opencv.h>
#include"Eigen/Dense"
using namespace Eigen;
using namespace dlib;
using namespace cv;
using namespace std;

std::vector<cv::Point2f> facedetect(string pname)                                             //生产特征点
{
    std::vector<cv::Point2f> pointlist;
    Mat imgt = imread(pname);
    try
    {
        frontal_face_detector detector = get_frontal_face_detector();

        shape_predictor sp;//定义个shape_predictor类的实例

        deserialize("E:\\Maths\\Digital Image Processing\\HW\\HW2\\test\\Project2\\shape_predictor_68_face_landmarks.dat") >> sp;//设置人脸识别的数据包，路径为文件所在路径

        image_window win, win_faces;
        {
            array2d<rgb_pixel> img;//注意变量类型 rgb_pixel 三通道彩色图像
            load_image(img,pname);
            std::vector<dlib::rectangle> dets = detector(img);//检测人脸，获得边界框

            std::vector<full_object_detection> shapes;//注意形状变量的类型，full_object_detection
            for (unsigned long j = 0; j < dets.size(); ++j)
            {
                full_object_detection shape = sp(img, dets[j]);//预测姿势，注意输入是两个，一个是图片，另一个是从该图片检测到的边界框
                for (int i = 0; i < 68; i++)
                {
                    pointlist.push_back(cv::Point2f(shape.part(i).x(), shape.part(i).y()));
                }
                shapes.push_back(shape);
            }

            pointlist.push_back(cv::Point2f(0,0));
            pointlist.push_back(cv::Point2f(0, imgt.rows - 1));
            pointlist.push_back(cv::Point2f(imgt.cols - 1, 0));
            pointlist.push_back(cv::Point2f(imgt.cols - 1, imgt.rows - 1));
            pointlist.push_back(cv::Point2f(imgt.cols/2, 0));
            pointlist.push_back(cv::Point2f(0, imgt.rows/2));
            pointlist.push_back(cv::Point2f(imgt.cols - 1, imgt.rows/2));
            pointlist.push_back(cv::Point2f(imgt.cols/2, imgt.rows - 1));

            return pointlist;

        }
    }
    catch (exception& e)
    {
        cout << "\nexception thrown!" << endl;
        cout << e.what() << endl;
    }
}

std::vector<Vec3i> delaunayTri(Mat& img, std::vector<Point2f> points)   //将三角关系转化为索引关系
{
    Size size = img.size();
    Rect rect(0, 0, size.width, size.height);
    std::vector<Vec3i> delaunayTri;
    Subdiv2D subdiv(rect);
    for (std::vector<Point2f>::iterator it = points.begin(); it != points.end(); it++)
    {
        subdiv.insert(*it);
    }
    std::vector<Vec6f> triangleList;
    subdiv.getTriangleList(triangleList);

    std::vector<Point> pt(3);
    for (int i = 0; i < triangleList.size(); i++)
    {
        Vec6f t = triangleList[i];
        pt[0] = Point(cvRound(t[0]), cvRound(t[1]));
        pt[1] = Point(cvRound(t[2]), cvRound(t[3]));
        pt[2] = Point(cvRound(t[4]), cvRound(t[5]));
        if (rect.contains(pt[0]) && rect.contains(pt[1]) && rect.contains(pt[2]))
        {
            int ji = 0;
            cv::Vec3i ind;
            for (int j = 0; j < 3; ++j)
            {
                for (size_t k = 0; k < points.size(); k++)
                {
                    if (abs(pt[j].x - points[k].x) < 0.1   &&  abs(pt[j].y - points[k].y) < 0.1)
                    {
                        ind(j) = k;
                        ji++;
                    }
                }
            }
            if (ji > 3)//重合点检查，如果出现问号，则此图无法作为输入
            {
                cout << "?" << endl;
            }
            if (ji == 3)
            {
                delaunayTri.push_back(ind);
            }
        }
    }

    return delaunayTri;
}

MatrixXd Wraptransform(VectorXd before, VectorXd after)                                         //求变换矩阵
{
    MatrixXd M = MatrixXd::Zero(6, 6);
    MatrixXd Out = MatrixXd::Zero(3, 3);

    for (int i = 0; i < 3; i++)
    {
        M(2 * i, 0) = before(2 * i);
        M(2 * i, 1) = before(2 * i + 1);
        M(2 * i, 2) = 1;
        M(2 * i + 1, 3) = before(2 * i);
        M(2 * i + 1, 4) = before(2 * i + 1);
        M(2 * i + 1, 5) = 1;
    }
    VectorXd solve = M.colPivHouseholderQr().solve(after);

    for (int i = 0; i < 6; i++)
    {
        Out(i / 3, i % 3) = solve(i);
    }
    Out(2, 2) = 1;
    return Out;
}

Mat facemorph(float alpha, std::vector<cv::Point2f> input_pointlist, std::vector<cv::Point2f> output_pointlist, Mat img1, Mat img2, std::vector<Vec3i> trianglelist)
{

    int width = img1.cols;                   //图像的水平边
    int height = img1.rows;                  //图像的竖直边

    std::vector<cv::Point2f> morph_pointlist;
    for (int i = 0; i < 76; i++)
    {
        cv::Point2f morphpoint = (1 - alpha)*input_pointlist[i] + alpha*output_pointlist[i];
        morph_pointlist.push_back(morphpoint);
    }
    Mat record = img1.clone();
    Mat show = img1.clone();
    cv::RNG rng(time(0));
    for (int i = 0; i < trianglelist.size(); i++)
    {

        std::vector<Point> fill;
        fill.push_back(morph_pointlist[trianglelist[i](0)]);
        fill.push_back(morph_pointlist[trianglelist[i](1)]);
        fill.push_back(morph_pointlist[trianglelist[i](2)]);
        fillConvexPoly(record, fill, i);
        fillConvexPoly(show, fill, cv::Scalar(rng.uniform(0, 255), rng.uniform(0, 255), rng.uniform(0, 255)));
    }

    std::vector<MatrixXd> before_transferlist, after_transferlist;
    for (int i = 0; i < trianglelist.size(); i++)
    {
        MatrixXd before_transfer, after_transfer;
        VectorXd before = VectorXd::Zero(6);
        VectorXd after = VectorXd::Zero(6);
        VectorXd mid = VectorXd::Zero(6);
        for (int j = 0; j < 3; j++)
        {
            before(2 * j) = input_pointlist[trianglelist[i](j)].x;
            before(2 * j + 1) = input_pointlist[trianglelist[i](j)].y;
            after(2 * j) = output_pointlist[trianglelist[i](j)].x;
            after(2 * j + 1) = output_pointlist[trianglelist[i](j)].y;
            mid(2 * j) = morph_pointlist[trianglelist[i](j)].x;
            mid(2 * j + 1) = morph_pointlist[trianglelist[i](j)].y;
        }
        before_transfer = Wraptransform(mid, before);
        after_transfer = Wraptransform(mid, after);
        before_transferlist.push_back(before_transfer);
        after_transferlist.push_back(after_transfer);
    }

    Mat morphimage = Mat::zeros(height, width, CV_8UC3);
    for (int i = 0; i<height; i++)
    {
        for (int j = 0; j<width; j++)            //j是x，i是y
        {
            int num;
            num = record.at<Vec3b>(i, j)[0];

            Vector3d coordinate(j, i, 1);
            Vector3d beforecoordinate = before_transferlist[num] * coordinate;
            Vector3d aftercoordinate = after_transferlist[num] * coordinate;
            int beforex = int(beforecoordinate(0));
            int beforey = int(beforecoordinate(1));
            int afterx = int(aftercoordinate(0));
            int aftery = int(aftercoordinate(1));
            morphimage.at<Vec3b>(i, j) = (1 - alpha)*img1.at<Vec3b>(beforey, beforex) + alpha*img2.at<Vec3b>(aftery, afterx);
        }
    }

    return morphimage;
}

int main()
{
    //设置输入图像的路径
    string adress1 = "E:\\Maths\\Digital Image Processing\\Practice\\xia1.jpg";
    string adress2 = "E:\\Maths\\Digital Image Processing\\Practice\\wuyanzu1.jpg";

    Scalar points_color(0, 0, 255);

    std::vector<Vec6f> tri1;
    std::vector<cv::Point2f> input_pointlist = facedetect(adress1);                      //生成特征点
    Mat img1 = imread(adress1);
    std::vector<Vec6f> tri2;
    std::vector<cv::Point2f> output_pointlist = facedetect(adress2);                      //生成特征点
    Mat img2 = imread(adress2);
    Mat test = img1.clone();
    std::vector<Vec3i> trianglelist = delaunayTri(img1, input_pointlist);                  //生成三角索引关系，两张图共用此索引关系

    std::vector<Mat> pic;
    pic.push_back(img1);
    for (double alpha = 0.01; alpha <0.99; alpha = alpha + 0.01)
    {
        Mat morphimage = facemorph(alpha, input_pointlist, output_pointlist, img1, img2, trianglelist);
        pic.push_back(morphimage);
    }
    pic.push_back(img2);
    cout << "Finish!" << endl;
    VideoWriter output_src("facemorph.avi", CV_FOURCC('M', 'J', 'P', 'G'), 24, pic[0].size(), 1);   //输出.avi文件
    for (auto c : pic)
    {
        output_src << c;
    }
}


#ifndef CAFFE_REID_BASIC_MATCH_HPP_
#define CAFFE_REID_BASIC_MATCH_HPP_

#include <vector>
#include <map>
#include <boost/shared_ptr.hpp>
#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"

#include "caffe/tracker/bounding_box.hpp"
#include "caffe/tracker/basic.hpp"

#include "caffe/util/math_functions.hpp"

namespace caffe {

/**
 * 该文件提供了用于分配ID的所有方法
 * 包括：
 * １．定义对象的数据结构
 * ２．定义处理融合／分裂等过程的方法
 * 注意：该系列方法仅作为参考，不建议直接使用。
 * 请在阅读掌握源码基础上使用上述方法，请做适当修改后使用。
 * 这些代码仅作为辅助参考用，禁止直接使用。
 */

using namespace std;
// 关节点信息
template <typename Dtype>
struct Joint {
  Dtype x;
  Dtype y;
  Dtype v;
};

/**
 * 提议对象：当前帧检测到的所有提议对象
 */
template <typename Dtype>
struct pProp {
  // bbox位置
  BoundingBox<Dtype> bbox;
  // 关节点信息kps
  vector<Joint<Dtype> > kps;
  // 可见点数num_vis
  int num_vis;
  // 置信度评分score
  Dtype score;
  // 特征模板指针：默认不包含模板，当匹配成功或加入到目标池后，其模板会自动生成
  Dtype* fptr = NULL;

  // 是否已匹配，默认没有进行匹配
  bool is_matched = false;
  // 是否是分裂对象，默认不是分裂对象
  // 分裂对象：当多个提议对象与同一个目标对象匹配时，即存在分裂行为，此时提议对象的分裂标记有效
  bool is_split = false;
};

/**
 * 单个对象的数据结构，包含：
 * １．ID
 * ２．特征模板
 * ３．相似度：与前景对象之间的相似度
 */
template <typename Dtype>
struct pPerson {
  // ID
  int id;
  // 特征模板
  shared_ptr<Blob<Dtype> > pT;
  // 相似度
  Dtype simi;
};

/**
 * 目标对象的数据结构：可视化中的目标对象
 * 注意：该目标对象可能保存多个相互遮挡的对象
 */
template <typename Dtype>
struct pTarget {
  // 位置
  BoundingBox<Dtype> bbox;
  // kps（前景）
  vector<Joint<Dtype> > kps;
  // 可见数量（关节点可见数量）
  int num_vis;
  // 得分
  Dtype score;

  // 前景特征模板
  shared_ptr<Blob<Dtype> > pT;

  // 前景目标
  pPerson<Dtype> front;
  // 遮挡目标列表：可能包含多个遮挡对象
  vector<pPerson<Dtype> > back;

  // 回收标志: 新创建目标默认不回收
  // 回收：当目标对象没有匹配，开始准备回收
  bool is_trash = false;
  // 回收时间戳，默认为０
  int miss_frame = 0;

  //　默认值：匹配对象为０
  int matched_id = 0;
  // 匹配相似度，默认为０
  Dtype matched_similarity = 0.;
  // 合并标记，默认是不可合并
  // 合并：当多个目标对象与同一个提议对象匹配，则构成合并
  bool is_merged = false;
  // 可分裂标记，默认不可分裂
  // 分裂：当多个提议对象与同一个目标对象匹配，则构成分裂
  bool can_split = false;
  // 支持分裂数，默认为０
  // 分裂数：内部包含的遮挡对象数
  int num_split = 0;

  // 默认目标没有遮挡
  // 遮挡：当该对象与其他对象之间的最大IOU超过阈值，则认为存在遮挡
  bool is_occluded = false;
};

// APIs
/**
 * 定义了一系列方法来处理遮挡问题
 */

/**
 * 初始化目标对象
 * @param targets [目标对象集合]
 */
template <typename Dtype>
void initTargets(vector<pTarget<Dtype> >& targets);

/**
 * 匹配过程
 * @param targets   [目标对象集合]
 * @param proposals [提议对象集合]
 * @param iou_thre  [匹配IOU阈值]
 */
template <typename Dtype>
void pMatch(vector<pTarget<Dtype> >& targets, vector<pProp<Dtype> >& proposals, const Dtype iou_thre);

/**
 * 融合操作
 * @param targets [目标对象集合]
 */
template <typename Dtype>
void pMerge(vector<pTarget<Dtype> >& targets);

/**
 * 目标对象更新过程
 * @param targets   [目标对象集合]
 * @param proposals [提议对象集合]
 */
template <typename Dtype>
void updateTargets(vector<pTarget<Dtype> >& targets, vector<pProp<Dtype> >& proposals);

/**
 * 处理未匹配的目标对象
 * @param targets [目标对象集合]
 */
template <typename Dtype>
void processUnmatchedTargets(vector<pTarget<Dtype> >& targets);

/**
 * 根据特征模板计算两个对象之间的相似度
 * @param f_src [对象１]
 * @param f_dst [对象２]
 * @param thre  [有效阈值]
 * @param s     [相似度]
 */
template <typename Dtype>
void fsimilarity(const Blob<Dtype>& f_src, const Blob<Dtype>& f_dst, const Dtype thre, Dtype* s);

/**
 * 更新目标对象中的前景对象
 * @param targets [目标对象集合]
 * @param thre    [阈值]
 */
template <typename Dtype>
void foreTargetsUpdate(vector<pTarget<Dtype> >& targets, const Dtype thre);

/**
 * 更新目标对象的遮挡状态
 * @param targets   [目标对象集合]
 * @param proposals [提议对象集合]
 * @param thre      [阈值]
 */
template <typename Dtype>
void updateOccludedStatus(vector<pTarget<Dtype> >& targets, vector<pProp<Dtype> >& proposals, const Dtype thre);

/**
 * 模板更新
 * @param targets    [目标对象集合]
 * @param scale_str  [结构化颜色模板更新系数]
 * @param scale_area [区域颜色模板更新系数]
 */
template <typename Dtype>
void updateTemplates(vector<pTarget<Dtype> >& targets, const Dtype scale_str, const Dtype scale_area);

/**
 * 处理为匹配的提议对象
 * @param targets           [目标对象集合]
 * @param proposals         [提议对象集合]
 * @param iou_thre          [IOU阈值-1]
 * @param simi_thre         [相似度阈值-2]
 * @param thre_for_simi_cal [相似度计算阈值-3]
 */
template <typename Dtype>
void processUnmatchedProposals(vector<pTarget<Dtype> >& targets, vector<pProp<Dtype> >& proposals,
                               const Dtype iou_thre, const Dtype simi_thre, const Dtype thre_for_simi_cal);
/**
 * 生成新的目标对象
 * @param prop       [提议对象]
 * @param new_target [新的目标对象]
 */
template <typename Dtype>
void getNewTarget(pProp<Dtype>& prop, pTarget<Dtype>* new_target);

/**
 * 获取新的对象ID
 * @param  targets [目标对象集合]
 * @return         [新的ID]
 */
template <typename Dtype>
int getNewId(vector<pTarget<Dtype> >& targets);

/**
 * 打印目标对象信息
 * @param targets [目标对象集合]
 */
template <typename Dtype>
void printInfo(vector<pTarget<Dtype> >& targets);

}  // namespace caffe

#endif


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

#include <jni.h>
#include <stdint.h>
#include <stdio.h>
#include <array>
#include <cmath>

#include <glog/logging.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include "cpp/common/io.h"
#include "cpp/jni/jni_utils.h"
#include "cpp/jni/macros.h"
#include "cpp/photo/jpeg_encoder.h"
#include "cpp/video/stereo_mode.h"
#include "xmpmeta/gimage.h"
#include "xmpmeta/gpano.h"
#include "xmpmeta/pano_meta_data.h"
#include "xmpmeta/photo_sphere_writer.h"
#include "xmpmeta/vr_photo_writer.h"
#include "xmpmeta/xmp_writer.h"

#undef JNI_PACKAGE_NAME
#define JNI_PACKAGE_NAME com_google_vr180_media_photo

#undef JNI_CLASS_NAME
#define JNI_CLASS_NAME PhotoWriter

namespace {

static const double kRadiansToDegrees = 180.0 / M_PI;
static const int kJpegQuality = 100;

xmpmeta::PanoMetaData GetPanoMetaData(int width, int height, double fov_x,
                                      double fov_y,
                                      const Eigen::Vector3d& euler_angles) {
  xmpmeta::PanoMetaData meta;
  const double ppd_x = width / fov_x;
  const double ppd_y = height / fov_y;
  meta.full_height = static_cast<int>(180.0 * (ppd_x + ppd_y) / 2.0);
  meta.full_width = meta.full_height * 2;
  meta.cropped_width = width;
  meta.cropped_height = height;
  meta.cropped_left = (meta.full_width - meta.cropped_width) / 2;
  meta.cropped_top = (meta.full_height - meta.cropped_height) / 2;
  meta.initial_heading_degrees = 180;
  meta.pose_heading_degrees = 0.0;

  return meta;
}

// Write left image to jpeg and right image to xmp metadata.
bool WriteVRPhoto(const std::string& left, const std::string& right,
                  const xmpmeta::PanoMetaData& metadata,
                  const std::string& output_path) {
  const std::unique_ptr<xmpmeta::GPano> gpano =
      xmpmeta::GPano::CreateFromData(metadata);
  const std::unique_ptr<xmpmeta::XmpData> xmp_data =
      xmpmeta::CreateXmpData(/*create_extended=*/true);

  if (right.empty()) {
    // For mono photos, only package the GPano into the XMP metadata.
    if (!xmpmeta::WritePhotoSphereMetaToXmp(*gpano, xmp_data.get())) {
      return false;
    }
  } else {
    // Package right eye JPEG and GPano into the XMP metadata.
    const std::unique_ptr<xmpmeta::GImage> gimage =
        xmpmeta::GImage::CreateFromData(right, "image/jpeg");
    if (!xmpmeta::WriteVrPhotoMetaToXmp(*gimage, *gpano, nullptr,
                                        xmp_data.get())) {
      return false;
    }
  }

  // Write jpeg with xmp data.
  return xmpmeta::WriteLeftEyeAndXmpMeta(left, output_path, *xmp_data);
}

// Gets the euler angles about the world axes y,x,z.
Eigen::Vector3d GetEulerAngles(const Eigen::Matrix3d& coeff) {
  Eigen::Vector3d result;
  result[0] = std::atan2(coeff(0, 2), coeff(2, 2));
  const double c2 = Eigen::Vector2d(coeff(1, 1), coeff(1, 0)).norm();
  result[1] = std::atan2(-coeff(1, 2), c2);
  const double s1 = std::sin(result[0]);
  const double c1 = std::cos(result[0]);
  result[2] = std::atan2(s1 * coeff(2, 1) - c1 * coeff(0, 1),
                         c1 * coeff(0, 0) - s1 * coeff(2, 0));
  return result;
}

Eigen::Vector3d GetEulerAngles(const std::array<double, 3>& angle_axis) {
  const Eigen::Vector3d axis_vector(angle_axis.data());
  return GetEulerAngles(
      Eigen::AngleAxisd(axis_vector.norm(), axis_vector.normalized())
          .toRotationMatrix());
}

}  // namespace

extern "C" {
JNIEXPORT bool JNICALL JNI_METHOD(nativeWriteVRPhotoToFile)(
    JNIEnv* env, jobject obj, jbyteArray rgba_buffer, jint stereo_width,
    jint stereo_height, jint stride, jfloat fov_x, jfloat fov_y,
    jfloatArray angle_axis, jint stereo_mode, jstring joutput_path) {
  uint8* rgba_data =
      reinterpret_cast<uint8*>(env->GetByteArrayElements(rgba_buffer, nullptr));
  const std::string output_path = vr180::JavaStringToCppString(env, joutput_path);

  std::string left, right;
  bool success = false;
  if (fov_x == 0 || fov_y == 0) {
    // If we don't have a valid crop just save the raw image.
    success =
        vr180::EncodeRGBAJpeg(rgba_data, 0, 0, stereo_width, stereo_height,
                              stride, kJpegQuality, &left) &&
        vr180::SetFileContents(output_path, left);
  } else {
    // Get the photo metadata.
    const bool is_left_right =
        stereo_mode == static_cast<int>(vr180::StereoMode::LEFT_RIGHT);
    const int width = stereo_width / (is_left_right ? 2 : 1);
    const bool is_top_bottom =
        stereo_mode == static_cast<int>(vr180::StereoMode::TOP_BOTTOM);
    const int height = stereo_height / (is_top_bottom ? 2 : 1);
    const std::vector<float> fangle_axis =
        vr180::JavaFloatArrayToCppVector(env, angle_axis);
    const std::array<double, 3> orientation = {fangle_axis[0], fangle_axis[1],
                                               fangle_axis[2]};
    // Note: Xmpmeta does not currently support pitch and roll, so the full
    // orientation will not be saved.
    const xmpmeta::PanoMetaData metadata = GetPanoMetaData(
        width, height, fov_x, fov_y, GetEulerAngles(orientation));

    success =
        // Encode the left eye.
        vr180::EncodeRGBAJpeg(rgba_data, 0, 0, width, height, stride,
                              kJpegQuality, &left) &&
        // Encode the right eye iff the photo is a stereo photo.
        ((!is_left_right && !is_top_bottom) ||
         vr180::EncodeRGBAJpeg(rgba_data, stereo_width - width,
                               stereo_height - height, width, height, stride,
                               kJpegQuality, &right)) &&
        // Write the image(s) and metadata to file.
        WriteVRPhoto(left, right, metadata, output_path);
  }
  env->ReleaseByteArrayElements(rgba_buffer,
                                reinterpret_cast<jbyte*>(rgba_data), JNI_ABORT);
  return success;
}
}


/*    Copyright (c) 2010-2018, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 */

#include <algorithm>

#include <boost/make_shared.hpp>

#include "Tudat/SimulationSetup/EstimationSetup/createDopplerPartials.h"

namespace tudat
{

namespace observation_partials
{

using namespace ephemerides;

//! Function to generate one-way doppler partial wrt an initial position of a body.
std::shared_ptr< OneWayDopplerPartial > createOneWayDopplerPartialWrtBodyState(
        const observation_models::LinkEnds oneWayDopplerLinkEnds,
        const simulation_setup::NamedBodyMap& bodyMap,
        const std::string bodyToEstimate,
        const std::shared_ptr< PositionPartialScaling > oneWayDopplerScaler,
        const std::vector< std::shared_ptr< observation_partials::LightTimeCorrectionPartial > >&
        lightTimeCorrectionPartialObjects  )
{
    if( std::dynamic_pointer_cast< OneWayDopplerScaling >( oneWayDopplerScaler ) == nullptr )
    {
        throw std::runtime_error( "Error, expected one-way doppler scaling when making one-way doppler partial" );
    }

    // Create position partials of link ends for current body position
    std::map< observation_models::LinkEndType, std::shared_ptr< CartesianStatePartial > > positionPartials =
            createCartesianStatePartialsWrtBodyState( oneWayDopplerLinkEnds, bodyMap, bodyToEstimate );

    // Create one-doppler partials if any position partials are created (i.e. if any dependency exists).
    std::shared_ptr< OneWayDopplerPartial > oneWayDopplerPartial;
    if( positionPartials.size( ) > 0 )
    {
        oneWayDopplerPartial = std::make_shared< OneWayDopplerPartial >(
                    std::dynamic_pointer_cast< OneWayDopplerScaling >( oneWayDopplerScaler ),
                    positionPartials, std::make_pair(
                        estimatable_parameters::initial_body_state, std::make_pair( bodyToEstimate, "" ) ),
                    lightTimeCorrectionPartialObjects );
    }

    // Return doppler partial object (nullptr if no dependency exists).
    return oneWayDopplerPartial;
}

//! Function to create an object that computes the scaling of the state partials to obtain proper time rate partials
std::shared_ptr< OneWayDopplerProperTimeComponentScaling > createDopplerProperTimePartials(
        const std::shared_ptr< observation_models::DopplerProperTimeRateInterface > dopplerProperTimeInterface,
        const observation_models::LinkEnds oneWayDopplerLinkEnds,
        const observation_models::LinkEndType linkEndAtWhichPartialIsComputed )
{
    std::shared_ptr< OneWayDopplerProperTimeComponentScaling >  properTimeRateDopplerPartial = nullptr;
    if( dopplerProperTimeInterface == nullptr )
    {
        properTimeRateDopplerPartial = nullptr;
    }
    else if( std::dynamic_pointer_cast< observation_models::DirectFirstOrderDopplerProperTimeRateInterface >(
                 dopplerProperTimeInterface ) != nullptr )
    {
        bool computeStatePartials = ( oneWayDopplerLinkEnds.at( linkEndAtWhichPartialIsComputed ).first !=
                std::dynamic_pointer_cast< observation_models::DirectFirstOrderDopplerProperTimeRateInterface >(
                    dopplerProperTimeInterface )->getCentralBody( ) );
        properTimeRateDopplerPartial = std::make_shared< OneWayDopplerDirectFirstOrderProperTimeComponentScaling >(
                    std::dynamic_pointer_cast< observation_models::DirectFirstOrderDopplerProperTimeRateInterface >(
                        dopplerProperTimeInterface ), linkEndAtWhichPartialIsComputed, computeStatePartials );
    }
    else
    {
        std::cerr << "Warning, proper time contribution to Doppler observable not incorporated into Doppler partial " << std::endl;
        properTimeRateDopplerPartial = nullptr;
    }
    return properTimeRateDopplerPartial;

}

}

}


#include <cmath>
#include <iostream>
#include <ale_interface.hpp>
#include <glog/logging.h>
#include <boost/program_options.hpp>
//#include "prettyprint.hpp"
#include "../include/agent.hpp"
#include "../include/tools.hpp"

/*
  DeepQLearner(
    const ActionVect& legal_actions,
    double epsilon_start,
    const double epsilon_min,
    const double epsilon_decay,
    const int epsilon_explore_idx,
    const int replay_batch_size,
    const int replay_memory_capacity,
    const int replay_start_size,
    const std::string sampleStrategy,
    const int update_frequency,
    const double discount_factor,
    const std::string solver_param
    ):
*/

using namespace deepRL;
namespace po = boost::program_options;
double EpisodeLearning( ALEInterface& ale, deepRL::DeepQLearner& dqlearner, const bool update);
po::variables_map argmap;

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

  po::options_description desc("Valid options");
  desc.add_options()
  ("help,h", "Help message")
  ("gpu,u", po::value<bool>()->default_value(false), "Use GPU to brew Caffe")
	("gui,i", po::value<bool>()->default_value(false), "Open a GUI window")
  ("rom,o", po::value<std::string>()->default_value("./roms/breakout.bin"), "Rom file to play")
	("epsilon_start,s", po::value<double>()->default_value(1.0), "Epsilon start value for action selection")
	("epsilon_min,m", po::value<double>()->default_value(0.1), "The final epsilon")
	("epsilon_decay,d", po::value<double>()->default_value(0.98), "Epsilon decay ratio")
	("epsilon_explore_idx,x", po::value<int>()->default_value(100000), "Number of iterations needed for epsilon to reach epsilon_min")
	("replay_memory_capacity,c", po::value<int>()->default_value(50000), "Capacity of replay memory")
	("replay_start_size,r", po::value<int>()->default_value(500), "UEnough amount of transitions to start learning")
	("sampleStrategy,y", po::value<std::string>()->default_value("uniform")->required(), "Sampling strategy for transition")
	("update_frequency,f", po::value<int>()->default_value(1000), "Every #update_frequency set target_net_ = net_")
	("discount_factor,n", po::value<double>()->default_value(0.98), "Discount factor of future rewards (0,1]")
	("solver,v", po::value<std::string>()->default_value("./prototxt/aleSolver.prototxt")->required(), "Solver parameter file (*.prototxt)")
	("skip_frame,p", po::value<int>()->default_value(3), "Number of frames skipped")
	("model,l", po::value<std::string>()->default_value("")->required(), "Model file to load")
	("evaluate,u", po::value<bool>()->default_value(false), "Evaluation mode: only playing a game, no updates")
	("eval_epsilon,k", po::value<double>()->default_value(0.05), "Epsilon used in evaluate mode")
	("target_q_freq,g", po::value<int>()->default_value(1000), "Taregt_q_net_ update frequency")
  ("doubleQ,q", po::value<bool>()->default_value(false), "Enable double-Q learning");

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

  if (argmap.count("help")) {
      std::cout << desc << std::endl;
      return 0;
  }

  google::InitGoogleLogging(argv[0]);
  google::InstallFailureSignalHandler();
  google::LogToStderr();

  if (argmap["gpu"].as<bool>()) {
    caffe::Caffe::set_mode(caffe::Caffe::GPU);
  } else {
    caffe::Caffe::set_mode(caffe::Caffe::CPU);
  }

  ALEInterface ale(argmap["gui"].as<bool>());
  std::cout << "Init ale environment ok!" << std::endl;

  // Load the ROM file
  ale.loadROM(argmap["rom"].as<std::string>());

  std::cout << "Loading rom: " << argmap["rom"].as<std::string>() << " ok!" << std::endl;

  // Get the vector of legal actions
  const auto legal_actions = ale.getMinimalActionSet();

  std::cout << "Action space: " << legal_actions.size() << std::endl;

  // deepqlearner
  deepRL::DeepQLearner dqlearner(legal_actions,
				argmap["epsilon_start"].as<double>(),
				argmap["epsilon_min"].as<double>(),
				argmap["epsilon_decay"].as<double>(),
				argmap["epsilon_explore_idx"].as<int>(),
				deepRL::kMinibatchSize,
				argmap["replay_memory_capacity"].as<int>(),
				argmap["replay_start_size"].as<int>(),
				argmap["sampleStrategy"].as<std::string>(),
				argmap["update_frequency"].as<int>(),
				argmap["discount_factor"].as<double>(),
				argmap["solver"].as<std::string>(),
				argmap["evaluate"].as<bool>(),
        argmap["eval_epsilon"].as<double>(),
        argmap["target_q_freq"].as<int>(),
        argmap["doubleQ"].as<bool>());

  dqlearner.Initialize();

  if (argmap["model"].as<std::string>() !="null") {
    // Just evaluate the given trained model
    dqlearner.LoadPretrainedModel(argmap["model"].as<std::string>());
    std::cout << "Loading " << argmap["model"].as<std::string>() << std::endl;
  }
  //evaluate mode
  if (argmap["evaluate"].as<bool>()) {
    auto total_score = 0.0;
      const auto score =
          EpisodeLearning(ale, dqlearner, false);
      std::cout << "score: " << score << std::endl;
    return 0;
  }

  // learning mode
  for (auto episode = 0;; episode++) {
    dqlearner.Reset();
    const auto score = EpisodeLearning(ale, dqlearner, true);
    std::cout << "Episode: " << episode << " ,Score: " << score << std::endl;
    if (episode % 10 == 0) {
      // After every 10 episodes, evaluate the current strength
      const auto eval_score = EpisodeLearning(ale, dqlearner,false);
      std::cout << "evaluation score: " << eval_score << std::endl;
    }
  }
}

/**
 * one episode learning and return the total score
 */
double EpisodeLearning( ALEInterface& ale, deepRL::DeepQLearner& dqlearner, const bool update) {
  assert(!ale.game_over());
  std::deque<deepRL::FrameDataSp> past_frames;
  //dqlearner.replay_memory_.resetPool();

  auto total_score = 0.0;
  for (auto frame = 0; !ale.game_over(); ++frame) {
    //std::cout << "frame: " << frame << std::endl;
    const auto current_frame = deepRL::PreprocessScreen(ale.getScreen());

    past_frames.push_back(current_frame);
    if (past_frames.size() < deepRL::kInputFrameCount) {
      // If there are not past frames enough for DQN input, just select NOOP
      for (auto i = 0; i < argmap["skip_frame"].as<int>() + 1 && !ale.game_over(); ++i) {
        total_score += ale.act(PLAYER_A_NOOP);
      }
    } else {
      if (past_frames.size() > deepRL::kInputFrameCount) {
        past_frames.pop_front();
      }
      deepRL::InputFrames input_frames;
      std::copy(past_frames.begin(), past_frames.end(), input_frames.begin());
      const auto action = dqlearner.SelectAction(input_frames);
      auto immediate_score = 0.0;

      for (auto i = 0; i < argmap["skip_frame"].as<int>() + 1 && !ale.game_over(); ++i) {
        // Last action is repeated on skipped frames
        immediate_score += ale.act(action);
      }

      total_score += immediate_score;

      //clip reward for robust gradient update
      // Rewards for DQN are normalized as follows:
      // 1 for any positive score, -1 for any negative score, otherwise 0
      const auto reward =
          immediate_score == 0 ?
              0 :
              immediate_score /= std::abs(immediate_score);

      if (update) {
        // Add the current transition to replay memory
        const auto transition = ale.game_over() ?
            deepRL::Transition(input_frames, action, reward, boost::none) :
            deepRL::Transition(
                input_frames,
                action,
                reward,
                deepRL::PreprocessScreen(ale.getScreen()));
        dqlearner.replay_memory_.addTransition(transition);
	//std::cout << "Memorypool Size: " << dqlearner.replay_memory_.memory_size() << std::endl;
        // If the size of replay memory is enough, update DQN
        if (dqlearner.replay_memory_.memory_size() >= argmap["replay_start_size"].as<int>()
	           and dqlearner.numSteps()%argmap["update_frequency"].as<int>()==0 ) {
             dqlearner.MiniBatchUpdate();
        }
      }
    }
  }
  ale.reset_game();
  return total_score;
}


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

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_CONSTANT_DEFINITION_CONSTANT_HPP_INCLUDED
#define BOOST_SIMD_CONSTANT_DEFINITION_CONSTANT_HPP_INCLUDED

#include <boost/simd/config.hpp>
#include <boost/simd/as.hpp>
#include <boost/simd/detail/dispatch.hpp>
#include <boost/simd/detail/constant_traits.hpp>
#include <boost/simd/detail/dispatch/function/make_callable.hpp>
#include <boost/simd/detail/dispatch/hierarchy/functions.hpp>
#include <boost/simd/detail/dispatch/meta/scalar_of.hpp>
#include <cstdint>

namespace boost { namespace simd
{
  namespace tag
  {
    struct constant_ : boost::dispatch::constant_value_<constant_>
    {
      BOOST_DISPATCH_MAKE_CALLABLE(ext,constant_,boost::dispatch::constant_value_<constant_>);
    };
  }

  namespace ext
  {
    BOOST_DISPATCH_FUNCTION_DECLARATION(tag, constant_)
  }

  namespace detail
  {
    BOOST_DISPATCH_CALLABLE_DEFINITION(tag::constant_,constant);
  }

  template<typename Type, detail::bits_t<dispatch::scalar_of_t<Type>> Bits>
  BOOST_FORCEINLINE Type Constant()
  {
    using cst_t = typename detail::constantify<boost::dispatch::scalar_of_t<Type>,Bits>::type;
    return detail::constant( cst_t{}, boost::simd::as_<Type>{} );
  }
} }

#endif


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 *  @file DiscreteConditional.cpp
 *  @date Feb 14, 2011
 *  @author Duy-Nguyen Ta
 *  @author Frank Dellaert
 */

#include <gtsam/base/Testable.h>
#include <gtsam/base/debug.h>
#include <gtsam/discrete/DiscreteConditional.h>
#include <gtsam/discrete/Signature.h>
#include <gtsam/inference/Conditional-inst.h>

#include <algorithm>
#include <boost/make_shared.hpp>
#include <random>
#include <set>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

using namespace std;
using std::pair;
using std::stringstream;
using std::vector;
namespace gtsam {

// Instantiate base class
template class GTSAM_EXPORT
    Conditional<DecisionTreeFactor, DiscreteConditional>;

/* ************************************************************************** */
DiscreteConditional::DiscreteConditional(const size_t nrFrontals,
                                         const DecisionTreeFactor& f)
    : BaseFactor(f / (*f.sum(nrFrontals))), BaseConditional(nrFrontals) {}

/* ************************************************************************** */
DiscreteConditional::DiscreteConditional(size_t nrFrontals,
                                         const DiscreteKeys& keys,
                                         const ADT& potentials)
    : BaseFactor(keys, potentials), BaseConditional(nrFrontals) {}

/* ************************************************************************** */
DiscreteConditional::DiscreteConditional(const DecisionTreeFactor& joint,
                                         const DecisionTreeFactor& marginal)
    : BaseFactor(joint / marginal),
      BaseConditional(joint.size() - marginal.size()) {}

/* ************************************************************************** */
DiscreteConditional::DiscreteConditional(const DecisionTreeFactor& joint,
                                         const DecisionTreeFactor& marginal,
                                         const Ordering& orderedKeys)
    : DiscreteConditional(joint, marginal) {
  keys_.clear();
  keys_.insert(keys_.end(), orderedKeys.begin(), orderedKeys.end());
}

/* ************************************************************************** */
DiscreteConditional::DiscreteConditional(const Signature& signature)
    : BaseFactor(signature.discreteKeys(), signature.cpt()),
      BaseConditional(1) {}

/* ************************************************************************** */
DiscreteConditional DiscreteConditional::operator*(
    const DiscreteConditional& other) const {
  // Take union of frontal keys
  std::set<Key> newFrontals;
  for (auto&& key : this->frontals()) newFrontals.insert(key);
  for (auto&& key : other.frontals()) newFrontals.insert(key);

  // Check if frontals overlapped
  if (nrFrontals() + other.nrFrontals() > newFrontals.size())
    throw std::invalid_argument(
        "DiscreteConditional::operator* called with overlapping frontal keys.");

  // Now, add cardinalities.
  DiscreteKeys discreteKeys;
  for (auto&& key : frontals())
    discreteKeys.emplace_back(key, cardinality(key));
  for (auto&& key : other.frontals())
    discreteKeys.emplace_back(key, other.cardinality(key));

  // Sort
  std::sort(discreteKeys.begin(), discreteKeys.end());

  // Add parents to set, to make them unique
  std::set<DiscreteKey> parents;
  for (auto&& key : this->parents())
    if (!newFrontals.count(key)) parents.emplace(key, cardinality(key));
  for (auto&& key : other.parents())
    if (!newFrontals.count(key)) parents.emplace(key, other.cardinality(key));

  // Finally, add parents to keys, in order
  for (auto&& dk : parents) discreteKeys.push_back(dk);

  ADT product = ADT::apply(other, ADT::Ring::mul);
  return DiscreteConditional(newFrontals.size(), discreteKeys, product);
}

/* ************************************************************************** */
DiscreteConditional DiscreteConditional::marginal(Key key) const {
  if (nrParents() > 0)
    throw std::invalid_argument(
        "DiscreteConditional::marginal: single argument version only valid for "
        "fully specified joint distributions (i.e., no parents).");

  // Calculate the keys as the frontal keys without the given key.
  DiscreteKeys discreteKeys{{key, cardinality(key)}};

  // Calculate sum
  ADT adt(*this);
  for (auto&& k : frontals())
    if (k != key) adt = adt.sum(k, cardinality(k));

  // Return new factor
  return DiscreteConditional(1, discreteKeys, adt);
}

/* ************************************************************************** */
void DiscreteConditional::print(const string& s,
                                const KeyFormatter& formatter) const {
  cout << s << " P( ";
  for (const_iterator it = beginFrontals(); it != endFrontals(); ++it) {
    cout << formatter(*it) << " ";
  }
  if (nrParents()) {
    cout << "| ";
    for (const_iterator it = beginParents(); it != endParents(); ++it) {
      cout << formatter(*it) << " ";
    }
  }
  cout << "):\n";
  ADT::print("", formatter);
  cout << endl;
}

/* ************************************************************************** */
bool DiscreteConditional::equals(const DiscreteFactor& other,
                                 double tol) const {
  if (!dynamic_cast<const DecisionTreeFactor*>(&other)) {
    return false;
  } else {
    const DecisionTreeFactor& f(static_cast<const DecisionTreeFactor&>(other));
    return DecisionTreeFactor::equals(f, tol);
  }
}

/* ************************************************************************** */
DiscreteConditional::ADT DiscreteConditional::choose(
    const DiscreteValues& given, bool forceComplete) const {
  // Get the big decision tree with all the levels, and then go down the
  // branches based on the value of the parent variables.
  DiscreteConditional::ADT adt(*this);
  size_t value;
  for (Key j : parents()) {
    try {
      value = given.at(j);
      adt = adt.choose(j, value);  // ADT keeps getting smaller.
    } catch (std::out_of_range&) {
      if (forceComplete) {
        given.print("parentsValues: ");
        throw runtime_error(
            "DiscreteConditional::choose: parent value missing");
      }
    }
  }
  return adt;
}

/* ************************************************************************** */
DiscreteConditional::shared_ptr DiscreteConditional::choose(
    const DiscreteValues& given) const {
  ADT adt = choose(given, false);  // P(F|S=given)

  // Collect all keys not in given.
  DiscreteKeys dKeys;
  for (Key j : frontals()) {
    dKeys.emplace_back(j, this->cardinality(j));
  }
  for (size_t i = nrFrontals(); i < size(); i++) {
    Key j = keys_[i];
    if (given.count(j) == 0) {
      dKeys.emplace_back(j, this->cardinality(j));
    }
  }
  return boost::make_shared<DiscreteConditional>(nrFrontals(), dKeys, adt);
}

/* ************************************************************************** */
DecisionTreeFactor::shared_ptr DiscreteConditional::likelihood(
    const DiscreteValues& frontalValues) const {
  // Get the big decision tree with all the levels, and then go down the
  // branches based on the value of the frontal variables.
  ADT adt(*this);
  size_t value;
  for (Key j : frontals()) {
    try {
      value = frontalValues.at(j);
      adt = adt.choose(j, value);  // ADT keeps getting smaller.
    } catch (exception&) {
      frontalValues.print("frontalValues: ");
      throw runtime_error("DiscreteConditional::choose: frontal value missing");
    }
  }

  // Convert ADT to factor.
  DiscreteKeys discreteKeys;
  for (Key j : parents()) {
    discreteKeys.emplace_back(j, this->cardinality(j));
  }
  return boost::make_shared<DecisionTreeFactor>(discreteKeys, adt);
}

/* ****************************************************************************/
DecisionTreeFactor::shared_ptr DiscreteConditional::likelihood(
    size_t parent_value) const {
  if (nrFrontals() != 1)
    throw std::invalid_argument(
        "Single value likelihood can only be invoked on single-variable "
        "conditional");
  DiscreteValues values;
  values.emplace(keys_[0], parent_value);
  return likelihood(values);
}

/* ************************************************************************** */
#ifdef GTSAM_ALLOW_DEPRECATED_SINCE_V42
void DiscreteConditional::solveInPlace(DiscreteValues* values) const {
  ADT pFS = choose(*values, true);  // P(F|S=parentsValues)

  // Initialize
  DiscreteValues mpe;
  double maxP = 0;

  // Get all Possible Configurations
  const auto allPosbValues = frontalAssignments();

  // Find the maximum
  for (const auto& frontalVals : allPosbValues) {
    double pValueS = pFS(frontalVals);  // P(F=value|S=parentsValues)
    // Update maximum solution if better
    if (pValueS > maxP) {
      maxP = pValueS;
      mpe = frontalVals;
    }
  }

  // set values (inPlace) to maximum
  for (Key j : frontals()) {
    (*values)[j] = mpe[j];
  }
}

/* ************************************************************************** */
size_t DiscreteConditional::solve(const DiscreteValues& parentsValues) const {
  ADT pFS = choose(parentsValues, true);  // P(F|S=parentsValues)

  // Then, find the max over all remaining
  size_t max = 0;
  double maxP = 0;
  DiscreteValues frontals;
  assert(nrFrontals() == 1);
  Key j = (firstFrontalKey());
  for (size_t value = 0; value < cardinality(j); value++) {
    frontals[j] = value;
    double pValueS = pFS(frontals); // P(F=value|S=parentsValues)
    // Update solution if better
    if (pValueS > maxP) {
      maxP = pValueS;
      max = value;
    }
  }
  return max;
}
#endif

/* ************************************************************************** */
size_t DiscreteConditional::argmax() const {
  size_t maxValue = 0;
  double maxP = 0;
  assert(nrFrontals() == 1);
  assert(nrParents() == 0);
  DiscreteValues frontals;
  Key j = firstFrontalKey();
  for (size_t value = 0; value < cardinality(j); value++) {
    frontals[j] = value;
    double pValueS = (*this)(frontals);
    // Update MPE solution if better
    if (pValueS > maxP) {
      maxP = pValueS;
      maxValue = value;
    }
  }
  return maxValue;
}

/* ************************************************************************** */
void DiscreteConditional::sampleInPlace(DiscreteValues* values) const {
  assert(nrFrontals() == 1);
  Key j = (firstFrontalKey());
  size_t sampled = sample(*values);  // Sample variable given parents
  (*values)[j] = sampled;            // store result in partial solution
}

/* ************************************************************************** */
size_t DiscreteConditional::sample(const DiscreteValues& parentsValues) const {
  static mt19937 rng(2);  // random number generator

  // Get the correct conditional density
  ADT pFS = choose(parentsValues, true);  // P(F|S=parentsValues)

  // TODO(Duy): only works for one key now, seems horribly slow this way
  if (nrFrontals() != 1) {
    throw std::invalid_argument(
        "DiscreteConditional::sample can only be called on single variable "
        "conditionals");
  }
  Key key = firstFrontalKey();
  size_t nj = cardinality(key);
  vector<double> p(nj);
  DiscreteValues frontals;
  for (size_t value = 0; value < nj; value++) {
    frontals[key] = value;
    p[value] = pFS(frontals);  // P(F=value|S=parentsValues)
    if (p[value] == 1.0) {
      return value;  // shortcut exit
    }
  }
  std::discrete_distribution<size_t> distribution(p.begin(), p.end());
  return distribution(rng);
}

/* ************************************************************************** */
size_t DiscreteConditional::sample(size_t parent_value) const {
  if (nrParents() != 1)
    throw std::invalid_argument(
        "Single value sample() can only be invoked on single-parent "
        "conditional");
  DiscreteValues values;
  values.emplace(keys_.back(), parent_value);
  return sample(values);
}

/* ************************************************************************** */
size_t DiscreteConditional::sample() const {
  if (nrParents() != 0)
    throw std::invalid_argument(
        "sample() can only be invoked on no-parent prior");
  DiscreteValues values;
  return sample(values);
}

/* ************************************************************************* */
vector<DiscreteValues> DiscreteConditional::frontalAssignments() const {
  vector<pair<Key, size_t>> pairs;
  for (Key key : frontals()) pairs.emplace_back(key, cardinalities_.at(key));
  vector<pair<Key, size_t>> rpairs(pairs.rbegin(), pairs.rend());
  return DiscreteValues::CartesianProduct(rpairs);
}

/* ************************************************************************* */
vector<DiscreteValues> DiscreteConditional::allAssignments() const {
  vector<pair<Key, size_t>> pairs;
  for (Key key : parents()) pairs.emplace_back(key, cardinalities_.at(key));
  for (Key key : frontals()) pairs.emplace_back(key, cardinalities_.at(key));
  vector<pair<Key, size_t>> rpairs(pairs.rbegin(), pairs.rend());
  return DiscreteValues::CartesianProduct(rpairs);
}

/* ************************************************************************* */
// Print out signature.
static void streamSignature(const DiscreteConditional& conditional,
                            const KeyFormatter& keyFormatter,
                            stringstream* ss) {
  *ss << "P(";
  bool first = true;
  for (Key key : conditional.frontals()) {
    if (!first) *ss << ",";
    *ss << keyFormatter(key);
    first = false;
  }
  if (conditional.nrParents() > 0) {
    *ss << "|";
    bool first = true;
    for (Key parent : conditional.parents()) {
      if (!first) *ss << ",";
      *ss << keyFormatter(parent);
      first = false;
    }
  }
  *ss << "):";
}

/* ************************************************************************* */
std::string DiscreteConditional::markdown(const KeyFormatter& keyFormatter,
                                          const Names& names) const {
  stringstream ss;
  ss << " *";
  streamSignature(*this, keyFormatter, &ss);
  ss << "*\n" << std::endl;
  if (nrParents() == 0) {
    // We have no parents, call factor method.
    ss << DecisionTreeFactor::markdown(keyFormatter, names);
    return ss.str();
  }

  // Print out header.
  ss << "|";
  for (Key parent : parents()) {
    ss << "*" << keyFormatter(parent) << "*|";
  }

  auto frontalAssignments = this->frontalAssignments();
  for (const auto& a : frontalAssignments) {
    for (auto&& it = beginFrontals(); it != endFrontals(); ++it) {
      size_t index = a.at(*it);
      ss << DiscreteValues::Translate(names, *it, index);
    }
    ss << "|";
  }
  ss << "\n";

  // Print out separator with alignment hints.
  ss << "|";
  size_t n = frontalAssignments.size();
  for (size_t j = 0; j < nrParents() + n; j++) ss << ":-:|";
  ss << "\n";

  // Print out all rows.
  size_t count = 0;
  for (const auto& a : allAssignments()) {
    if (count == 0) {
      ss << "|";
      for (auto&& it = beginParents(); it != endParents(); ++it) {
        size_t index = a.at(*it);
        ss << DiscreteValues::Translate(names, *it, index) << "|";
      }
    }
    ss << operator()(a) << "|";
    count = (count + 1) % n;
    if (count == 0) ss << "\n";
  }
  return ss.str();
}

/* ************************************************************************ */
string DiscreteConditional::html(const KeyFormatter& keyFormatter,
                                 const Names& names) const {
  stringstream ss;
  ss << "
\n
  ";
  streamSignature(*this, keyFormatter, &ss);
  ss << "
\n";
  if (nrParents() == 0) {
    // We have no parents, call factor method.
    ss << DecisionTreeFactor::html(keyFormatter, names);
    return ss.str();
  }

  // Print out preamble.
  ss << "
\n
\n";

  // Print out header row.
  ss << "
";
  for (Key parent : parents()) {
    ss << "" << keyFormatter(parent) << "	";
  }
  auto frontalAssignments = this->frontalAssignments();
  for (const auto& a : frontalAssignments) {
    ss << "";
    for (auto&& it = beginFrontals(); it != endFrontals(); ++it) {
      size_t index = a.at(*it);
      ss << DiscreteValues::Translate(names, *it, index);
    }
    ss << "	";
  }
  ss << "
\n";

  // Finish header and start body.
  ss << "
\n
\n";

  // Output all rows, one per assignment:
  size_t count = 0, n = frontalAssignments.size();
  for (const auto& a : allAssignments()) {
    if (count == 0) {
      ss << "
";
      for (auto&& it = beginParents(); it != endParents(); ++it) {
        size_t index = a.at(*it);
        ss << "" << DiscreteValues::Translate(names, *it, index) << "	";
      }
    }
    ss << "" << operator()(a) << "	";  // value
    count = (count + 1) % n;
    if (count == 0) ss << "
\n";
  }

  // Finish up
  ss << "
\n
\n
";
  return ss.str();
}

/* ************************************************************************* */

}  // namespace gtsam