package pl.edu.pw.mini.sg.mixed.expander

import pl.edu.pw.mini.sg.rng.RandomGenerator
import pl.edu.pw.mini.sg.game.deterministic.TreeExpander
import pl.edu.pw.mini.sg.game.deterministic.single.{GameWithPayoff, SinglePlayerGame}
import pl.edu.pw.mini.sg.mcts.{MctsAlgorithm, MctsParameters}
import pl.edu.pw.mini.sg.mcts.tree.MctsNode



final case class MctsExpander(
  mctsParameters: MctsParameters,
  evaluationFunction: GWPEvaluation
) extends TreeExpander[Any, Any, GameWithPayoff] {

  override type SubtreeData[M,S] = Option[MctsNode[M,S]]

  override def expand[MM, SS, GG <: SinglePlayerGame[MM,SS,GG] with GameWithPayoff](
    state: GG, subtreeData: SubtreeData[MM,SS]
  )(implicit rng: RandomGenerator) : (MM, SubtreeData[MM,SS]) = {
    require(state.possibleMoves.nonEmpty)
    val tree: MctsNode[MM, SS] = subtreeData.getOrElse({
      mctsParameters.treeFactory.create[MM,SS](state.playerState, state.possibleMoves).root
    })
    MctsAlgorithm.training(tree, state, (x: GG) => (evaluationFunction.apply(x), None), mctsParameters, rng)
    val bestMove = tree.estimates.maxBy(_.estimate).move
    (bestMove, tree.successor(bestMove, state.makeMove(bestMove).playerState))
  }

  override def zeroSubtreeData[MM, SS]: SubtreeData[MM,SS] = None
}
