/*******************************************************************************
*
*       This file is part of the General Hidden Markov Model Library,
*       GHMM version __VERSION__, see http://ghmm.org
*
*       Filename: ghmm/ghmm/reestimate.c
*       Authors:  Bernhard Knab, Benjamin Georgi, Janne Grunau
*
*       Copyright (C) 1998-2004 Alexander Schliep
*       Copyright (C) 1998-2001 ZAIK/ZPR, Universitaet zu Koeln
*       Copyright (C) 2002-2004 Max-Planck-Institut fuer Molekulare Genetik,
*                               Berlin
*
*       Contact: schliep@ghmm.org
*
*       This library is free software; you can redistribute it and/or
*       modify it under the terms of the GNU Library General Public
*       License as published by the Free Software Foundation; either
*       version 2 of the License, or (at your option) any later version.
*
*       This library is distributed in the hope that it will be useful,
*       but WITHOUT ANY WARRANTY; without even the implied warranty of
*       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
*       Library General Public License for more details.
*
*       You should have received a copy of the GNU Library General Public
*       License along with this library; if not, write to the Free
*       Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
*
*
*       This file is version $Revision: 1961 $
*                       from $Date: 2007-11-07 12:10:10 -0500 (Wed, 07 Nov 2007) $
*             last change by $Author: grunau $.
*
*******************************************************************************/

#ifdef HAVE_CONFIG_H
#  include "../config.h"
#endif

#include <float.h>
#include <math.h>

#include "ghmm.h"
#include "mes.h"
#include "mprintf.h"
#include "reestimate.h"
#include "matrix.h"
#include "model.h"
#include "foba.h"
#include "ghmm_internals.h"

typedef struct local_store_t {
  double *pi_num;
  double pi_denom;
  double **a_num;
  double *a_denom;
  double **b_num;
  double **b_denom;
} local_store_t;


static double nologSum(double* vec, int len) {
  int i;
  double retval=vec[0];
  
  for (i=1; i<len; i++)
    retval += vec[i];

  return retval;
}


/*----------------------------------------------------------------------------*/
static int reestimate_free (local_store_t ** r, int N)
{
# define CUR_PROC "reestimate_free"
  int i;
  mes_check_ptr (r, return (-1));
  if (!*r)
    return (0);
  m_free ((*r)->pi_num);

  if ((*r)->a_num){
    for (i = 0; i < N; i++){
      m_free ((*r)->a_num[i]);
    }  
  }  
  m_free ((*r)->a_num);
  m_free ((*r)->a_denom);

  if ((*r)->b_num){
    for (i = 0; i < N; i++){
      m_free ((*r)->b_num[i]);
    }
  }
  m_free ((*r)->b_num);
  if ((*r)->b_denom){
    for (i = 0; i < N; i++){
      m_free ((*r)->b_denom[i]);
    }  
  }
  m_free ((*r)->b_denom);

  m_free (*r);
  return (0);
# undef CUR_PROC
}                               /* reestimate_free */

/*----------------------------------------------------------------------------*/
static local_store_t *reestimate_alloc (const ghmm_dmodel * mo)
{
# define CUR_PROC "reestimate_alloc"
  int i;
  local_store_t *r = NULL;
  ARRAY_CALLOC(r, 1);

  ARRAY_CALLOC(r->pi_num, mo->N);
  ARRAY_CALLOC(r->a_num, mo->N);
  for (i=0; i<mo->N; i++)
    ARRAY_CALLOC(r->a_num[i], mo->s[i].out_states);
  ARRAY_CALLOC(r->a_denom, mo->N);

  ARRAY_MALLOC(r->b_num, mo->N);
  ARRAY_MALLOC(r->b_denom, mo->N);

  if (mo->model_type & GHMM_kHigherOrderEmissions)
    for (i=0; i<mo->N; i++) {
      ARRAY_CALLOC(r->b_num[i], ghmm_ipow(mo, mo->M, mo->order[i] + 1));
      ARRAY_CALLOC(r->b_denom[i], ghmm_ipow (mo, mo->M, mo->order[i]));
    }
  else
    for (i=0; i<mo->N; i++) {
      ARRAY_CALLOC(r->b_num[i], mo->M);
      ARRAY_CALLOC(r->b_denom[i], 1);
    }

  return (r);
STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  reestimate_free (&r, mo->N);
  return (NULL);
# undef CUR_PROC
}                               /* reestimate_alloc */

/*----------------------------------------------------------------------------*/
static int reestimate_init(local_store_t * r, const ghmm_dmodel * mo) {
# define CUR_PROC "reestimate_init"

  int i, j, m, size;

  r->pi_denom = 0.0;

  for (i=0; i<mo->N; i++) {
    r->pi_num[i] = 0.0;
    r->a_denom[i] = 0.0;
    for (j=0; j<mo->s[i].out_states; j++)
      r->a_num[i][j] = 0.0;

    if (mo->model_type & GHMM_kHigherOrderEmissions) {
      size = ghmm_ipow(mo, mo->M, mo->order[i]);
      for (m=0; m<size; m++)
	r->b_denom[i][m] = 0.0;
      size *= mo->M;
      for (m=0; m<size; m++)
	r->b_num[i][m] = 0.0;
    } else {
      r->b_denom[i][0] = 0.0;
      for (m=0; m<mo->M; m++)
	r->b_num[i][m] = 0.0;
    }
  }
  return (0);
# undef CUR_PROC
}                               /* reestimate_init */

/*----------------------------------------------------------------------------*/
int ighmm_reestimate_alloc_matvek (double ***alpha, double ***beta, double **scale,
                             int T, int N)
{
# define CUR_PROC "ighmm_reestimate_alloc_matvek"
  int res = -1;
  *alpha = ighmm_cmatrix_stat_alloc (T, N);
  if (!(*alpha)) {
    GHMM_LOG_QUEUED(LCONVERTED);
    goto STOP;
  }
  *beta = ighmm_cmatrix_stat_alloc (T, N);
  if (!(*beta)) {
    GHMM_LOG_QUEUED(LCONVERTED);
    goto STOP;
  }
  ARRAY_CALLOC (*scale, T);
  res = 0;
STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  return (res);
# undef CUR_PROC
}                               /* ighmm_reestimate_alloc_matvek */

/*----------------------------------------------------------------------------*/
int ighmm_reestimate_free_matvek (double **alpha, double **beta, double *scale, int T)
{
# define CUR_PROC "ighmm_reestimate_free_matvek"
  ighmm_cmatrix_stat_free (&alpha);
  ighmm_cmatrix_stat_free (&beta);
  m_free (scale);
  return (0);
# undef CUR_PROC
}                               /* ighmm_reestimate_free_matvek */

/*----------------------------------------------------------------------------*/
void ghmm_dmodel_update_tie_groups (ghmm_dmodel * mo) {
#define CUR_PROC "ghmm_dmodel_update_tied_groups"
  int i, j, k;
  int bi_len;
  int nr=0;
  double *new_emissions;

  /* do nothing if there are no tied emissions */
  if (!(mo->model_type & GHMM_kTiedEmissions)) {
    GHMM_LOG(LWARN, "No tied emissions. Exiting.");
    return;
  }
  
  if (mo->model_type & GHMM_kHigherOrderEmissions) {
    ARRAY_MALLOC(new_emissions, ghmm_ipow(mo, mo->M, mo->maxorder+1));
  }
  else {
    ARRAY_MALLOC(new_emissions, mo->M);
  }

  for (i=0; i<mo->N; i++) {

    /* find tie group leaders */
    if (mo->tied_to[i] == i) {


      if (mo->model_type & GHMM_kHigherOrderEmissions)
	bi_len = ghmm_ipow(mo, mo->M, mo->order[i] + 1);
      else
	bi_len = mo->M;

      if (mo->model_type & GHMM_kSilentStates && mo->silent[i]) {
	GHMM_LOG_PRINTF(LWARN, LOC, "Tie group leader %d is silent.", i);
	nr = 0;
	/* initializing with zeros */
	for (k=0; k<bi_len; k++)
	  new_emissions[k] = 0.0;
      }
      else {
	nr = 1;
	/* initializing with tie group leader emissions */
	for (k=0; k<bi_len; k++)
	  new_emissions[k] = mo->s[i].b[k];
      }

      /* finding tie group members */
      for (j=i+1; j<mo->N; j++) {
        if (mo->tied_to[j] == i && (!(mo->model_type & GHMM_kHigherOrderEmissions)
				    || mo->order[i] == mo->order[j])) {
	  /* silent states have no contribution to the pooled emissions within a group */
	  if (!(mo->model_type & GHMM_kSilentStates) || (mo->silent[j] == 0)) {
            nr += 1;
            /* printf("  tie group member %d -> leader %d.\n",j,i); */
            /* summing up emissions in the tie group */
            for (k=0; k<bi_len; k++)
              new_emissions[k] += mo->s[j].b[k];
          }
          else {
	    GHMM_LOG_PRINTF(LWARN, LOC, "Tie group member %d is silent.", j);
	  }
        }
      }
      
      /* updating emissions */
      if (nr > 1)
        for (j=i; j < mo->N; j++) {
          /* states within one tie group are required to have the same order */
          if (mo->tied_to[j] == i && (!(mo->model_type & GHMM_kHigherOrderEmissions)
				      || mo->order[i] == mo->order[j])
	      && (!(mo->model_type & GHMM_kSilentStates) || (mo->silent[j] == 0)))
            for (k = 0; k < bi_len; k++) {
              mo->s[j].b[k] = new_emissions[k] / nr;
              /* printf("s(%d)[%d] -> %f / %f = %f\n", j, k, new_emissions[k], nr,mo->s[j].b[k]);   */
            }
	}
      else {
	GHMM_LOG_PRINTF(LINFO, LOC, "The tie group with leader %d has only one non-silent state. Kind of pointless!", i);
      }
    }
  }

STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  m_free (new_emissions);
#undef CUR_PROC
}           /* ghmm_dmodel_update_tied_groups */

/*----------------------------------------------------------------------------*/
static int reestimate_setlambda (local_store_t * r, ghmm_dmodel * mo)
{
# define CUR_PROC "reestimate_setlambda"
  int res = -1;
  int h, i, j, m, l, j_id, positive;
  double factor, p_i;
  int hist, col, size, reachable;

  mes_check_0 (r->pi_denom, goto STOP);

  for (i = 0; i < mo->N; i++) {
    reachable = 1;
    positive = 0;

    /* Pi */
    mo->s[i].pi = r->pi_num[i] / r->pi_denom;

    /* A */
    /* note: denom. might be 0; never reached state? */
    p_i = 0.0;
    if (r->a_denom[i] < GHMM_EPS_PREC) {
      for (h = 0; h < mo->s[i].in_states; h++) {
        p_i += mo->s[i].in_a[h];
      }
      if (p_i == 0.0) {
        if (h == 0)
          GHMM_LOG_PRINTF(LINFO, LOC, "State %d can't be reached (no in_states)", i);
        else
          GHMM_LOG_PRINTF(LINFO, LOC, "State %d can't be reached (prob = 0.0)", i);
        reachable = 0;
      }
      factor = 0.0;
    }
    else
      factor = (1 / r->a_denom[i]);


    for (j = 0; j < mo->s[i].out_states; j++) {
      /* TEST: denom. < numerator */
      if ((r->a_denom[i] - r->a_num[i][j]) <= -GHMM_EPS_PREC) {
        GHMM_LOG(LCONVERTED, " numerator > denom!\n");
      }
      mo->s[i].out_a[j] = r->a_num[i][j] * factor;
      if (r->a_num[i][j] >= GHMM_EPS_PREC)
        positive = 1;
      /* important: also update in_a  */
      l = 0;
      j_id = mo->s[i].out_id[j];
      while (l < mo->s[j_id].in_states)
        if (mo->s[j_id].in_id[l] == i)
          break;
        else
          l++;
      if (l == mo->s[j_id].in_states) {
        GHMM_LOG(LCONVERTED, "no corresponding in_a to out_a found!\n");
        goto STOP;
      }
      mo->s[j_id].in_a[l] = mo->s[i].out_a[j];
    }

    /*if (!positive) {
       str = ighmm_mprintf(NULL, 0, 
       "All numerator a[%d][j] == 0 (denom=%.4f, P(in)=%.4f)!\n", 
       i, r->a_denom[i], p_i);
       GHMM_LOG(LCONVERTED, str);
       m_free(str);
       } */

    /* if fix, continue to next state */
    if (mo->s[i].fix)
      continue;

    /* B */
    if (mo->model_type & GHMM_kHigherOrderEmissions)
      size = ghmm_ipow (mo, mo->M, mo->order[i]);
    else
      size = 1;
    /* If all in_a's are zero, the state can't be reached.
       Set all b's to -1.0 */
    if (!reachable) {
      for (hist = 0; hist < size; hist++) {
        col = hist * mo->M;
        for (m = col; m < col + mo->M; m++) {
          mo->s[i].b[m] = -1.0;
        }
      }
    }
    else {
      for (hist = 0; hist < size; hist++) {
        /* If the denominator is very small, we have not seen many emissions
           in this state with this history.
           We are conservative and just skip them. */
        if (r->b_denom[i][hist] < GHMM_EPS_PREC)
          continue;
        else
          factor = (1.0 / r->b_denom[i][hist]);

        positive = 0;
        /* TEST: denom. < numerator */
        col = hist * mo->M;
        for (m = col; m < col + mo->M; m++) {
          if ((r->b_denom[i][hist] - r->b_num[i][m]) <= -GHMM_EPS_PREC) {
            GHMM_LOG_PRINTF(LCONVERTED, LOC, "numerator b (%.4f) > denom (%.4f)!\n",
                            r->b_num[i][m], r->b_denom[i][hist]);
          }
          
          mo->s[i].b[m] = r->b_num[i][m] * factor;
          if (mo->s[i].b[m] >= GHMM_EPS_PREC)
            positive = 1;
        }

        if (!positive) {
          GHMM_LOG_PRINTF(LCONVERTED, LOC, "All numerator b[%d][%d-%d] == 0 (denom = %g)!\n",
                          i, col, col + mo->M, r->b_denom[i][hist]);
        }
      }                           /* for each history */
    }
  }                             /* for (i = 0 .. < mo->N)  */

  res = 0;
  if (mo->model_type & GHMM_kTiedEmissions)
    ghmm_dmodel_update_tie_groups (mo);
STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  return (res);
# undef CUR_PROC
}                               /* reestimate_setlambda */

/*----------------------------------------------------------------------------*/
static int reestimate_one_step (ghmm_dmodel * mo, local_store_t * r, int seq_number,
				int *seq_length, int **O, double *log_p,
				double *seq_w)
{
# define CUR_PROC "reestimate_one_step"
  int res = -1;
  int k, i, j, t, j_id, valid=0;
  int e_index;
  int errors;
  double **alpha = NULL;
  double **beta = NULL;
  double *scale = NULL;
  int T_k=0;
  double gamma;
  double log_p_k;

  /* first set maxorder to zero if model_type & kHigherOrderEmissions is FALSE 

     TODO XXX use model->maxorder only
     if model_type & kHigherOrderEmissions is TRUE */

  if (!(mo->model_type & GHMM_kHigherOrderEmissions))
    mo->maxorder = 0;

  *log_p = 0.0;
  /* loop over all sequences */
  for (k = 0; k < seq_number; k++) {
    mo->emission_history = 0;
    T_k = seq_length[k];        /* current seq. length */

    /* initialization of  matrices and vector depends on T_k */
    if (ighmm_reestimate_alloc_matvek (&alpha, &beta, &scale, T_k, mo->N) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto FREE;
    }
    if (ghmm_dmodel_forward (mo, O[k], T_k, alpha, scale, &log_p_k) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto FREE;
    }

    if (log_p_k != +1) {        /* O[k] can be generated */
      *log_p += log_p_k;
      valid = 1;
      
      if (ghmm_dmodel_backward (mo, O[k], T_k, beta, scale) == -1) {
        GHMM_LOG_QUEUED(LCONVERTED);
        goto FREE;
      }

      /* loop over all states */
      for (i = 0; i < mo->N; i++) {
        /* Pi */
        r->pi_num[i] += seq_w[k] * alpha[0][i] * beta[0][i];
        r->pi_denom  += seq_w[k] * alpha[0][i] * beta[0][i];

        for (t=0; t < T_k-1; t++) {
          /* B */
          if (!mo->s[i].fix) {
            e_index = get_emission_index(mo, i, O[k][t], t);
            if (e_index != -1) {
              gamma = seq_w[k] * alpha[t][i] * beta[t][i];
              r->b_num[i][e_index] += gamma;
              r->b_denom[i][e_index / (mo->M)] += gamma;
            }
          }
          update_emission_history(mo, O[k][t]);

          /* A */
          r->a_denom[i] += (seq_w[k] * alpha[t][i] * beta[t][i]);
          for (j=0; j < mo->s[i].out_states; j++) {
            j_id = mo->s[i].out_id[j];
            e_index = get_emission_index(mo, j_id, O[k][t+1], t+1);
            if (e_index != -1)
              r->a_num[i][j] += (seq_w[k] * alpha[t][i] * mo->s[i].out_a[j]
                                 * mo->s[j_id].b[e_index] * beta[t+1][j_id]
                                 * (1.0 / scale[t + 1]));       /* c[t] = 1/scale[t] */
          }
        }
        /* B: last iteration for t==T_k-1 */
        t = T_k - 1;
        if (!mo->s[i].fix) {
          e_index = get_emission_index (mo, i, O[k][t], t);
          if (e_index != -1) {
            gamma = seq_w[k] * alpha[t][i] * beta[t][i];
            r->b_num[i][e_index] += gamma;
            r->b_denom[i][e_index / (mo->M)] += gamma;
          }
        }
      }
    }                           /* if (log_p_k != +1) */
    else {
      GHMM_LOG_PRINTF(LWARN, LOC, "O(%d) can't be built from model mo!\n", k);
    }

    ighmm_reestimate_free_matvek(alpha, beta, scale, T_k);
  }                             /* for (k = 0; k < seq_number; k++) */

  if (valid) {
    /* new parameter lambda: set directly in model */

    if (reestimate_setlambda(r, mo) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto STOP;
    }

    if ((errors = ghmm_dmodel_check(mo))) {
      GHMM_LOG_PRINTF(LERROR, LOC, "Reestimated model is invalid, "
                      "model_check found %d errors", -errors);
      goto STOP;
    }
  }
  else {                        /* NO sequence can be built from model mo! */
    *log_p = +1;
  }
  res = 0;

STOP:
   return (res);
FREE:
   ighmm_reestimate_free_matvek(alpha, beta, scale, T_k);
   return (res);
# undef CUR_PROC
}                               /* reestimate_one_step */


/*---------------------------------------------------------------------------*/
static int reestimate_one_step_lean (ghmm_dmodel * mo, local_store_t * r,
                                     int seq_number, int *seq_length,
                                     int **seqs, double *log_p, double *seq_w)
{
# define CUR_PROC "reestimate_one_step_lean"
  int res = -1;
  int i, t, k, e_index, len;
  int m, j, j_id, g, g_id, l, s, h;
  int size;
  int * O;
  double c_t;
  
  double * alpha_last_col=NULL;
  double * alpha_curr_col=NULL;
  double * switching_tmp;
  double scale, old_scale, scalingf;
  double * summands=NULL;
  local_store_t * * curr_est=NULL, * * last_est=NULL, * * switch_lst;

  /* allocating memory for two columns of alpha matrix */
  ARRAY_CALLOC (alpha_last_col, mo->N);
  ARRAY_CALLOC (alpha_curr_col, mo->N);

  /* allocating 2*N local_store_t */
  ARRAY_CALLOC (last_est, mo->N);
  for (i=0; i<mo->N; i++)
    last_est[i] = reestimate_alloc(mo);
  ARRAY_CALLOC (curr_est, mo->N);
  for (i=0; i<mo->N; i++)
    curr_est[i] = reestimate_alloc(mo);


  /* temporary array to hold logarithmized summands
     for sums over probabilities */
  ARRAY_CALLOC (summands, m_max(mo->N,ghmm_ipow(mo,mo->M,mo->maxorder+1))+1);


  for (k=0; k<seq_number; k++) {
    len = seq_length[k];
    O = seqs[k];
   
    ghmm_dmodel_forward_init(mo, alpha_last_col, O[0], &scale);
    if (scale < GHMM_EPS_PREC) {
      /* means: first symbol can't be generated by hmm */
      *log_p = +1;
      goto STOP;
    }
    *log_p = - log(1/scale);
    
    for (t=1; t<len; t++) {
      old_scale = scale;
      scale = 0.0;
      update_emission_history(mo, O[t-1]);
    
      /* iterate over non-silent states */
      for (i=0; i<mo->N; i++) {
	/* printf("  akt_ state %d\n",i);*/
      
	e_index = get_emission_index(mo, i, O[t], t);
	if (e_index != -1){
	  alpha_curr_col[i] = ghmm_dmodel_forward_step(&mo->s[i], alpha_last_col,
					       mo->s[i].b[e_index]);
	  scale += alpha_curr_col[i];
	}
	else
	  alpha_curr_col[i] = 0;
      }
    
      if (scale < GHMM_EPS_PREC) {
	GHMM_LOG(LERROR, "scale smaller than eps");
	/* O-string  can't be generated by hmm */
	*log_p = +1.0;
	break;
      }
      c_t = 1/scale;
      for (i=0; i<mo->N; i++)
	alpha_curr_col[i] *= c_t;
    
      /*sum log(c[t]) scaling values to get  log( P(O|lambda) ) */
      *log_p -= log(c_t);

      scalingf = 1/old_scale;
      for (m=0; m<mo->N; m++) {
	for (i=0; i<mo->N; i++) {

	  /* computes estimates for the numerator of transition probabilities */
	  for (j=0; j<mo->s[i].out_states; j++) {
	    j_id = mo->s[i].out_id[j];

	    for (g=0; g<mo->s[j_id].in_states; g++) {
	      g_id = mo->s[j_id].out_id[g];
	      e_index = get_emission_index(mo, g_id, O[t], t);
	       /* scales all summands with the current */
	      summands[g] = last_est[m]->a_num[i][j] * mo->s[j_id].in_a[g]
		* mo->s[g_id].b[e_index] * scalingf;
	    }
	    if (j_id == m) {
	      e_index = get_emission_index(mo, j_id, O[t], t);
	      /* alpha is scaled. no other scaling necessary */
	      summands[g++] = alpha_last_col[i] * mo->s[i].out_a[j]
		* mo->s[j_id].b[e_index];
	    }
	    curr_est[m]->a_num[i][j] = nologSum(summands, g);
	  }
#ifdef calculate_denominators_extra
	  /* computes denominator of transition probabilities */	  
	  for (g=0; g<mo->s[m].in_states; g++) {
	    g_id = mo->s[m].in_id[g];
	    e_index = get_emission_index(mo, m, O[t], t);
	    /* scales all summands with the current factor */
	    summands[g] = last_est[m]->a_denom[i] * mo->s[m].in_a[g]
	      * mo->s[m].b[e_index] * scalingf;
	  }
	  if (i == m) {
	    for (l=0; l<mo->s[i].out_states; l++)
	      if (mo->s[i].out_id[l] == i)
		break;
	    if (l<mo->s[i].out_states) {
	      e_index = get_emission_index(mo, i, O[t], t);
	      /* alpha is scaled. no other scaling necessary */
	      summands[++g] = alpha_last_col[i] 
		* mo->s[i].out_a[l] * mo->s[m].b[e_index];
	    }
	  }
	  curr_est[m]->a_denom[i] = nologSum(summands, g);
#endif

	  /* computes estimates for the numerator of emmission probabilities*/
	  if (mo->model_type & GHMM_kHigherOrderEmissions)
	    size = ghmm_ipow(mo,mo->M, mo->order[i]);
	  else
	    size = 1;
	  for (h=0; h<size; h++)
	    for (s=h*mo->M; s<h*mo->M+mo->M; s++) {
	      for (g=0; g<mo->s[m].in_states; g++) {
		g_id = mo->s[m].out_id[g];
		e_index = get_emission_index(mo, g_id, O[t], t);
		/* scales all summands with the last scaling factor
		   of alpha */
		summands[g] = last_est[m]->b_num[i][s] * mo->s[m].in_a[g]
		  * mo->s[g_id].b[e_index] * scalingf;
	      }
	      curr_est[m]->b_num[i][s] = nologSum(summands, g);
	    }

	  e_index = get_emission_index(mo, m, O[t], t);
	  if (i == m) {
	    for (l=0; l<mo->s[m].out_states; l++)
	      if (mo->s[m].out_id[l] == m)
		break;
	    if (l<mo->s[m].out_states)
	      /* alpha is scaled. no other scaling necessary */
	      curr_est[m]->b_num[i][e_index] += alpha_last_col[i]
		* mo->s[m].out_a[l] * mo->s[m].b[e_index];
	  }
	  
	}
      }

      /* switching pointers of alpha_curr_col and alpha_last_col */
      switching_tmp  = alpha_last_col;
      alpha_last_col = alpha_curr_col;
      alpha_curr_col = switching_tmp;

      switch_lst = last_est;
      last_est   = curr_est;
      curr_est   = switch_lst;
    }

    /* filling the usual reestimate arrays by summing all states */
    for (m=0; m<mo->N; m++) {
      curr_est[m]->pi_denom = 0;
      for (i=0; i<mo->N; i++) {
	/* PI */
	/* XXX calculate the estimates for pi numerator */
	curr_est[m]->pi_num[i] = mo->s[i].pi;
	r->pi_num[i] += seq_w[k] * curr_est[m]->pi_num[i];
	curr_est[m]->pi_denom += curr_est[m]->pi_num[i];

	/* A */
	curr_est[m]->a_denom[i] = 0;
	for (j=0; j<mo->s[i].out_states; j++) {
	  r->a_num[i][j] += seq_w[k] * curr_est[m]->a_num[i][j];
	  curr_est[m]->a_denom[i] += curr_est[m]->a_num[i][j];
	}
	r->a_denom[i] += seq_w[k] * curr_est[m]->a_denom[i];
	
	/* B */
	for (h=0; h<size; h++) {
	  curr_est[m]->b_denom[i][h] = 0;
	  for (s=h*mo->M; s<h*mo->M+mo->M; s++) {
	    r->b_num[i][s] += seq_w[k] * curr_est[m]->b_num[i][s];
	    curr_est[m]->b_denom[i][h] += curr_est[m]->b_num[i][s];
	  }
	  r->b_denom[i][h] += seq_w[k] * curr_est[m]->b_denom[i][h];
	}
      }
      /* PI */
      r->pi_denom += seq_w[k] * curr_est[m]->pi_denom;
    }
    
  }

  if (*log_p == 1.0)
    res = -1;
  else
    res = 0;

STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  /* deallocation */
  if (last_est)
    for (i=0; i<mo->N; i++)
      reestimate_free(&(last_est[i]), mo->N);
  m_free(last_est);

  if (curr_est)
    for (i=0; i<mo->N; i++)
      reestimate_free(&(curr_est[i]), mo->N);
  m_free(curr_est);

  m_free(alpha_last_col);
  m_free(alpha_curr_col);
  m_free(summands);

  return res;  
#undef CUR_PROC
}


/*============================================================================*/
int ghmm_dmodel_baum_welch (ghmm_dmodel * mo, ghmm_dseq * sq)
{
# define CUR_PROC "ghmm_dmodel_baum_welch"

  return ghmm_dmodel_baum_welch_nstep (mo, sq, GHMM_MAX_ITER_BW, GHMM_EPS_ITER_BW);
# undef CUR_PROC
}                               /* ghmm_dmodel_baum_welch */


/*============================================================================*/
int ghmm_dmodel_baum_welch_nstep (ghmm_dmodel * mo, ghmm_dseq * sq, int max_step,
                                 double likelihood_delta)
{
# define CUR_PROC "ghmm_dmodel_baum_welch"
  int n, k, valid;
  double log_p, log_p_old, log_p_k, diff;
  local_store_t *r = NULL;
  int res = -1;

  /* local store for all iterations */
  r = reestimate_alloc (mo);
  if (!r) {
    GHMM_LOG_QUEUED(LCONVERTED);
    goto STOP;
  };

  log_p_old = -DBL_MAX;
  n = 1;

  /* main loop Baum-Welch-Alg. */
  while (n <= max_step) {
    
    if (1) {
      if (reestimate_one_step(mo, r, sq->seq_number, sq->seq_len, sq->seq,
			      &log_p, sq->seq_w) == -1) {
	GHMM_LOG_PRINTF(LCONVERTED, LOC, "reestimate_one_step false (%d.step)\n", n);
	goto STOP;
      }
    }
    else {
      if (reestimate_one_step_lean(mo, r, sq->seq_number, sq->seq_len, sq->seq,
			      &log_p, sq->seq_w) == -1) {
	GHMM_LOG_PRINTF(LCONVERTED, LOC, "reestimate_one_step false (%d.step)\n", n);
	goto STOP;
      }
    }

    /*if (n == 1)*/
    /*printf("%8.5f (-log_p input model)\n", -log_p); */
    /*else*/
    /*printf("%8.5f (-log_p)\n", -log_p); */

    if (log_p == +1) {
      GHMM_LOG(LERROR, "Reestimate stopped: No sequence can be built "
               "from model mo!");
      break;
    }

    diff = log_p - log_p_old;
    /* error in convergence ? */
    if (diff < -GHMM_EPS_PREC) {
      GHMM_LOG_PRINTF(LCONVERTED, LOC, "No convergence: log P < log P-old! (n=%d)\n", n);
      goto STOP;
    }
    else if (log_p > GHMM_EPS_PREC) {
      GHMM_LOG_PRINTF(LCONVERTED, LOC, "No convergence: log P > 0! (n=%d)\n", n);
      goto STOP;
    }

    /* stop iterations? */
    if (diff < fabs ((double) likelihood_delta * log_p)) {
      GHMM_LOG_PRINTF(LINFO, LOC, "Convergence after %d steps", n);
      break;
    }
    else {
      /* for next iteration */
      log_p_old = log_p;
      reestimate_init (r, mo);  /* sets all fields to zero */
      n++;
    }
  }                             /* while (n <= MAX_ITER) */

  /* log_p of reestimated model */
  log_p = 0.0;
  valid = 0;
  for (k = 0; k < sq->seq_number; k++) {
    if (ghmm_dmodel_logp (mo, sq->seq[k], sq->seq_len[k], &log_p_k) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto STOP;
    }
    if (log_p_k != +1) {
      log_p += log_p_k;
      valid = 1;
    }
  }
  if (!valid)
    log_p = +1;
  /*printf("%8.5f (-log_p optimized model)\n", -log_p);*/

  /* check new parameter for plausibility */
  /* if (ghmm_dmodel_check(mo) < 0) { GHMM_LOG_QUEUED(LCONVERTED); goto STOP; } */
  res = 0;

STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  reestimate_free (&r, mo->N);
  return res;
# undef CUR_PROC
}                               /* ghmm_dmodel_baum_welch_nstep */





/**============================ Labeled HMMs ================================**/
/*----------------------------------------------------------------------------*/
static int reestimate_one_step_label(ghmm_dmodel *mo, local_store_t *r,
                                     int seq_number, int *seq_length, int **O,
                                     int **label, double *log_p, double *seq_w)
{
# define CUR_PROC "reestimate_one_step_label"

  int k, i, j, t, j_id;
  int errors;
  int e_index, T_k, valid=0;
  double **alpha = NULL;
  double **beta = NULL;
  double *scale = NULL;
  double gamma;
  double log_p_k;

  /* first set maxorder to zero if model_type & kHigherOrderEmissions is FALSE 

     TODO XXX use model->maxorder only
     if model_type & GHMM_kHigherOrderEmissions is TRUE */

  if (!(mo->model_type & GHMM_kHigherOrderEmissions))
    mo->maxorder = 0;

  *log_p = 0.0;

  /* loop over all sequences */
  for (k = 0; k < seq_number; k++) {
    mo->emission_history = 0;
    T_k = seq_length[k];        /* current seq. length */

    /* initialization of  matrices and vector depends on T_k */
    if (ighmm_reestimate_alloc_matvek (&alpha, &beta, &scale, T_k, mo->N) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto FREE;
    }

    if (ghmm_dmodel_label_forward (mo, O[k], label[k], T_k, alpha, scale, &log_p_k)
        == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto FREE;
    }

    if (log_p_k != +1) {        /* O[k] can be generated */
      *log_p += log_p_k;
      valid = 1;

      if (ghmm_dmodel_label_backward (mo, O[k], label[k], T_k, beta, scale, &log_p_k)
          == -1) {
        GHMM_LOG_QUEUED(LCONVERTED);
        goto FREE;
      }

      /* loop over all states */
      for (i = 0; i < mo->N; i++) {
        /* Pi */
        r->pi_num[i] += seq_w[k] * alpha[0][i] * beta[0][i];
        r->pi_denom += seq_w[k] * alpha[0][i] * beta[0][i];

        for (t = 0; t < T_k - 1; t++) {
          /* B */
          if (!(mo->s[i].fix) && (mo->label[i] == label[k][t])) {
            e_index = get_emission_index (mo, i, O[k][t], t);
            if (e_index != -1) {
              gamma = seq_w[k] * alpha[t][i] * beta[t][i];
              r->b_num[i][e_index] += gamma;
              r->b_denom[i][e_index / (mo->M)] += gamma;
            }
          }
          update_emission_history (mo, O[k][t]);

          /* A */
          r->a_denom[i] += seq_w[k] * alpha[t][i] * beta[t][i];
          for (j = 0; j < mo->s[i].out_states; j++) {
            j_id = mo->s[i].out_id[j];
            if (label[k][t + 1] != mo->label[j_id])
              continue;
            e_index = get_emission_index (mo, j_id, O[k][t + 1], t + 1);
            if (e_index != -1)
              r->a_num[i][j] += (seq_w[k] * alpha[t][i] * mo->s[i].out_a[j] *
                                 mo->s[j_id].b[e_index] * beta[t + 1][j_id] *
                                 (1.0 / scale[t + 1]));
          }
        }
        /* B: last iteration for t==T_k-1 */
        t = T_k - 1;
        if (!(mo->s[i].fix) && (mo->label[i] == label[k][t])) {
          e_index = get_emission_index (mo, i, O[k][t], t);
          if (e_index != -1) {
            gamma = seq_w[k] * alpha[t][i] * beta[t][i];
            r->b_num[i][e_index] += gamma;
            r->b_denom[i][e_index / (mo->M)] += gamma;
          }
        }
      }
    }
    else {
      GHMM_LOG_PRINTF(LWARN, LOC, "warning: sequence %d can't be built from model", k);
    }

    ighmm_reestimate_free_matvek (alpha, beta, scale, T_k);
  }                             /* for (k = 0; k < seq_number; k++) */

  if (valid) {
    /* new parameter lambda: set directly in model */
    if (reestimate_setlambda (r, mo) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto STOP;
    }

    if ((errors = ghmm_dmodel_check(mo))) {
      GHMM_LOG_PRINTF(LERROR, LOC, "Reestimated model is invalid, "
                      "model_check found %d errors", -errors);
      goto STOP;
    }
  }
  else {                        /* NO sequence can be built from model mo! */
    *log_p = +1;
  }

  return (0);
FREE:
  ighmm_reestimate_free_matvek (alpha, beta, scale, T_k);
STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  return (-1);
# undef CUR_PROC
}                               /* reestimate_one_step_label */

/*============================================================================*/
int ghmm_dmodel_label_baum_welch (ghmm_dmodel * mo, ghmm_dseq * sq)
{
# define CUR_PROC "ghmm_dl_baum_welch"

  return ghmm_dmodel_label_baum_welch_nstep (mo, sq, GHMM_MAX_ITER_BW, GHMM_EPS_ITER_BW);
# undef CUR_PROC
}                               /* ghmm_dmodel_baum_welch */

/*============================================================================*/
int ghmm_dmodel_label_baum_welch_nstep (ghmm_dmodel * mo, ghmm_dseq * sq,
                                       int max_step, double likelihood_delta)
{
# define CUR_PROC "ghmm_dl_baum_welch_nstep"
  int n, k, valid;
  double log_p, log_p_old, log_p_k, diff;
  local_store_t *r = NULL;
  int res = -1;

  /* local store for all iterations */
  r = reestimate_alloc (mo);
  if (!r) {
    GHMM_LOG_QUEUED(LCONVERTED);
    goto STOP;
  };

  log_p_old = -DBL_MAX;
  n = 1;

  /* main loop Baum-Welch-Alg. */
  while (n <= max_step) {

    if (reestimate_one_step_label
        (mo, r, sq->seq_number, sq->seq_len, sq->seq, sq->state_labels,
         &log_p, sq->seq_w) == -1) {
      GHMM_LOG_PRINTF(LCONVERTED, LOC, "reestimate_one_step_label false (%d.step)\n", n);
      goto STOP;
    }

    if (n == 1)
      GHMM_LOG_PRINTF(LINFO, LOC, "%8.5f (-log_p input model)\n", -log_p);
    else
      GHMM_LOG_PRINTF(LINFO, LOC, "%8.5f (-log_p)\n", -log_p);

    if (log_p == +1) {
      GHMM_LOG_PRINTF(LERROR, LOC, "Reestimate stopped: No sequence can be "
                      "built from model mo!\n");
      break;
    }

    diff = log_p - log_p_old;
    /* error in convergence ? */
    if (diff < -GHMM_EPS_PREC) {
      GHMM_LOG_PRINTF(LCONVERTED, LOC, "No convergence: log P < log P-old! (n = %d)\n", n);
      goto STOP;
    }
    else if (log_p > GHMM_EPS_PREC) {
      GHMM_LOG_PRINTF(LCONVERTED, LOC, "No convergence: log P > 0! (n = %d)\n", n);
      goto STOP;
    }

    /* stop iterations? */
    if (diff < fabs ((double) likelihood_delta * log_p)) {
      GHMM_LOG_PRINTF(LINFO, LOC, "Convergence after %d steps\n", n);
      break;
    }
    else {
      /* for next iteration */
      log_p_old = log_p;
      reestimate_init (r, mo);  /* sets all fields to zero */
      n++;
    }
  }                             /* while (n <= MAX_ITER) */

  /* log_p of reestimated model */
  log_p = 0.0;
  valid = 0;
  for (k = 0; k < sq->seq_number; k++) {
    if (ghmm_dmodel_label_logp (mo, sq->seq[k], sq->state_labels[k], sq->seq_len[k],
                         &log_p_k) == -1) {
      GHMM_LOG_QUEUED(LCONVERTED);
      goto STOP;
    }

    if (log_p_k != +1) {
      log_p += log_p_k;
      valid = 1;
    }
  }

  if (!valid)
    log_p = +1;
  GHMM_LOG_PRINTF(LINFO, LOC, "%8.5f (-log_p optimized model)", -log_p);

  /* check new parameter for plausibility */
  /* if (ghmm_dmodel_check(mo) < 0) { GHMM_LOG_QUEUED(LCONVERTED); goto STOP; } */
  res = 0;

STOP:     /* Label STOP from ARRAY_[CM]ALLOC */
  reestimate_free (&r, mo->N);
  return res;
# undef CUR_PROC
}                               /* ghmm_dl_baum_welch_nstep */
