/* itgtrain.cpp
 *
 * Contains class implementaion used for training ITG model.
 *
 * $Id: itgtrain.cpp,v 1.10 2008/05/19 01:22:35 zhanghao Exp $
 */

#include "itgtrain.h"
#include <cmath>
#include <sys/time.h>

using namespace std;

Trainer::Trainer(ITGGrammar* ig, ITGCorp* ic) : ITG(ig, ic)
{
  chart = new ITGChart(MAXT,MAXV);

  counts_file = runid + ".cnt";    
  rule_probs_file  = runid + ".rules.prb";
  t_probs_file = runid + ".t.prb";

  en = new WORDTYPE[MAXT];
  cn = new WORDTYPE[MAXV];


  CLONING = Globals::Get<bool>("cloning");


  MAXITER = Globals::Get<int>("iterations");

  string rules_file = Globals::Get<string>("grammarfile");
  string ttable_file = Globals::Get<string>("probfile");
}


Trainer::~Trainer()
{
  delete chart;

  delete[] en;
  delete[] cn;

}


/** Trainer::tinit()
 *
 *  Initializes the chart for phrase-based training.
 */
void Trainer::tinit() {
  pi = 0.0;

  // set cells to zeros
  chart->clean(T, V);

  // base cases for phrases
  for (int deltas = 0; deltas <= T; deltas++)  {
    for (int s = 0; s <= T - deltas; s++)  {
      int t = s + deltas;
      WORDTYPE *en_part = splice(en,s,t);
      for (int deltau = 0; deltau <= V; deltau++)  {
        if (deltas == 0 && deltau == 0)
          continue;
        for (int u = 0; u <= V - deltau; u++)  {
          int v = u + deltau;
          WORDTYPE *cn_part = splice(cn,u,v);

          double prob = grammar->phr2phr_prob(en_part,cn_part);

          chart->incr_inside(s, t, u, v, grammar->T, prob);

	  if (Globals::Get<int>("debuglevel") >= 3) {
	    if (prob > 0.0)  {
	      fprintf(logfile, 
		      "prob[%d,%d,%d,%d] = %.2g\n",s,t,u,v,chart->inside(s,t,u,v,grammar->T));
	      fprintf(logfile, "  en = ");
	      fprint_array(logfile,en_part);
	      fprintf(logfile, "\n");
	      fprintf(logfile, "  cn = ");
	      fprint_array(logfile, cn_part);
	      fprintf(logfile, "\n");
	      fflush(logfile);
	    }
	  }
          
          delete[] cn_part;
        }
      }
      delete[] en_part;
    }
  }

  //  base case of the outer probabilities
  chart->incr_outside(0, T, 0, V, grammar->start_symbol, 1.0);

  return; 
}



void Trainer::parse() {

  /* parse bottom-up, with the guarantee that inside probs for smaller
   *  spans in L1 will be built before larger spans in L1
   */

  for (int deltas = 0; deltas <= T; deltas++)  {

    for (int s = 0; s <= T - deltas; s++)  {
      int t = s + deltas;

      for (int deltau = 0; deltau <= V; deltau++)  {
        if (deltas == 0 && deltau == 0)
          continue;

        for (int u = 0; u <= V - deltau; u++)  {
          int v = u + deltau;

          for (int ruleno = 0; ruleno < grammar->B; ruleno++) {
            Rule *rule = grammar->binary_rules[ruleno];
            int lhs =  rule->lhs_int;
            int rhsl = rule->rhsl_int;
            int rhsr = rule->rhsr_int;

            for (int S = s; S <= t; S++)  {

              for (int U = u; U <= v; U++)  {
#ifdef CLONECODE
                bool do_clone;
#endif

                double betal, betar;

                if (rule->inverted) {
                  betal = chart->inside(s, S, U, v, rhsl);
                  if (betal <= 0.0)  continue;  // this saves a lot of time
                  betar = chart->inside(S, t, u, U, rhsr);
#ifdef CLONECODE
                  do_clone = (CLONING && s != S && S != t && u == U);
#endif
                } else {
                  betal = chart->inside(s, S, u, U, rhsl);
                  if (betal <= 0.0)  continue;
                  betar = chart->inside(S, t, U, v, rhsr);
#ifdef CLONECODE
                  do_clone = (CLONING && s != S && S != t && v == U);
#endif
                }

                /* normal case */

                double mass = rule->prob * betal * betar;
#ifdef CLONECODE
                if (do_clone)
                  mass *= (1 - grammar->pins);
#endif
                // regular inside probability mass
#ifdef CLONECODE
                chart->incr_inside(s, t, u, v, lhs, mass, do_clone, false);
#else
		chart->incr_inside(s, t, u, v, lhs, mass, false, false);
#endif
		
		if (Globals::Get<int>("debuglevel") >= 4) {
		  fprintf(logfile, "  * inside prob of [%s;%d,%d,%d,%d] += %g = %g %g %g\n", 
			  grammar->nonterminals[lhs].data(),
			  s, t, u, v,  mass, rule->prob, betal, betar);
		}

#ifdef CLONECODE
                /* cloning */

                if (do_clone) {
                  double beta_c = chart->clone(S,t);
                  if (beta_c > 0.0) {
                    double mass = betal * beta_c * grammar->pins;
                    // increment the clone mass
                    chart->incr_inside(s, t, u, v, lhs, mass, true, true);
                  }
                }
#endif
              } // end for splitting U
            } // end for splitting S

#ifdef CLONECODE
            /* update the clone for the current (s,t) span in L1.
             * Since our outer loop grows in (s,t), the mass for the
             * current (s,t) will have been built entirely when we
             * call chart->clone(S,t) on it, when considering larger
             * intervals
             */

            if (CLONING && t > s && v > u)  {
              double mass = chart->inside(s,t,u,v,lhs);
              chart->incr_clone(s,t,mass);
            }
#endif
          }//end for all binary rules

          /* unary rules */
	  //warning: assuming the unary rules are not chained up in this order
          for (int r = 0; r < grammar->U; r++){
            Rule *rule = grammar->unary_rules[r];
            int lhs = rule->lhs_int;
            int rhs = rule->rhsl_int;
            double mass = rule->prob * chart->inside(s, t, u, v, rhs);
            chart->incr_inside(s, t, u, v, lhs, mass);
	    if (Globals::Get<int>("debuglevel") >= 4) {
	      fprintf(logfile, "  * inside prob of [%s;%d,%d,%d,%d] += %g = %g %g\n", 
		      grammar->nonterminals[lhs].data(),
		      s, t, u, v,  mass, rule->prob, chart->inside(s, t, u, v, rhs));
	    }
          }

	  if (Globals::Get<int>("debuglevel") >= 3) {
	    for (int i=0; i<grammar->_nonterminal_cnt; i++){
	      fprintf(logfile, "  * inside prob of [%s;%d,%d,%d,%d] = %g\n", 
		     grammar->nonterminals[i].data(),
		     s, t, u, v,  chart->inside(s, t, u, v, i));
	    }
	  }
        }//end for u
      }//end for deltau
    }//end for s
  }//end for deltas

  /* top-down proceeds in much the same manner */

  for (int deltas = T; deltas > 0; deltas--) {

    for (int s = 0; s <= T - deltas; s++) {
      int t = s + deltas;

      for (int deltau = V; deltau > 0; deltau--) {
        if (deltas == 0 && deltau == 0)
          continue;

        for (int u = 0; u <= V - deltau; u++) {
          int v = u + deltau;

          // unary rules
          for (int ruleno = 0; ruleno < grammar->U; ruleno++) {
            Rule *rule = grammar->unary_rules[ruleno];
            int lhs = rule->lhs_int;
            int rhs = rule->rhsl_int;
	    double alpha = chart->outside(s, t, u, v, lhs);
	    if (alpha <= 0.0) continue;
            double mass = rule->prob * alpha;
            chart->incr_outside(s, t, u, v, rhs, mass);
          }

          //binary rules
          for (int ruleno = 0; ruleno < grammar->B; ruleno++) {
            Rule *rule = grammar->binary_rules[ruleno];
            int lhs =  rule->lhs_int;
            int rhsl = rule->rhsl_int;
            int rhsr = rule->rhsr_int;
            int inv = rule->inverted;
            double pr = rule->prob;

            double alpha = chart->outside(s, t, u, v, lhs);

            if (alpha <= 0.0)  continue;

            for (int S = s; S <= t; S++) {

              for (int U = u; U <= v; U++) {
                if (inv) {
                  double betal = chart->inside(s, S, U, v, rhsl);
                  if (betal <= 0.0)  continue;
                  chart->incr_outside(S, t, u, U, rhsr, pr * betal * alpha);
                  double betar = chart->inside(S, t, u, U, rhsr);
                  if (betar <= 0.0)  continue;
                  chart->incr_outside(s, S, U, v, rhsl, pr * betar * alpha);
                } else {
                  //beta(sSuU, rhsl)
                  double betal = chart->inside(s, S, u, U, rhsl);
                  if (betal <= 0.0)  continue;
                  chart->incr_outside(S, t, U, v, rhsr, pr * betal * alpha);
                  double betar = chart->inside(S, t, U, v, rhsr);
                  if (betar <= 0.0)  continue;
                  chart->incr_outside(s, S, u, U, rhsl, pr * betar * alpha);
                }
              }//end for splitting U
            }//end for splitting S
          }//end of binary rules

	  if (Globals::Get<int>("debuglevel") >= 3) {
	    for (int i = 0; i < grammar->_nonterminal_cnt; i++) {
	      //             if (chart->outside(s,t,u,v,i) > defaultprob)
              fprintf(logfile, "  * outside prob of [%s;%d,%d,%d,%d] = %g\n", 
		      grammar->nonterminals[i].c_str(),
		      s, t, u, v, chart->outside(s, t, u, v, i));
	    }
	  }
        }//end for u
      }//end for deltau
    }//end for s
  }//end for deltas

  return;
}


void Trainer::count() {

  //sentence pair probability
  pi = chart->inside(0, T, 0, V, grammar->start_symbol);

#ifdef DEBUG
  cout << "* pi = " << pi << endl;
#endif

  loglike += log(pi);

  for (int s = 0; s <= T; s++){
    for (int t = s; t <= T; t++){

      WORDTYPE *en_part = splice(en,s,t);

      for (int u = 0; u <= V; u++){
        for (int v = u; v <= V; v++){

          // skip null/null
          if (v - u == 0 && t - s == 0)
            continue;

          WORDTYPE *cn_part = splice(cn,u,v);

          //a rectangle (s, t, u, v)
          for (int lhs = 0; lhs < grammar->_nonterminal_cnt; lhs++)  {
            double alpha = chart->outside(s, t, u, v, lhs);
            if (alpha <= 0.0)  continue;

            double beta = chart->inside(s, t, u, v, lhs);
            if (beta <= 0.0)  continue;

            double alb = alpha*beta;

            //terminals
            if (lhs == grammar->T)  {
              double pr = alpha * grammar->phr2phr_prob(en_part, cn_part) / pi;
	      if (Globals::Get<int>("debuglevel") >= 3) {
		fprintf(logfile, "- span = %d,%d,%d,%d\n", s, t, u, v);
		fprintf(logfile, "** N_cnt[T] = %g\n", grammar->N_cnt[lhs]);
		fprintf(logfile, "** alpha = %g\n", alpha);
		fprintf(logfile, "** beta = %g\n", beta);
		fprintf(logfile, "** pi = %g\n", pi);
	      }
              grammar->phr2phr_cntsum(en_part, cn_part, pr);

            //nonterminals
            } else {
              //update the count for the lhs
              grammar->N_cnt[lhs] += alb/pi;

              //rules
              for (int rhsl = 0; rhsl < grammar->_nonterminal_cnt; rhsl++){

                //UNARY rules
                double pr = grammar->rulprob(lhs, rhsl);
                if (pr > 0) {
                  double beta0 = chart->inside(s, t, u, v, rhsl);
                  grammar->rulcntsum(lhs, rhsl, alpha * pr * beta0 / pi);
                }

                for (int rhsr = 0; rhsr < grammar->_nonterminal_cnt; rhsr++){
                  //BINARY rules
                  double sPr = grammar->rulprob(lhs, rhsl, rhsr, 0);
                  double iPr = grammar->rulprob(lhs, rhsl, rhsr, 1);

                  if (! (sPr > 0.0 || iPr > 0.0))  continue;

                  double sBetasum = 0.0;
                  double iBetasum = 0.0;
                  for (int S = s; S <= t; S++){
                    for (int U = u; U <= v; U++){

                      //beta(sSuU, rhsl)
                      double beta1 = chart->inside(s, S, u, U, rhsl);
                      //beta(sSUv, rhsl)
                      double beta3 = chart->inside(s, S, U, v, rhsl);

                      //beta(StUv, rhsr)
                      double beta2 = chart->inside(S, t, U, v, rhsr);
                      //beta(StuV, rhsr)
                      double beta4 = chart->inside(S, t, u, U, rhsr);

                      sBetasum += beta1 * beta2;
                      iBetasum += beta3 * beta4;

                    }//end for splitting 
                  }
                  grammar->rulcntsum(lhs,rhsl,rhsr,0,alpha*sPr*sBetasum/pi);
                  grammar->rulcntsum(lhs,rhsl,rhsr,1,alpha*iPr*iBetasum/pi);
		  if (Globals::Get<int>("debuglevel") >= 3) {
		    fprintf(logfile, "- span = %d,%d,%d,%d\n", s, t, u, v);
		    fprintf(logfile, "** alpha = %g\n", alpha);
		    fprintf(logfile, "** sBetasum = %g, sPr = %g\n", sBetasum, sPr);
		    fprintf(logfile, "** iBetasum = %g, iPr = %g\n", iBetasum, iPr);
		    fprintf(logfile, "** pi = %g\n", pi);
		  }
                
                }//end for rhs.r
              }//end for rhs.l
            }//end if
          }//end for lhs

          delete[] cn_part;

        }//end for v
      }//end for u

      delete[] en_part;

    }//end for t
  }//end for s

#ifdef CLONECODE
  //cloning
  //copy values to the grammar for dumping / normalizing
  grammar->clone_mass += chart->clone_mass;
  grammar->total_mass += chart->total_mass;
//   printf("* Copying to grammar: clone = %.5f, total = %.5f, pi = %.5f\n",
//          grammar->clone_mass, grammar->total_mass, pi);
#endif

}

void Trainer::run()
{
  int sent1 = Globals::Get<int>("sentence_from");
  int sentN = Globals::Get<int>("sentence_to");

  for (int iter = 0; iter < MAXITER; iter++){
   // printf("Iteration %d\n", iter+1);
    fprintf(logfile, "Iteration %d\n", iter+1);
    fflush(logfile);
    corpus->reset();//move the scanner to the beginnings of the corpra
    _sentno=0;//number of sentence pairs processed so far
    while (readin(corpus)){//one pair
      _sentno++;
      if (sent1 && _sentno < sent1)
	continue;
      if (sentN && _sentno > sentN) {
	cout << endl;
	cout << "*************************************************\n";
	cout << "Quitting after processing sentences " 
	     << sent1 << " to " << sentN << endl;
	cout << "*************************************************\n";
	break;
      }

      gettimeofday(&start, 0);
      tinit();//initialzation of the training mode chart
      parse();
      count();
      gettimeofday(&finish, 0);

      long compTime = (finish.tv_sec - start.tv_sec);
      fprintf(logfile, "pair %d: <%d, %d>\t%ld seconds, pi = %.20e\n", 
              _sentno, T, V, compTime, pi);
      fflush(logfile);
    }
    fprintf(logfile, "log likelihood = %.20e\n", loglike);
    fflush(logfile);
    dumpbefore();//dump the expected counts

    if (! Globals::Get<bool>("distributed_EM") || MAXITER > 1 ) 
      update(); // a) normalizes grammar, b) resets loglikelihood
  }

  if (! Globals::Get<bool>("distributed_EM")) 
     dumpafter();// un-comment if to be run in single machine mode

}
