package pl.edu.pw.mini.jk.sg.mixeduct.transfer
import pl.edu.pw.mini.jk.sg.game.common.DefenderMove
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableDefenderState
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.probabilityfunctions.NodeProbabilityWeightFunction
import pl.edu.pw.mini.jk.sg.uct.tree.TreeJoiner
import pl.edu.pw.mini.jk.sg.uct.tree.TreeJoinerFunctions
import java.{lang => jl, util => ju}
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.filters.BestPartFilterFactory
import pl.edu.pw.mini.jk.sg.uct.tree.UctNode
import org.apache.commons.math3.random.RandomGenerator
import com.google.common.collect.ImmutableList
import pl.edu.pw.mini.jk.sg.mixeduct.util.TreeUtils
import pl.edu.pw.mini.jk.sg.mixeduct.TreeLogConfig
import pl.edu.pw.mini.jk.sg.uct.tree.PayoffTransformer

final class BestPartTransfer[DM <: DefenderMove, DS <: AdvanceableDefenderState[DM]](
  roundLimit: Int,
  portion: Double,
  nodeProbabilityWeightFunction: NodeProbabilityWeightFunction,
  treeDescentLimitCoefficient: Double,
  historySize: Int,
  treeJoiner: TreeJoiner,
  treeJoinerFunctions: TreeJoinerFunctions,
  treeQualityFunction: UctNode[DS, DM] => Double,
  bestPartFilterFactory: BestPartFilterFactory,
  preAddTransformer: PayoffTransformer
) extends AttackersInformationTransfer[DM, DS] {
  val pruneFunction = new BestPartPruneFunctions(portion)
  val bestPartFilter = bestPartFilterFactory.createInstance(treeQualityFunction)

  override def newAttackerTree(oldTree: UctNode[DS,DM], rng: RandomGenerator): UctNode[DS,DM] = {
    bestPartFilter.registerTree(oldTree.pruneTree(pruneFunction).transformStatistics(preAddTransformer));
    oldTree.createEmptyTree;
  }

  override def getDefenderMoveProbabilities(rng: RandomGenerator, tlc: ju.function.Consumer[_ >: UctNode[_ <: DS,_ <: DM]]): ju.Map[ImmutableList[DM],java.lang.Double] = {
    val tree = treeJoiner.joinTrees(treeJoinerFunctions, bestPartFilter.getFilteredCollecton())
    val probs = TreeUtils.moveSequenceWeights(
      tree, roundLimit,
      nodeProbabilityWeightFunction, treeDescentLimitCoefficient
    );
    tlc.accept(tree)
    require(probs.size() > 0);
    probs;
  }
}
