#include <stdio.h>
#include <math.h>

#define STATES 5
#define ACTIONS 2

double actions[ACTIONS][STATES][STATES] = {
  {
    { 0.0, 0.0, 1.0, 0.0, 0.0 },
    { 0.1, 0.0, 0.0, 0.9, 0.0 },
    { 1.0, 0.0, 0.0, 0.0, 0.0 },
    { 0.0, 0.0, 0.0, 0.0, 1.0 },
    { 1.0, 0.0, 0.0, 0.0, 0.0 }
  },
  {
    { 0.0, 1.0, 0.0, 0.0, 0.0 },
    { 1.0, 0.0, 0.0, 0.0, 0.0 },
    { 0.0, 0.0, 0.0, 0.0, 1.0 },
    { 0.0, 0.0, 1.0, 0.0, 0.0 },
    { 0.0, 0.0, 1.0, 0.0, 0.0 }
  }
};

double reward[ACTIONS][STATES] = {
  { 1.0, 0.0, 0.0, 5.0, 0.0 },
  { 0.0, 0.0, 0.0, 0.0, 0.0 }
};

#define DISCOUNT 0.6
#define EPSILON 0.01

/* Store the value function for MAXITERATIONS iterations. */
#define MAXITERATIONS 10000
double values[MAXITERATIONS][STATES];

void showvalues(int ite) {
  int i;
  printf("%i:",ite);
  for(i=0;i<STATES;i++) {
    printf(" %.4f",values[ite][i]);
  }
  printf("\n");
}

int main() {
  int i,j,k,iteration;
  double delta;
  for(i=0;i<STATES;i++) {
    values[0][i] = 0.0;
  }
  iteration=0;
  showvalues(iteration);
  do {
    delta = 0.0; /* Biggest change in value function. */
    iteration += 1; 
    for(i=0;i<STATES;i++) {
      double maxval = -1E10; /* Highest reward obtainable in state i */
      int action;
      for(j=0;j<ACTIONS;j++) {
        double val = reward[j][i];
        for(k=0;k<STATES;k++) {
          double p = actions[j][i][k];
          val += DISCOUNT*p*values[iteration-1][k];
        }
        if(val > maxval) {
          action = j;
          maxval = val;
        }
      }
      //      printf("state %i action %i.\n",i,action);
      values[iteration][i] = maxval;

      if(fabs(values[iteration][i]-values[iteration-1][i]) > delta)
	delta =  fabs(values[iteration][i]-values[iteration-1][i]);
    }
    showvalues(iteration);
  } while (delta > EPSILON*(1-DISCOUNT)/2*DISCOUNT);

  /* See which action produce the policy. */
  for(i=0;i<STATES;i++) {
    int action;
    double bestval = -1E20;
    int bestaction = -1;
    for(j=0;j<ACTIONS;j++) {
      double val = reward[j][i];
      for(k=0;k<STATES;k++) {
        double p = actions[j][i][k];
        val += DISCOUNT*p*values[iteration][k];
      }
      if(fabs(val-values[iteration][i]) < fabs(bestval-values[iteration][i])) {
        bestval = val;
        bestaction = j;
      }
    }
    printf("state %i has action %i.\n",i,bestaction);
  }
  exit(0);
}