package pl.edu.pw.mini.jk.sg.uct

import org.apache.commons.collections4.Factory
import org.apache.commons.math3.random.RandomGenerator
import pl.edu.pw.mini.jk.sg.game.common.SecurityGame
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.statsupdater.{StatsUpdater, StatsUpdaterFactory}
import pl.edu.pw.mini.jk.sg.uct.stopcondition.{ CombinedStopCondition, GameLengthStopCondition, GameStopCondition }
import pl.edu.pw.mini.jk.sg.uct.tree.UctNode
import scala.annotation.tailrec

//TODO: przenieść fabrykę do parametrów do learning sequence.
//Następnie: Zrobić UCTPlayer typem kontrawariantnym
final class UctPlayer[M, S <: AdvanceableState[M], G <: SinglePlayerGame[M, _ <: S]](
  gameFactory: Factory[_ <: G],
  calculateBoundFunction: SelectionFunction,
  statsUpdaterFactory: StatsUpdaterFactory[_ >: S, _ >: M],
  gameStopCondition: GameStopCondition[_ >: G],
  simulations: Int,
  random: RandomGenerator
) {
  def doLearningSequence(timestepDescent: Int, uctTreeRoot: UctNode[S, M]): Unit = {
    for (i <- 0 to simulations) {
      val game2 = gameFactory.create();
      var node = uctTreeRoot;
      while(game2.getTimeStep < timestepDescent) {
        val move = node.getBestMove(random);
        game2.makeMove(move);
        node.getSuccessor(move) match {
          case Some(n) => node = n;
          case None => node = node.addSuccessor(move, game2.getCurrentState(),
            game2.getCurrentState().getPossibleMoves());
        }
      }
      playSingleGame(game2, node);
    }
  }

  private def playSingleGame(game: G, startingNode: UctNode[S, M]): Unit = {
    require(game.getCurrentState.equals(startingNode.getState))

    val statsUpdater: StatsUpdater[S, M] = statsUpdaterFactory.create();

    @tailrec
    def selection(currentNode: UctNode[S, M]): Unit = {
      if (gameStopCondition.isGameStopCondition(game)) {
        Unit
      } else {
        require(game.getCurrentState.equals(currentNode.getState))

        val selectedMove = currentNode.selectOrRandomUnvisidedMove(calculateBoundFunction, random)
        val payoff = game.makeMove(selectedMove).payoff
        statsUpdater.registerDecision(currentNode, selectedMove)
        if (payoff != 0) { statsUpdater.updatePayoffs(payoff, game.getCurrentState) }

        currentNode.getSuccessor(selectedMove) match {
          case Some(n) => selection(n)
          case None => expansion(currentNode, selectedMove)
        }
      }
    }

    def expansion(expandedNode: UctNode[S, M], move: M): Unit = {
      expandedNode.addSuccessor(move, game.getCurrentState, game.possibleMoves)
      mc()
    }

    @tailrec
    def mc(): Unit = {
      if (gameStopCondition.isGameStopCondition(game))
        Unit
      else {
        val mcMoveChoice = game.randomMove(random);
        statsUpdater.registerMcMove(game.getCurrentState, mcMoveChoice);
        val payoff = game.makeMove(mcMoveChoice).payoff;
        if (payoff != 0) {
          statsUpdater.updatePayoffs(payoff, game.getCurrentState);
        }
        mc()
      }
    }

    selection(startingNode)
  }
}


object UctPlayer {
  def combinedStopConditionWithRoundLimit[G <: SecurityGame](conditions: List[GameStopCondition[G]], gameDepth: Int): GameStopCondition[G] = {
    return new CombinedStopCondition(
      (new GameLengthStopCondition(gameDepth.toDouble, i => i.toDouble)) :: conditions
    )
  }

  def fromConfig[M, S <: AdvanceableState[M], G <: SecurityGame with SinglePlayerGame[M, S]](
    gameFactory: Factory[G],
    gameDepth: Int,
    i2UctConfig: I2UctConfig[_ >: G],
    rng: RandomGenerator
  ) = new UctPlayer[M, S, G](
    gameFactory,
    i2UctConfig.selectionFormula,
    i2UctConfig.statsUpdaterFactory,
    combinedStopConditionWithRoundLimit(i2UctConfig.stopConditions, gameDepth),
    i2UctConfig.getSimulationsInLevel, rng
  )
}
