#include <stdio.h>

#define STATES 3

#define ACTIONS 2

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

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

int policy[ACTIONS];

#define ITERATIONS 20
#define DISCOUNT 0.5

double values[2][STATES];

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

void latexshowvalues(int ite) {
  int i;
  printf("%i",ite);
  for(i=0;i<STATES;i++) {
    printf(" & %.2f",values[ite][i]);
  }
  printf(" \\\\ \n");
  fflush(stdout);
}

/* Evaluate the policy iteratively: this is like value iteration except
   that no maximization over actions is made. This iterative computation
   (instead of exact solution of the linear equations representing
   the value of the policy) means that the main advantage of policy
   iteration, exactness, is lost. (We just wanted to demonstrate how
   policy iteration works, and were too lazy to start impementing
   algorithms for solving sets of linear equations.) */

void evaluate() {
  int i,j,k,iteration;
  for(i=0;i<STATES;i++) {
    values[0][i] = 0.0;
  }
  for(iteration=1;iteration<=ITERATIONS;iteration++) {
    for(i=0;i<STATES;i++) {
      double val = reward[policy[i]][i];
	for(k=0;k<STATES;k++) {
	  double p = actions[policy[i]][i][k];
	  val += DISCOUNT*p*values[(iteration-1)&1][k];
	}
	values[iteration&1][i] = val;
	/* Here and above we use a trick: only 2 vectors are needed
	   for representing the values, each of which is in turn
	   for the last (or the second) last value vector. */
    }
  }
}

int main() {
  int iteration;
  int i,j,k;
  int change;
  for(i=0;i<ACTIONS;i++) {
    policy[i] = 1;
  }
  do {
    evaluate(); /* Evaluate the current policy. */
    showvalues(0);
    change = 0;
    for(i=0;i<STATES;i++) { /* Choose the best action for every state. */
      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[0][k];
	}
	if(val > bestval) {
	  bestval = val;
	  bestaction = j;
	}
      }
      if(policy[i] != bestaction) change = 1; /* Policy got improved. */
      policy[i] = bestaction;
      printf("state %i has action %i.\n",i,bestaction);
    }
  } while(change);
}