#include "AI_Tree.h"
#include <math.h>

using namespace std;

MQV_Map& AI_Tree::get_qvalues()
{
	AI_SearchHeuristic* sh = _player->_search_heuristic;
	AI_StateInspector* si = _player->_state_inspector;
	sh->begin_traversal();

	if (_root == 0)	_root = new AI_Tree_Node(0, sh->branch_var());

	AI_Tree_Node* cn = _root;

	while (!cn->is_leaf())
		cn = next(cn, sh, si);

	++cn->num_visits;
	return cn->mqv_map;
}

void AI_Tree::update_qvalue(const AI_State& state, MOVE_ID move, double nq)
{
	assert(_root != 0);
	AI_Tree_Node* cn = _root;

	int index = 0;
	while (!cn->is_leaf())
		cn = cn->children[state_value_to_index(cn->branch_var, state[index++])];
	
	double lr = _player->_lr;
	double oq = cn->mqv_map[move];
	
	cn->mqv_map[move] = oq*(1-lr) + nq*lr;
}

void AI_Tree::load(PLAYER_ID pid)
{
}

void AI_Tree::save(PLAYER_ID pid)
{
}

void AI_Tree::clear()
{
	if (_root!=0)
		rec_delete(_root);

	_root = 0;
}

void AI_Tree::rec_delete(AI_Tree_Node* cn)
{
	tn_iter it = cn->children.begin();
	tn_iter end = cn->children.end();

	while (it != end)
	{
		rec_delete((*it).second);
		++it;
	}

	delete cn;
}

AI_Tree::AI_Tree_Node* AI_Tree::next(AI_Tree_Node* cn, AI_SearchHeuristic* sh, AI_StateInspector* si)
{
	assert(cn->branch_var != NULL_VAR);
	++cn->num_visits;

	STATE_VAR_ID sid = cn->branch_var;
	STATE_VALUE sv = si->get_state_var(sid);

	int index = state_value_to_index(sid, sv);
	sh->next(sv);

	tn_iter it;
	if ((it = cn->children.find(index)) == cn->children.end())
	{
		AI_Tree_Node* sn = new AI_Tree_Node(cn, sh->branch_var());
		cn->children[index] = sn;
		cn = sn;
		if (cn->is_leaf())
			sh->initial_heuristic(cn->mqv_map);
	} else {
		cn = (*it).second;
	}

	return cn;
}