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

import org.apache.commons.math3.random.RandomGenerator
import pl.edu.pw.mini.jk.sg.game.common.DefenderVisibleState
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.probabilityfunctions.NodeProbabilityWeightFunction
import pl.edu.pw.mini.jk.sg.mixeduct.util.TreeUtils
import pl.edu.pw.mini.jk.sg.uct.tree.{PruneTreeFunctions, UctNode}

import java.{util => ju}

import pl.edu.pw.mini.jk.sg.uct.UCTDecisionStatistics
final class FullTreeTransfer[S <: AdvanceableState[M] with DefenderVisibleState, M](
  roundLimit: Int,
  nodeProbabilityWeightFunction: NodeProbabilityWeightFunction,
  pruneFunction: PruneTreeFunctions,
  treeDescentFilterCoefficient: Double,
  statisticsReduction : Option[UCTDecisionStatistics => UCTDecisionStatistics]
) extends AttackersInformationTransfer[M, S] {
  var tree: Option[UctNode[S, M]] = None

  override def newAttackerTree(oldTree: UctNode[S, M], rng: RandomGenerator) = {
    tree = Some({ val tr = oldTree
      .pruneTree(pruneFunction)
      val tr2 = statisticsReduction.map{r =>  tr.reduceStatistics(r)}.getOrElse(tr)
      tr2
    })
    tree.get
  }

  override def getDefenderMoveProbabilities(rng: RandomGenerator, tlc: ju.function.Consumer[_ >: UctNode[_ <: S,_ <: M]]) =
    tree match {
      case Some(t) => {
        tlc.accept(t)
        TreeUtils.moveSequenceWeights(t, roundLimit, nodeProbabilityWeightFunction, treeDescentFilterCoefficient)
      }
      case None => throw new IllegalStateException("Cannot calculate probabilities without registered tree")
    }
}
