/* ngram.cpp
 *
 * $Id: ngram.cpp,v 1.2 2011/10/16 22:36:43 gildea Exp $
 */

#include <cmath>

#include "ngram.h"
#include "array.h"

using namespace std;

const char* SEP = " \n\r\t";
int MAXLINELEN = 1000;

/*  Returns the n-gram probability of a phrase, using
 *  whatever n-gram model is defined.
 */
double NGram::prob(WORDTYPE* phr) {
  double pr = 0.0;
  int phrlen = phraselen(phr);

  if (phrlen >= N) {
    if (N == 2) {
      for (int i = 0; i < phrlen - 1; i++) {
        pr += bigram(phr[i], phr[i+1]);
      }
    } else if (N == 3) {
      for (int i = 0; i < phrlen - 2; i++) {
        pr += trigram(phr[i], phr[i+1], phr[i+2]);
      }
    }
  }

//   cout << "* ngram.prob(" << phr << ") = " << pr << endl;

  return pr;
}

double NGram::unigram(WORDTYPE w) {
  return uninonz(w) ? unigram_tbl[w] : -99;
}

double NGram::bigram(WORDTYPE m, WORDTYPE n) {
  double prb = 0.0;

  if (m == 0 || n == 0)
    prb = 0.0;
  else if (binnonz(m,n))
    prb = bigram_tbl[m][n];
  else 
    prb = unigram_w_tbl[m] + unigram(n);

  return prb;
}

double NGram::trigram(WORDTYPE l, WORDTYPE m, WORDTYPE n) {
  // this check is very important!
  if (l == 0 || m == 0 || n == 0)
    return 0;

  if (trinonz(l, m, n))
    return trigram_tbl[l][m][n];

  if (biwnonz(l, m))
    return bigram_w_tbl[l][m] + bigram(m, n);

  return bigram(m, n);
}

int NGram::uninonz(int n){
  unigram_table::iterator xitr = unigram_tbl.find(n);
  if (xitr != unigram_tbl.end())
    return 1;

  return 0;
}

int NGram::binnonz(int m, int n){
  bigram_table::iterator xitr;
  unigram_table::iterator yitr;

  xitr = bigram_tbl.find(m);
  if (xitr != bigram_tbl.end()){
    yitr = xitr->second.find(n);
    if (yitr != xitr->second.end())
      return 1;
  }

  return 0;
}


int NGram::trinonz(int l, int m, int n){
  trigram_table::iterator xitr;
  bigram_table::iterator yitr;
  unigram_table::iterator zitr;

  xitr = trigram_tbl.find(l);
  if (xitr != trigram_tbl.end()) {
    yitr = xitr->second.find(m);
    if (yitr != xitr->second.end()) {
      zitr = yitr->second.find(n);
      if (zitr != yitr->second.end())
        return 1;
    }
  }

  return 0;
}

int NGram::biwnonz(int m, int n){
  bigram_table::iterator xitr;
  unigram_table::iterator yitr;

  xitr = bigram_w_tbl.find(m);
  if (xitr != bigram_w_tbl.end()){
    yitr = xitr->second.find(n);
    if (yitr != xitr->second.end())
      return 1;
  }

  return 0;
}


void NGram::load(string _lmfile) {
  if (loaded && lmfile == _lmfile)
    return;

  lmfile = _lmfile;

  char line[MAXLINELEN];
  double logpr, w_logpr;
  char *token;

  int counts[3];  
  counts[0] = counts[1] = counts[2] = 0;

  FILE *lf = fopen(lmfile.data(), "r");
  if (NULL == lf) {
    cerr << "* FATAL: couldn't open ngram file '" << lmfile << "'\n";
    exit(1);
  }

  stringstream s;
  if (Globals::Get<int>("verbosity") >= 1) {
    s << "Loading ngrams from '" << lmfile << "'";
    print_header(s.str());
  }

  /* read the unigrams */
  // skip past the header
  do {
    fgets(line, MAXLINELEN, lf);
    token = strtok(line, SEP);
  } while (!token || strcmp(token, "\\1-grams:") != 0);

  //read the 1 grams
  int uni;
  while (fscanf(lf, "%lf\t%d%lf\n", &logpr, &uni, &w_logpr)==3){
    //  unigram_tbl[uni]=exp(unilog*log(10)); //actual probability
    unigram_tbl[uni] = logpr * log(10.0); //log probability, natural log
    //  unigram_w_tbl[uni]=exp(uniwlog*log(10)); //actual probability
    unigram_w_tbl[uni] = w_logpr * log(10.0); //log weight
    counts[0]++;
  }

  if (Globals::Get<int>("verbosity") >= 1) {
    printf("* read %d unigrams...", counts[0]);
    fflush(stdout);
  }

  /* read the bigrams */
  // skip the header
  do {
    fgets(line, MAXLINELEN, lf);
    token = strtok(line, SEP);
  } while (!token || strcmp(token, "\\2-grams:") != 0);

  // read the 2 grams
  int nsucc, bin1, bin2;
  do {
    if (fgets(line, MAXLINELEN, lf) == NULL)
      break;
    nsucc = sscanf(line, "%lf\t%d %d%lf\n", &logpr, &bin1, &bin2, &w_logpr);
    if (nsucc >= 3)
      bigram_tbl[bin1][bin2] = logpr * log(10.0); //log prob
    if (nsucc == 4)
      bigram_w_tbl[bin1][bin2] = w_logpr * log(10.0);

    counts[1]++;

  } while (nsucc>=3);

  if (Globals::Get<int>("verbosity") >= 1) {
    printf("read %d bigrams...", counts[1]);
    fflush(stdout);
  }

  // read the 3-grams
  if (N >= 3) {
    do {
      fgets(line, MAXLINELEN, lf);
      token = strtok(line, SEP);
    } while (!token || strcmp(token, "\\3-grams:")!=0);

    int tri1, tri2, tri3;
    while (fscanf(lf, "%lf\t%d%d%d\n", &logpr, &tri1, &tri2, &tri3)==4){
      //  unigram_tbl[uni]=exp(unilog*log(10)); //actual probability
      trigram_tbl[tri1][tri2][tri3] = logpr * log(10.0); //log probability, natural log
      counts[2]++;
    }

    if (Globals::Get<int>("verbosity") >= 1) {
      printf("read %d trigrams...", counts[2]);
      fflush(stdout);
    }
  }

  if (Globals::Get<int>("verbosity") >= 1) {
    printf("done.\n");
    fflush(stdout);
  }

  fclose(lf);

  loaded = true;
}
