/* itgdec.cpp
 *
 * Decoder base class implementation.
 *
 * $Id: itgdec.cpp,v 1.12 2008/02/03 00:55:35 post Exp $
 */

#include "itgdec.h"
#include "globals.h"
#include <math.h>

using namespace std;

/* maximum over lhs preterminals of translation prob */
double lhs_max_prob(lhs_to_prob_table &yitr) {
  double result = -INFINITY;
  lhs_to_prob_table::iterator l_iter;
  for (l_iter = yitr.begin(); l_iter != yitr.end(); l_iter++) {
    if (l_iter->second > result)
      result = l_iter->second;
  }
  return result;
}

Decoder::Decoder(ITGGrammar* ig, ITGCorp* ic) : ITG(ig, ic) {
  Globals* globals = Globals::Instance();

  // max number of insertion candidates
  MAXINS = globals->get<int>("maxins");
  // max number of translation candidates per source word
  MAXTRS = globals->get<int>("maxtrs");
  
  // MAXV=mt*MAXTRS;//E(0, T)
  MAXT = Globals::Get<int>("maxsentlen");
  MAXV = DECODEMAXV;
//   cout << "** MAXT = " << MAXT << "   MAXV = " << MAXV << endl;

  en = new WORDTYPE[MAXV];
  en[0] = 0;//special use for initialization in dinit
  cn = new WORDTYPE[MAXT];
}

/** Decoder::findkbest(phrase, K, & kbests)
 *
 *  Finds the k-best translation for the length _lspan_ phrase
 *  starting at position _t_.  Upon completion, _kbests_ will contain
 *  a list of up to k translations, sorted in decreasing order of
 *  probability.  The prob of a translation is the maximum over
 *  possible lhs preterminals for the phrase-to-phrase grammar rule,
 *  if more than one rule generates the same phrase pair.
 *
 *  If a single word phrase is not found, we treat it as <unk>.
 *  Multi-word phrases that aren't found in the t-table are not put
 *  into the chart.
 */
int Decoder::findkbest(WORDTYPE *phrase, unsigned K, list<WORDTYPE*> &kbest) { 

  phrase_table::iterator l2itr;
  phrase_entry::iterator l1itr;
  list<WORDTYPE*>::iterator kitr;

  int phrlen = phraselen(phrase);

  // search for phrase in table
  kbest.clear();
  l2itr = grammar->probs_table.find(phrase);
  if (l2itr == grammar->probs_table.end()) {  // didn't find phrase
    if (phrlen == 1) // pretend we're the unknown word
      l2itr = grammar->probs_table.find(grammar->unkphr);
    else // phrases or insertions
      return 0;

    // maybe unk wasn't in the table
    if (l2itr == grammar->probs_table.end())
      return 0;
  }

  // look at all the l1 phrases that pair with this l2 phrase
  for (l1itr = l2itr->second.begin(); l1itr != l2itr->second.end(); l1itr++){
    // insert into sorted list if it fits
    for (kitr = kbest.begin(); kitr != kbest.end(); kitr++){
      // for both old and new phrase, take max prob over possible preterminals
      if (lhs_max_prob(grammar->probs_table[phrase][*kitr]) < lhs_max_prob(l1itr->second))
        break;
    }
    if ( (kbest.size() < K) || (kitr != kbest.end()) )
      kbest.insert(kitr, l1itr->first);

    // the insertion might have made the size of the kbest list greater than K
    if ( kbest.size() > K ) {
      kbest.pop_back();
    }
  }

  /*
  cerr << "* kbest(" << phrase << ") - " << "[" << kbest.size() << "]\n";
  for (list<WORDTYPE*>::iterator kitr = kbest.begin(); kitr != kbest.end(); kitr++) {
    cerr << "  " << transliterate(*kitr,&l1vocab) << " (" << grammar->phr2phr_prob(*kitr,phrase) << ")\n";
  }
  */

  return kbest.size();
}

void dump_kbest(WORDTYPE* phrase, list <WORDTYPE*>& kbest)
{
  cerr << "* kbest(" << phrase << ") - " << "[" << kbest.size() << "]\n";
  for (list<WORDTYPE*>::iterator kitr = kbest.begin(); kitr != kbest.end(); kitr++) {
    cerr << "  " << transliterate(*kitr,&l1vocab) << " (" << grammar->phr2phr_prob(*kitr,phrase) << ")\n";
  }
}


/*	LM helper functions.  A window of ngram state slides across
 *	the english side of the rule from left to right.  Each word
 *	from the rule is shifted onto the right edge of this window.
 *	Words in the left edge lm state have their ngram prob added
 *	into cost of new item.  Words in right edge replace existing
 *	words in the window, but do not have their cost counted.  If
 *	lm state is narrower than window size, the constit is shorter
 *	than the ngram size, and the left and right lm state are the
 *	same, so the replacement doesn't change anything.  If the lm
 *	state is larger than the window size, it will replace
 *	everything in the window.
 *
 *	Zero means no word - calls to ngram() will return 0 if any
 *	of the word id params are zero.
 */

double Decoder::push_lm_window(int *win, int word, int lm_size, 
                                  int *uchoice) {
  if (word == 0)
    return 0.0;
  bool set_uchoice = (win[1] == 0);	/* haven't seen first n-1 words */
  for (int i = 1; i < lm_size; i++)	/* slide array back one position */ 
    win[i-1] = win[i];		
  win[lm_size-1] = word;		/* push onto end of array */
  if (set_uchoice) {
    int j = 0;				/* freeze left lm state for new item */
    while (!win[j])
      j++;
    for (int i = 0; i < lm_size-1 && j+i < lm_size; i++) {
      uchoice[i] = win[j+i];
    }
  }
  return (lm_size == 2) ?
    ngram->bigram(win[lm_size-2], win[lm_size-1]) : 
    ngram->trigram(win[lm_size-3],
                  win[lm_size-2],
                  win[lm_size-1]) ;
}

double Decoder::compute_lm(int *ul, int *ur, int *vl, int *vr,
                              const Rule *rule,
                              int *uchoice, int *vchoice) {
  int lm_size = ngram->N;
  int lm_window[3];		/* sliding window of ngram context */
  double logprob = 0;
  int rule_pos = 0;
  for (int i = 0; i < lm_size; i++) 	/* initialize */
    lm_window[i] = 0;

  if (rule->e_rhs)	// terminals at beginning of rule
    for (rule_pos = 0; rule->e_rhs[rule_pos] > 0; rule_pos++)
      logprob += push_lm_window(lm_window, rule->e_rhs[rule_pos], lm_size, uchoice);

  // first child nonterminal
  for (int i = 0; i < lm_size-1; i++)		// left edge
    logprob += push_lm_window(lm_window, ul[i], lm_size, uchoice);
  for (int i = 0; i < lm_size-1 && vl[i]; i++)	// right edge
    lm_window[lm_size-1-i] = vl[i];

  if (rule->e_rhs)	// middle terminals in rule
    for (rule_pos++; rule->e_rhs[rule_pos] > 0; rule_pos++)
      logprob += push_lm_window(lm_window, rule->e_rhs[rule_pos], lm_size, uchoice);

  if (rule->binary) {
    // second child nonterminal
    for (int i = 0; i < lm_size-1 && ur[i]; i++)	// left edge
      logprob += push_lm_window(lm_window, ur[i], lm_size, uchoice);
    for (int i = 0; i < lm_size-1 && vr[i]; i++)	// right edge
      lm_window[lm_size-1-i] = vr[i];
    
    if (rule->e_rhs)	// final terminals in rule
      for (rule_pos++; rule->e_rhs[rule_pos] > 0; rule_pos++)
        logprob += push_lm_window(lm_window, rule->e_rhs[rule_pos], lm_size, uchoice);
  }

  // copy righthand lm state of new item
  for (int i = 0; i < lm_size-1; i++) {
    WORDTYPE w = lm_window[lm_size-1-i];
    vchoice[i] = w;
  }
  
  if (Globals::Get<int>("debuglevel") >= 4) {
    fprintf(logfile, "compute lm");
    fprintf(logfile, "[%d %d .. %d %d] ", ul[0], ul[1], vl[1], vl[0] );
    if (rule->binary)
      fprintf(logfile, "[%d %d .. %d %d] ", ur[0], ur[1], vr[1], vr[0] );
    fprintf(logfile, "=> [%d %d .. %d %d] %g\n", uchoice[0], uchoice[1],
            vchoice[1], vchoice[0], logprob);
  }

  return logprob;
}
