package pl.edu.pw.mini.jk.sg.mixeduct;

import java.util.Map;
import java.util.Optional;

import org.apache.commons.math3.random.RandomGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import pl.edu.pw.mini.jk.lib.Logging;
import pl.edu.pw.mini.jk.lib.PmfUtils;
import pl.edu.pw.mini.jk.lib.PmfUtil2;
import pl.edu.pw.mini.jk.sg.game.DefenderAttackerGame;
import pl.edu.pw.mini.jk.sg.game.common.Cloneable;
import pl.edu.pw.mini.jk.sg.game.common.DefenderMove;
import pl.edu.pw.mini.jk.sg.game.common.SecurityGame;
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableAttackerState;
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableDefenderState;
import pl.edu.pw.mini.jk.sg.gteval.GameMatrices;
import pl.edu.pw.mini.jk.sg.mixeduct.attacker.AttackerStrategy;
import pl.edu.pw.mini.jk.sg.mixeduct.attacker.AttackerWithHistory;
import pl.edu.pw.mini.jk.sg.mixeduct.attackerresponse.AttackerSolver;
import pl.edu.pw.mini.jk.sg.mixeduct.attackerresponse.AttackerSolverResult;
import pl.edu.pw.mini.jk.sg.mixeduct.resulttrends.ResultTrends;
import pl.edu.pw.mini.jk.sg.mixeduct.startattackers.AttackerProbabilityBuilder;
import pl.edu.pw.mini.jk.sg.mixeduct.stopcondition.StopCondition;
import pl.edu.pw.mini.jk.sg.mixeduct.strategychooser.StrategyChooser;
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.AttackersInformationTransfer;
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.strategyfilter.StrategyFilter;
import pl.edu.pw.mini.jk.sg.mixeduct.util.TreePrinter;
import pl.edu.pw.mini.jk.sg.uct.UctPlayer;
import pl.edu.pw.mini.jk.sg.uct.tree.UctNode;

public class MixedUctStrategy {

    public static <DM extends DefenderMove, AM, AS extends AdvanceableAttackerState<AM>, DS extends AdvanceableDefenderState<DM>, G extends Cloneable<G> & SecurityGame & DefenderAttackerGame<AM, DM, AS, DS, ?>> MixedUctTrainingResultFactory<DM, AM> uctTrainedPlayerStrategy(
			final MixedUctConfig muConfig, final GameMatrices<AM, DM, AS, DS, G> gameMatrices,
			final RandomGenerator rng, final TreeLogConfig treeLogConfig) {
	int iterationCount = 0;
	MixedUctTrainingResultFactory<DM, AM> bestResult = null;

	do {
	    final MixedUctTrainingResultFactory<DM, AM> res = uctTrainedPlayerStrategySingleRun(muConfig, gameMatrices, rng, iterationCount, treeLogConfig);

	    if(bestResult == null)
		bestResult = res;
	    else
		bestResult = MixedUctTrainingResultFactory.joinRestarts(bestResult, res);
	    iterationCount = bestResult.getIterations();
	} while (iterationCount < muConfig.iterationBudget() && bestResult.getRestarts() < muConfig.restartLimit());

	return bestResult;
    }
    
	public static <DM extends DefenderMove, AM, AS extends AdvanceableAttackerState<AM>, DS extends AdvanceableDefenderState<DM>, G extends Cloneable<G> & SecurityGame & DefenderAttackerGame<AM, DM, AS, DS, ?>> MixedUctTrainingResultFactory<DM, AM> uctTrainedPlayerStrategySingleRun(
			final MixedUctConfig muConfig, final GameMatrices<AM, DM, AS, DS, G> gameMatrices,
			final RandomGenerator rng, final int startIteration, final TreeLogConfig treeLogConfig) {

		final Logger defenderResultLogger = LoggerFactory.getLogger("petrus.DefenderResult");
		final Logger defenderStrategyLogger = LoggerFactory.getLogger("petrus.DefenderStrategy");
		final Logger attStrategyLogger = LoggerFactory.getLogger("petrus.AttackerStrategy");
		final Logger avAttStrategyLogger = LoggerFactory.getLogger("petrus.AveragedAttackerStrategy");
		final Logger treeLogger = LoggerFactory.getLogger("petrus.UctTree");
		final Optional<Integer> loggingPrecission = Optional.of(4);

		final AttackersInformationTransfer<DM, DS> attackersInformationTransfer = muConfig.getInformationTransfer()
		    .create(gameMatrices.getGameLength(), gameMatrices, muConfig.i2UctConfig().uctTreeFactory().treeJoiner());
		final StrategyFilter<DM> strategyFilter = muConfig.getStrategyFilter().getStrategyFilter();
		final StrategyChooser<DM, AM> strategyChooser = muConfig.getStrategyChooser().create(gameMatrices);
		final AttackerProbabilityBuilder attackerProbabilityBuilder = muConfig.getStartAttacker();
		final ResultTrends resultTrends = muConfig.getResultTrends().createResultTrends();
		final AttackerSolver attackerCalculator = muConfig.getAttackerSolver();
		
		UctNode<DS, DM> uctTree = muConfig.i2UctConfig().uctTreeFactory().create(gameMatrices.getDefenderInitialState());
		final StopCondition<DM,AM> stopCondition = muConfig.getStopConditionFactory().create();
		final AttackerWithHistory<AM> attackerWithHistory = muConfig.getAttackerHistory().getAttackerWithHistory(rng, attackerProbabilityBuilder.buildProbabilities(gameMatrices));

		
		int i;
		for (i = startIteration; !stopCondition.isStopCondition() && i < muConfig.iterationBudget(); i++) {
			final int iteration = i;
			final AttackerStrategy<AM> attackerProbabilities = attackerWithHistory.getInstance();

			Logging.logIfDebug(avAttStrategyLogger, () -> String.format("I: %d AAS: %s", iteration,
					PmfUtils.pmfToSortedString(attackerProbabilities.getPmf(), loggingPrecission)));

			final UctPlayer<DM, DS, DefenderSecurityGameView<DM, AM, DS, G>> uctPlayer = UctPlayer
		    .fromConfig(
				DefenderSecurityGameView.transformFactory(gameMatrices.getNoaiFactory(),									  attackerProbabilities),
							gameMatrices.getGameLength(), muConfig.getI2UctConfig(), rng);

			for(int j =0; j< gameMatrices.getGameLength(); j++) {
			    uctPlayer.doLearningSequence(j, uctTree);
			    if(muConfig.interStepPruneFunction().isDefined()) {
				uctTree = uctTree.pruneTree(muConfig.interStepPruneFunction().get());
			    }
			}
			uctTree = attackersInformationTransfer.newAttackerTree(uctTree, rng);
			
			final IterationResult<DM, AM, AS, DS, G> iterationResult;
			{
			final Map<ImmutableList<DM>, Double> defenderMoveProbabilitiesNF = attackersInformationTransfer
			    .getDefenderMoveProbabilities(rng, treeLogConfig.log(i));
			Logging.logIfDebug(defenderStrategyLogger, () -> String.format("I: %d NF: %s", iteration,
					PmfUtils.pmfToSortedString(defenderMoveProbabilitiesNF, loggingPrecission)));

			final ImmutableMap<ImmutableList<DM>, Double> defenderMoveProbabilities = PmfUtil2.normalizePmf(strategyFilter
															.filterStrategy(defenderMoveProbabilitiesNF));
			Logging.logIfDebug(defenderStrategyLogger, () -> String.format("I: %d F: %s", iteration,
					PmfUtils.pmfToSortedString(defenderMoveProbabilities, loggingPrecission)));
			
			
			final AttackerSolverResult<AM> attackerStrategy = attackerCalculator.calculateAttackerStrategy(defenderMoveProbabilities, gameMatrices,
					rng);
			iterationResult = new IterationResult<DM, AM, AS, DS, G>(defenderMoveProbabilities, attackerStrategy, gameMatrices);
			}


			if (attStrategyLogger.isDebugEnabled()) {
				attStrategyLogger.debug("I: {} AS: {}", i,
							PmfUtils.pmfToSortedString(iterationResult.attackerStrategy().strategy(), loggingPrecission));
			}
			
			if (defenderResultLogger.isDebugEnabled()) {
				defenderResultLogger.debug("DefScore: {} {} {}", i,
							   iterationResult.givenPayoff().defenderPayoff(),
							   iterationResult.givenPayoff().attackerPayoff());
			}
			strategyChooser.registerStrategy(iterationResult);
			resultTrends.registerResult(iterationResult.givenPayoff().defenderPayoff());
			stopCondition.registerTrainingResult(iterationResult);
			attackerWithHistory.registerIterationResult(i, iterationResult, attackerProbabilities);
		}

		return new MixedUctTrainingResultFactory<DM, AM>(i, 1, strategyChooser.getBestIterationResult());
	}
}
