// file: bw_comp_1.cc
//

// isip include files
//
#include "bw_train.h"
#include "bw_train_constants.h"

// method: comp_back_prob_cc
//
// arguments:
//  float_8*** back_prob : (output) the forward probabilities
//  float_8& utt_prob : (output) the utterance probability
//  float_8* max_back : (output) the maximum value at each frame
//  float_8** max_mback : (output) the maximum value for each model at each
//                                 frame
//  int_2* upper : (input) the upper model bound
//  int_2* lower : (input) the lower model bound
//  float_4 width : (input) the beam width
//  int_4 num_vect : (input) the number of total feature vectors
//  float_4*** transitions : (input) the old transition matrix
//  State** states : (input) the whole states sequence
//  int_4 num_mod : (input) the number of input models
//  int_4* model_list : (input) the index of each model
//  Model** models : (input) the whole models sequence
//  int_4* trans_map : (input) the mapping from model index to trans index
//  float_8** vectors : (input) the observations
//  int_4** st_map : (input) the mapping from model index and local state
//                             index to global state index
//  int_4 num_phy_st : (input) the number of total physical states
//  float_8*** state_scores : (output/input) the state scores array
//
// return: a logical flag to indicate success
//
// this method computes backward probability by looping all the time frames,
// all the models and all the states within the model
//
logical_1 comp_back_prob_cc(float_8*** back_prob_a, float_8& utt_prob_a,
			    float_8* max_back_a, float_8** max_mback_a,
			    int_2* upper_a,
			    int_2* lower_a, float_4 beam_width_a,
			    int_4 num_vect_a,
			    float_4*** transitions_a, Train_State** states_a,
			    int_4 num_mod_a, int_4* model_list_a,
			    Train_Model** models_a, int_4* trans_map_a,
			    float_8** vectors_a, int_4** st_map_a,
			    int_4 num_phy_st_a, float_8*** state_scores_a) {
  
  // local variables
  //
  int_4 loc_st = (int_4)0;
  int_4 pre_loc_st = (int_4)0;
  int_4 tmp_mix = (int_4)0;
  float_8 tmp_score = (float_8)0.0;
  float_8 tmp_trans = (float_8)0.0;
  float_8 tmp_trans1 = (float_8)0.0;
  float_8 tmp_prob = (float_8)0.0;
  int_4 start_model = 0;
  int_4 end_model = 0;
  int_2 max_model_ind = 0;
  logical_1 yes_prune = ISIP_TRUE;
  
  // initialization
  //
  utt_prob_a = BW_LOG_ZERO;
  
  // start computing the backward probability, loop over the time
  //
  for (int_4 t = num_vect_a; t >= 1; t--) {

    // set the model bounds for this particular
    // time frame
    //
    int_2 temp;
    if (t == num_vect_a) {
      temp = num_mod_a;
    }
    else {
      temp = upper_a[t+1];
    }

    // since we hypothesize from the end the index of start model
    // value will be greater than that of end model value
    // 
    start_model = temp;
    end_model = lower_a[t];
    
    // loop over all models
    //
    for (int_2 q = start_model; q >= end_model; q--) {
      
      // loop over the states for each model
      //

      // get the number of states for the next model if the
      // current model is not the last model
      //
      if (q < num_mod_a) {
	pre_loc_st = models_a[model_list_a[q]]->get_num_states_cc();
      }

      // get the number of states for current model under evaluation
      //
      loc_st = models_a[model_list_a[q-1]]->get_num_states_cc();
      for (int_4 i = loc_st; i >= 1; i--) {
	
	// set the initial conditions
	// check if this is the last time frame
	//
	if (t == num_vect_a) {
	  
	  // if this is the last time frame
	  // backward probability of last model's last state in the
	  // last time frame is 1.0
	  //
	  if ((q == num_mod_a) && (i == loc_st)) {
	    back_prob_a[t][q][i] = BW_LOG_ONE;
	  }

	  // if this is not the last model
	  //
	  else {
	    
	    // check if this the final state
	    //
	    if (i == loc_st) {
	      tmp_trans = transitions_a[trans_map_a[model_list_a[q+1]]]
		[0][pre_loc_st-1];
	      
	      back_prob_a[t][q][i] = back_prob_a[t][q+1][pre_loc_st] +
		tmp_trans;
	    }

	    // check if this is the start state for a model
	    //
	    else if (i == 1) {
	      
	      // compute the summation from state 2 to state N-1
	      //
	      tmp_score = BW_LOG_ZERO;
	      for (int_4 j = 2; j <= loc_st-1; j++) {
		
		// compute the output probability
		//
		if (state_scores_a[t][q][j] != BW_LOG_ZERO) {
		  tmp_prob = state_scores_a[t][q][j];
		}
		else {

		  // evaluate the state score
		  //
		  tmp_prob = states_a[st_map_a[model_list_a[q-1]][j-1]]
		    ->eval_score_cc(vectors_a[t-1], t, tmp_mix);
		  state_scores_a[t][q][j] = tmp_prob;
		}

		// compute the transition probability
		//
		tmp_trans = transitions_a[trans_map_a[model_list_a[q-1]]]
		  [0][j-1];
		tmp_score = log_add_cc(tmp_score, (tmp_trans + tmp_prob +
						   back_prob_a[t][q][j]));
	      }
	      back_prob_a[t][q][i] = tmp_score;
	    }
	    else {
	      
	      // normal condition
	      // 
	      tmp_trans = transitions_a[trans_map_a[model_list_a[q-1]]]
		[i-1][loc_st-1];
	      back_prob_a[t][q][i] = tmp_trans + back_prob_a[t][q][loc_st];
	    }
	  }
	}
	
	// compute the major part of backward probability
	// when not in last time frame
	//
	else {

	  // if in last model's last state and
	  // not in last time frame
	  //
	  if ((q == num_mod_a) && (i == loc_st)) {
	    back_prob_a[t][q][i] = BW_LOG_ZERO;
	  }
	  else {

	    // if in last state of a model other than the
	    // last model
	    //
	    if (i == loc_st) {
	      
	      // compute the transition
	      //
	      tmp_trans = transitions_a[trans_map_a[model_list_a[q]]]
		[0][pre_loc_st-1];
	      
	      tmp_score = log_add_cc(back_prob_a[t+1][q+1][1],
				     (back_prob_a[t][q+1][pre_loc_st] +
				      tmp_trans));
	      
	      back_prob_a[t][q][i] = tmp_score;
	    }

	    // if this the first state and not in last
	    // time frame
	    //
	    else if (i == 1) {
	      
	      // compute the summation
	      //
	      tmp_score = BW_LOG_ZERO;
	      for (int_4 j = 2; j <= loc_st-1; j++) {
		
		// compute the transition
		//
		tmp_trans = transitions_a[trans_map_a[model_list_a[q-1]]]
		  [0][j-1];
		
		// compute the output probability
		//
		if (state_scores_a[t][q][j] != BW_LOG_ZERO) {
		  tmp_prob = state_scores_a[t][q][j];
		}
		else {
		  tmp_prob = states_a[st_map_a[model_list_a[q-1]][j-1]]
		    ->eval_score_cc(vectors_a[t-1], t, tmp_mix);
		  state_scores_a[t][q][j] = tmp_prob;
		}

		tmp_score = log_add_cc(tmp_score, (tmp_trans + tmp_prob +
						   back_prob_a[t][q][j]));
	      }
	      back_prob_a[t][q][i] = tmp_score;
	    }
	    else {
	      // if in emitting states of a some model 
	      //
	      // compute the first transition
	      //
	      tmp_trans = transitions_a[trans_map_a[model_list_a[q-1]]]
		[i-1][loc_st-1];
	      
	      // compute the summation
	      //
	      tmp_score = BW_LOG_ZERO;
	      for (int_4 j = 2; j <= loc_st-1; j++) {
		
		// compute the second transition
		//
		tmp_trans1 = transitions_a[trans_map_a[model_list_a[q-1]]]
		  [i-1][j-1];
		
		// compute the output probability
		//
		if (state_scores_a[t+1][q][j] != BW_LOG_ZERO) {
		  tmp_prob = state_scores_a[t+1][q][j];
		}
		else {
		  tmp_prob = states_a[st_map_a[model_list_a[q-1]][j-1]]
		    ->eval_score_cc(vectors_a[t], t+1, tmp_mix);
		  state_scores_a[t+1][q][j] = tmp_prob;
		}

		tmp_score = log_add_cc(tmp_score, (tmp_trans1 + tmp_prob +
						   back_prob_a[t+1][q][j]));
	      }
	      tmp_score = log_add_cc(tmp_score, (tmp_trans +
						 back_prob_a[t][q][loc_st]));
	      back_prob_a[t][q][i] = tmp_score;
	    }
	  }
	} // end of computing major part of backward probability
	
	// find the maximum value for this model
	//
	if (max_mback_a[t][q] < back_prob_a[t][q][i]) {
	  max_mback_a[t][q] = back_prob_a[t][q][i];
	}
      } // end of looping over each state
      
      // find the maximum value at this frame
      //
      if (max_back_a[t] < max_mback_a[t][q]) {
	max_back_a[t] = max_mback_a[t][q];
	max_model_ind = q;
      }
    } // end of looping over each model
    
    // beta pruning
    // if the difference is greater than beam width then do not take that
    // model into account while calculating the backward probability.
    // so set it to NULL
    //
    yes_prune = ISIP_FALSE;
    while (max_back_a[t] - max_mback_a[t][start_model] > beam_width_a) {
      max_mback_a[t][start_model] = BW_LOG_ZERO;
      yes_prune = ISIP_TRUE;
      start_model--;
      upper_a[t] = start_model;
    }

    // decrement the start model if required
    //
    while (start_model > upper_a[t]) {
      start_model--;
    }

    // reset the model boundaries
    //
    upper_a[t] = start_model;
    
    while (max_back_a[t] - max_mback_a[t][end_model] > beam_width_a) {
      yes_prune = ISIP_TRUE;
      max_mback_a[t][end_model] = BW_LOG_ZERO;

      // the backward probability needs to be reset to BW_LOG_ZERO
      // for each state of pruned end model
      //
      int_4 num_st = models_a[model_list_a[end_model-1]]->get_num_states_cc();
      for ( int_4 i = 0; i <= num_st; i++) {
	back_prob_a[t][end_model][i] = BW_LOG_ZERO;
      }
      lower_a[t-1] = end_model;      
      end_model++;
    }
  } // end of looping over each frame
  
  // assign the utterance probability
  //
  if (back_prob_a[1][1][1] != BW_LOG_ZERO) {
    utt_prob_a = back_prob_a[1][1][1];
  }
  else {
    fprintf(stdout, "Warning: beta computation failed.\n\n");
    return(ISIP_FALSE);
  }
  
  // exit gracefully
  //
  return(ISIP_TRUE);
}
