// file: ht_update_0.cc
//
// isip include files
//
#include "hmm_train.h"
#include "hmm_train_constants.h"

// method: update_cc
//
// arguments:
//  float_4*** transitions : (output) the transition data need to be updated
//  Train_State** states : (output) the states data need to be updated
//  float_8*** mean : (input) the sum of data in each state
//  float_8*** covar : (input) the sum of the covar in each state
//  int_4** st_count : (input) the count of the all states
//  int_*** trans_count : (input) the appearing time of the arcs in one
//                                transition matrix
//  int_4* trans_size : (input) size of each transition matrix
//  int_4 num_feat : (input) the number of the features
//  int_4 num_st : (input) the number of the states
//  int_4 num_mix : (input) the number of the mixtures
//  int_4 num_trans : (input) the number of the transition
//  float_8 var_floor : (input) variance floor
//
// return: a logical flag to indicate success
//
logical_1 update_cc(float_4*** transitions_a, Train_State** states_a,
		    float_8*** mean_a, float_8*** covar_a,
		    int_4** st_count_a, int_4*** trans_count_a,
		    int_4* trans_size_a, int_4 num_feat_a,
		    int_4 num_st_a, int_4 num_mix_a,
		    int_4 num_trans_a, float_8 var_floor_a) {

  // local variable
  //
  int_4 sum = (int_4)0;
  float_4* temp_scale = new float_4[num_mix_a];
  float_4** temp_mean = new float_4*[num_mix_a];
  float_4** temp_covar = new float_4*[num_mix_a];
  float_4* temp_weight = new float_4[num_mix_a];
  int_4 weight_sum = (int_4)0;

  float_4* temp_scale1 = (float_4*)NULL;
  float_4** temp_mean1 = (float_4**)NULL;
  float_4** temp_covar1 = (float_4**)NULL;
  float_4* temp_weight1 = (float_4*)NULL;
    
  for (int_4 i = 0; i < num_mix_a; i++) {
    temp_mean[i] = new float_4[num_feat_a];
    temp_covar[i] = new float_4[num_feat_a];
    temp_weight[i] = (float_4)0;
  }
  
  // update the means and variances in states array
  //
  for (int_4 i = 1; i < num_st_a; i++) {

    // initialize the temp variables
    //
    temp_scale1 = states_a[i]->get_scale_cc();
    temp_mean1 = states_a[i]->get_mean_cc();
    temp_covar1 = states_a[i]->get_covar_cc();
    temp_weight1 = states_a[i]->get_weights_cc();

    // initialize the means and variances to values from previous pass
    // of estimation
    //
    memcpy(temp_scale, temp_scale1, num_mix_a * sizeof(float_4));
    memcpy(temp_weight, temp_weight1, num_mix_a * sizeof(float_4));

    for (int_4 aa = 0; aa < num_mix_a; aa++) {
      memcpy(temp_mean[aa], temp_mean1[aa], num_feat_a * sizeof(float_4));
      memcpy(temp_covar[aa], temp_covar1[aa], num_feat_a * sizeof(float_4));
    }

    // compute the new mean and variance
    //
    for (int_4 j = 0; j < num_mix_a; j++) {
      for (int_4 l = 0; l < num_feat_a; l++) {

	// accumulate only if mixture has been accessed atleast once
	//
	if (st_count_a[i][j] > (int_4)0) {

	  // compute the correct covariance by subtracting the square
	  // of the mean
	  //
	  mean_a[i][j][l] = mean_a[i][j][l]/(float_8)st_count_a[i][j];
	  covar_a[i][j][l] = covar_a[i][j][l]/(float_8)st_count_a[i][j];
	  covar_a[i][j][l] = (covar_a[i][j][l]) -
	    (mean_a[i][j][l]*mean_a[i][j][l]);
	}
      }
    }

    // update the mean and variance for the current state
    //
    for (int_4 tt = 0; tt < num_mix_a; tt++) {

      // accumulate only if mixture has been accessed atleast once
      //
      if (st_count_a[i][tt] > (int_4)0) {
	temp_scale[tt] = (float_4)0.0;
	for (int_4 ss = 0; ss < num_feat_a; ss++) {
	  temp_mean[tt][ss] = (float_4)mean_a[i][tt][ss];
	  
	  // accumulate covariance only if mixture has been accessed
	  // atleast twice
	  //
	  if (st_count_a[i][tt] > (int_4)1) {

	    // variance flooring
	    //
	    if (covar_a[i][tt][ss] < var_floor_a) {
	      temp_covar[tt][ss] = 1.0 / var_floor_a;
	    } 
	    else {
	      temp_covar[tt][ss] = (float_4)1/(float_4)covar_a[i][tt][ss];
	    }
	  }

	  // accumulate scale factor
	  //
	  temp_scale[tt] += log((1.0/temp_covar[tt][ss]) * ISIP_TWOPI);
	}
      }
    }    
    
    // update the mixture weight for all the states
    //
    weight_sum = (int_4)0;
    for (int_4 mm = 0; mm < num_mix_a; mm++) {
      weight_sum += st_count_a[i][mm];
    }

    if (weight_sum > (int_4)0) {
      for (int_4 nn = 0; nn < num_mix_a; nn++) {

	// find new weight only if mixture has been accessed
	// (note that this does not guarantee weights summing to unity)
	//
	if (st_count_a[i][nn] > (int_4)0) {
	  temp_weight[nn] =
	    log((float_4)st_count_a[i][nn]/(float_4)weight_sum);
	}
      }
    }
    
    // update state information if state has been accessed
    //
    if(weight_sum > (int_4)1) {
      states_a[i]->set_mean_cc(temp_mean);
      states_a[i]->set_covar_cc(temp_covar);
      states_a[i]->set_scale_cc(temp_scale);
      states_a[i]->set_weights_cc(temp_weight);
    }

    // reset pointers
    //
    temp_scale1 = (float_4*)NULL;
    temp_mean1 = (float_4**)NULL;
    temp_covar1 = (float_4**)NULL;
    temp_weight1 = (float_4*)NULL;
  }
  

  // update the transition matrix
  //
  for (int_4 i = 1; i < num_trans_a; i++) {
    for (int_4 start_st = 0; start_st < trans_size_a[i]; start_st++) {
      for (int_4 end = 0; end < trans_size_a[i]; end++) {
	sum += trans_count_a[i][start_st][end];
      }
      for (int_4 end_st = 0; end_st < trans_size_a[i]; end_st++) {
	if (sum > (int_4)0) {
	  transitions_a[i][start_st][end_st] =
	    (float_4)trans_count_a[i][start_st][end_st]/(float_4)sum;
	}
       	else {
	  transitions_a[i][start_st][end_st] =
	    exp(transitions_a[i][start_st][end_st]);
       	}
      }
      
      // reset the sum
      //
      sum = (int_4)0;
    }
  }

 
  // free memory
  //
  for (int_4 i = 0; i < num_mix_a; i++) {
    delete [] temp_mean[i];
    delete [] temp_covar[i];
  }
  delete [] temp_mean;
  delete [] temp_covar;
  delete [] temp_scale;
  delete [] temp_weight;
  
  // exit gracefully
  //
  return ISIP_TRUE;
}
























