// FiniteTD.H; copyright by Susumu Katayama
// Logarithmic-time implementation of Temporal-difference learning (Katayama & Kobayashi 1999, Katayama 2000) with finite states.
// You may use this file within your personal use, provided that you leave these comment lines untouched.
// You may not redistribute this file without my permission if some parts have been changed (me = S. Katayama).


// This version should be a little faster for some cases than that for TD with infinitely countable states.
// There still exists room for optimization exploiting the fact
// that half of the (d,s)'s are equal to (1,0).
#ifndef FINITETD_H
#define FINITETD_H

#include <vector>
#include <utility>
using namespace std;

//#include <vector.h>
//#include <pair.h>

class FiniteTD {
  friend istream &operator>>(istream &istr, FiniteTD &ftd);
  friend ostream &operator<<(ostream &ostr, const FiniteTD &ftd);

private:
  unsigned size;
  vector<pair<double,double> > ds; // the latter half is z and p.
  double prevV;

public:
  // Well, I guess implementation of cuv and cv should be moved to FiniteTD.C ....
  double cuv(unsigned x, double dd, double ss) {
    pair<double,double> ds10(1,0);
    unsigned y = 1;
    vector<pair<double,double> >::iterator pb = ds.begin(), p=pb+1;
    while (y < size) { // cerr << dd << "\t" << ss << endl;
      ss = ss * p->first + p->second;
      dd *= p->first;
      *p = ds10;
      p += y;
      y <<= 1;
      unsigned xand1 = x & 1;
      if (xand1) {
	p->second += ss * p->first; // cerr << p->second << endl;
	p->first  *= dd;
	p++;
	y++;
      } else {
	p++;
	p->second += ss * p->first; // cerr << p->second << endl;
	p->first  *= dd;
	p--;
      }
      x >>= 1;
    }
    p->second += ss * p->first; // cerr << p->second << endl;
#ifdef ACCUMULATE
    p->first *= dd;
#else
    (*p).first = 1;
#endif
    return (*p).second;
  }
  void ct(double dd, double ss) {
    ds[1].second += ss * ds[1].first;
    ds[1].first *= dd;
  }
  double cv(unsigned x, double ss = 0) const {
    unsigned y = 1;
    vector<pair<double,double> >::const_iterator pb = ds.begin(), p=pb+1;
    ss = ss * p->first + p->second;
    while (y < size) {
      unsigned xand1 = x & 1;
      p += y + xand1;
      ss = ss * p->first + p->second;
      y = (y<<1) | xand1; //  "y = p - pb" would not be so fast.
      x >>= 1;
    }
    return ss;
  }

  FiniteTD() {
  }
  FiniteTD(unsigned sz, double v0) {
    sz--;
    for (size=1; sz>0; sz>>=1)
      size<<=1;
    pair<double,double> zp0(0,v0); //cerr << size;
    ds.resize(size*2);
    vector<pair<double,double> >::iterator p = ds.begin()+size, pe = ds.end();
    for ( ; p!=pe; p++)
      *p = zp0;
    prevV = v0;
  }
  double next(double reward, double gamma, double alpha, double lambda, unsigned obs) {
    prevV = cuv(obs, lambda*gamma, alpha * (reward + gamma*cv(obs) - prevV));
    return prevV;
  }
};

  /*
      unsigned child = y<<1;
      ds[child].second += ss * ds[child].first;
      ds[child].first *= dd;
      */

istream &operator>>(istream &istr, FiniteTD &ftd);
ostream &operator<<(ostream &ostr, const FiniteTD &ftd);

#endif