import java.text.*;

/** This class implements a Hidden Markov Model, as well as
the Baum-Welch Algorithm for training HMMs.
@author Holger Wunsch ([email protected])
*/
public class HMM {
  /** number of states */
  public int numStates;

  /** size of output vocabulary */
  public int sigmaSize;

  /** initial state probabilities */
  public double pi[];

  /** transition probabilities */
  public double a[][];

  /** emission probabilities */
  public double b[][];

  /** initializes an HMM.
  @param numStates number of states
  @param sigmaSize size of output vocabulary
  */

 public HMM(int numStates, int sigmaSize) {
   this.numStates = numStates;
   this.sigmaSize = sigmaSize;

   pi = new double[numStates];
   a = new double[numStates][numStates];
   b = new double[numStates][sigmaSize];
 }

 /** implementation of the Baum-Welch Algorithm for HMMs.
  @param o the training set
  @param steps the number of steps
 */
  public void train(int[] o, int steps) {
     int T = o.length;
     double[][] fwd;
     double[][] bwd;

     double pi1[] = new double[numStates];
     double a1[][] = new double[numStates][numStates];
     double b1[][] = new double[numStates][sigmaSize];

     for (int s = 0; s < steps; s++) {
      /* calculation of Forward- und Backward Variables from the
 current model */
      fwd = forwardProc(o);
      bwd = backwardProc(o);

      /* re-estimation of initial state probabilities */
      for (int i = 0; i < numStates; i++)
    pi1[i] = gamma(i, 0, o, fwd, bwd);

      /* re-estimation of transition probabilities */
      for (int i = 0; i < numStates; i++) {
    for (int j = 0; j < numStates; j++) {
     double num = 0;
     double denom = 0;
     for (int t = 0; t <= T - 1; t++) {
       num += p(t, i, j, o, fwd, bwd);
       denom += gamma(i, t, o, fwd, bwd);
     }
   a1[i][j] = divide(num, denom);
  }
    }

  /* re-estimation of emission probabilities */
  for (int i = 0; i < numStates; i++) {
for (int k = 0; k < sigmaSize; k++) {
  double num = 0;
  double denom = 0;

  for (int t = 0; t <= T - 1; t++) {
    double g = gamma(i, t, o, fwd, bwd);
    num += g * (k == o[t] ? 1 : 0);
    denom += g;
  }
  b1[i][k] = divide(num, denom);
}
  }
  pi = pi1;
  a = a1;
  b = b1;
}
}


  /** calculation of Forward-Variables f(i,t) for state i at time
  t for output sequence O with the current HMM parameters
  @param o the output sequence O
  @return an array f(i,t) over states and times, containing
          the Forward-variables.
  */

  public double[][] forwardProc(int[] o) {
int T = o.length;
double[][] fwd = new double[numStates][T];

/* initialization (time 0) */
for (int i = 0; i < numStates; i++)
  fwd[i][0] = pi[i] * b[i][o[0]];

/* induction */
for (int t = 0; t <= T-2; t++) {
  for (int j = 0; j < numStates; j++) {
fwd[j][t+1] = 0;
for (int i = 0; i < numStates; i++)
  fwd[j][t+1] += (fwd[i][t] * a[i][j]);
fwd[j][t+1] *= b[j][o[t+1]];
  }
}

return fwd;
}