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

import scala.collection.immutable.HashMap

import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.tree.{TreeJoiner, TreeJoinerFunctions, UctNode, UctNodeStatisticsContainer}



object LimitedHistoryTreeJoiner extends TreeJoiner {
  override def joinTrees[S <: AdvanceableState[M], M](functions: TreeJoinerFunctions)(trees: Iterable[UctNode[S,M]]) : UctNode[S,M] = {
    val castTrees = trees.map(_.asInstanceOf[LimitedNode[S, M]])

    val globalStats = functions.calculateGlobalStats(castTrees.flatMap(t => t.tr.nodes.values.flatMap(c => c.statistics.map(_._2))))

    val outMap = castTrees.map(n => HashMap(n.tr.nodes.map{case (k, v) => (k, Vector(v))}.toSeq :_*))
      .reduce(_.merged(_){case ((k1, v1), (k2, v2)) => (k1, v1++v2)})
    .map{case (k, v) => (k, UctNodeStatisticsContainer.joinStats(v, functions, trees.size)(globalStats))}

    new LimitedHistoryTree(
      castTrees.head.tr.mhf,
      castTrees.head.tr.initialState,
      scala.collection.mutable.HashMap(outMap.toSeq :_*)
    ).getRoot
  }
}
