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

import java.util.Collection
import java.util.Comparator
import java.util.function.Function
import java.util.stream.Collectors
import com.google.common.collect.MinMaxPriorityQueue
import pl.edu.pw.mini.jk.sg.game.common.DefenderMove
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.tree.UctNode

final class BestTreeFilter[DM <: DefenderMove, DS <: AdvanceableState[DM]](
  bestTreeCount: Int,
  treeQualityFunction: UctNode[DS, DM] => Double,
) extends BestPartFilter[DM, DS] {
  private val trees  : MinMaxPriorityQueue[WeightedTree[DM, DS]] = MinMaxPriorityQueue
    .orderedBy(Comparator.comparingDouble((wt: WeightedTree[DM, DS]) => wt.weight).reversed())
    .maximumSize(bestTreeCount).create();

  override def getFilteredCollecton() : Collection[UctNode[DS, DM]] =
    trees.stream().map[UctNode[DS, DM]](t => t.tree).collect(Collectors.toList())

  override def registerTree(tree: UctNode[DS, DM]) : Collection[UctNode[DS, DM]] = {
    trees.add(new WeightedTree(tree, treeQualityFunction.apply(tree)));
    getFilteredCollecton();
  }
}
