package policies.softmax;

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

import features.feature_sets.BaseFeatureSet;
import function_approx.LinearFunction;
import game.Game;
import game.types.play.RoleType;
import main.collections.FVector;
import main.collections.FastArrayList;
import metadata.ai.features.Features;
import metadata.ai.features.trees.FeatureTrees;
import other.context.Context;
import other.move.Move;
import other.trial.Trial;
import search.mcts.MCTS;

/**
 * A Softmax Policy that can automatically initialise itself by
 * using the Playout features embedded in a game's metadata.
 * 
 * @author Dennis Soemers
 */
public final class SoftmaxFromMetadataPlayout extends SoftmaxPolicy
{
	
	//-------------------------------------------------------------------------
	
	/** Softmax policy we wrap around; can change into a linear or decision tree based policy depending on metadata */
	private SoftmaxPolicy wrappedSoftmax = null;
	
	//-------------------------------------------------------------------------
	
	/**
	 * Constructor
	 * @param epsilon Epsilon for epsilon-greedy feature-based playouts. 1 for uniform, 0 for always softmax
	 */
	public SoftmaxFromMetadataPlayout(final double epsilon)
	{
		friendlyName = "Softmax Policy (Playout features from Game metadata)";
		this.epsilon = epsilon;
	}
	
	//-------------------------------------------------------------------------

	@Override
	public void initAI(final Game game, final int playerID)
	{
		final Features featuresMetadata = game.metadata().ai().features();
		if (featuresMetadata != null)
		{
			final List<BaseFeatureSet> featureSetsList = new ArrayList<BaseFeatureSet>();
			final List<LinearFunction> linFuncs = new ArrayList<LinearFunction>();
			
			wrappedSoftmax = new SoftmaxPolicyLinear();
			wrappedSoftmax.epsilon = epsilon;
			wrappedSoftmax.playoutActionLimit = 200;
			
			for (final metadata.ai.features.FeatureSet featureSet : featuresMetadata.featureSets())
			{
				if (featureSet.role() == RoleType.Shared)
					((SoftmaxPolicyLinear)wrappedSoftmax).addFeatureSetWeights(0, featureSet.featureStrings(), featureSet.playoutWeights(), featureSetsList, linFuncs);
				else
					((SoftmaxPolicyLinear)wrappedSoftmax).addFeatureSetWeights(featureSet.role().owner(), featureSet.featureStrings(), featureSet.playoutWeights(), featureSetsList, linFuncs);
			}
			
			((SoftmaxPolicyLinear) wrappedSoftmax).featureSets = featureSetsList.toArray(new BaseFeatureSet[featureSetsList.size()]);
			((SoftmaxPolicyLinear) wrappedSoftmax).linearFunctions = linFuncs.toArray(new LinearFunction[linFuncs.size()]);
		}
		else
		{
			// TODO no distinction between selection and playout here
			final FeatureTrees featureTrees = game.metadata().ai().trainedFeatureTrees();
			wrappedSoftmax = SoftmaxPolicyLogitTree.constructPolicy(featureTrees, epsilon);
			wrappedSoftmax.playoutActionLimit = 200;
		}
		
		wrappedSoftmax.initAI(game, playerID);
		
		super.initAI(game, playerID);
	}
	
	@Override
	public boolean supportsGame(final Game game)
	{
		// We support any game with appropriate features in metadata
		if (game.metadata().ai() != null)
		{
			if (game.metadata().ai().features() != null)
			{
				final Features featuresMetadata = game.metadata().ai().features();
				if (featuresMetadata.featureSets().length == 1 && featuresMetadata.featureSets()[0].role() == RoleType.Shared)
					return true;
				else
					return (featuresMetadata.featureSets().length == game.players().count());
			}
			else if (game.metadata().ai().trainedFeatureTrees() != null)
			{
				return true;
			}
			else
			{
				return false;
			}
		}
		
		return false;
	}
	
	//-------------------------------------------------------------------------
	
	@Override
	public Trial runPlayout(final MCTS mcts, final Context context) 
	{
		return wrappedSoftmax.runPlayout(mcts, context);
	}

	@Override
	public boolean playoutSupportsGame(final Game game) 
	{
		return supportsGame(game);
	}

	@Override
	public int backpropFlags() 
	{
		return 0;
	}

	@Override
	public void customise(final String[] inputs) 
	{
		System.err.println("customise() not implemented for SoftmaxFromMetadataPlayout!");
	}

	@Override
	public float computeLogit(final Context context, final Move move) 
	{
		return wrappedSoftmax.computeLogit(context, move);
	}

	@Override
	public FVector computeDistribution(final Context context, final FastArrayList<Move> actions, final boolean thresholded) 
	{
		return wrappedSoftmax.computeDistribution(context, actions, thresholded);
	}

	@Override
	public Move selectAction(final Game game, final Context context, final double maxSeconds, final int maxIterations, final int maxDepth) 
	{
		return wrappedSoftmax.selectAction(game, context, maxSeconds, maxIterations, maxDepth);
	}
}
