package net.woodyfolsom.msproj.policy;

import java.util.ArrayList;
import java.util.List;

import net.woodyfolsom.msproj.Action;
import net.woodyfolsom.msproj.GameState;
import net.woodyfolsom.msproj.Player;
import net.woodyfolsom.msproj.tree.AMAFProperties;
import net.woodyfolsom.msproj.tree.GameTreeNode;
import net.woodyfolsom.msproj.tree.MonteCarloProperties;

public class MonteCarloAMAF extends MonteCarloUCT {

	public MonteCarloAMAF(Policy movePolicy, long searchTimeLimit) {
		super(movePolicy, searchTimeLimit);
	}

	@Override
	public void update(GameTreeNode<MonteCarloProperties> node, Rollout rollout) {
		GameTreeNode<MonteCarloProperties> currentNode = node;
		//List<Action> subTreeActions = new ArrayList<Action>(rollout.getPlayout());
		
		List<Action> playout = rollout.getPlayout();
		int reward = rollout.getReward();
		while (currentNode != null) {
			AMAFProperties nodeProperties = (AMAFProperties)currentNode.getProperties();
			
			//Always update props for the current node
			nodeProperties.setWins(nodeProperties.getWins() + reward);
			nodeProperties.setVisits(nodeProperties.getVisits() + 1);
			nodeProperties.setAmafWins(nodeProperties.getAmafWins() + reward);
			nodeProperties.setAmafVisits(nodeProperties.getAmafVisits() + 1);
			
			GameTreeNode<MonteCarloProperties> parentNode = currentNode.getParent();
			if (parentNode != null) {
				Player playerToMove = parentNode.getGameState().getPlayerToMove();
				for (Action actionFromParent : parentNode.getActions()) {
					if (playout.contains(actionFromParent)) {
						GameTreeNode<MonteCarloProperties> subTreeChild = parentNode.getChild(actionFromParent);
						//Don't count AMAF properties for the current node twice
						if (subTreeChild == currentNode) {
							continue;
						}
						
						AMAFProperties siblingProperties = (AMAFProperties)subTreeChild.getProperties();
						//Only update AMAF properties if the sibling is reached by the same action with the same player to move
						if (rollout.hasPlay(playerToMove,actionFromParent)) {
							siblingProperties.setAmafWins(siblingProperties.getAmafWins() + reward);
							siblingProperties.setAmafVisits(siblingProperties.getAmafVisits() + 1);
						}
					}
				}
			}

			currentNode = currentNode.getParent();
		}
	}

	@Override
	protected GameTreeNode<MonteCarloProperties> createRootNode(GameState rootGameState) {
		return new GameTreeNode<MonteCarloProperties>(
				rootGameState, new AMAFProperties());
	}
	
	@Override
	protected double getNodeScore(GameTreeNode<MonteCarloProperties> gameTreeNode) {
		//double nodeVisits = gameTreeNode.getParent().getProperties().getVisits();
		double parentAmafVisits = ((AMAFProperties)gameTreeNode.getParent().getProperties()).getAmafVisits();
		
		double nodeScore;
		if (gameTreeNode.getGameState().isTerminal()) {
			nodeScore = 0.0;
		} else {
			/*
			MonteCarloProperties properties = gameTreeNode.getProperties();
			nodeScore = (double) (properties.getWins() / properties
					.getVisits())
					+ (TUNING_CONSTANT * Math.sqrt(Math.log(nodeVisits)
							/ gameTreeNode.getProperties().getVisits()));
							 * 
							 */
			AMAFProperties properties = (AMAFProperties) gameTreeNode.getProperties();
			nodeScore = (double) (properties.getAmafWins() / properties
					.getAmafVisits())
					+ (TUNING_CONSTANT * Math.sqrt(Math.log(parentAmafVisits)
							/ properties.getAmafVisits()));
		
		}
		return nodeScore;
	}
	
	@Override
	protected List<GameTreeNode<MonteCarloProperties>> addNewChildren(
			GameTreeNode<MonteCarloProperties> node, Action action,
			GameState successorState) {
		List<GameTreeNode<MonteCarloProperties>> newChildren = new ArrayList<GameTreeNode<MonteCarloProperties>>();

		GameTreeNode<MonteCarloProperties> newChild = new GameTreeNode<MonteCarloProperties>(
				successorState, new AMAFProperties());

		newChildren.add(newChild);
		node.addChild(action, newChild);
		return newChildren;
	}
}