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

import pl.edu.pw.mini.jk.sg.uct.tree.UnmodifiableDecisionStatisticsView
import scala.collection.JavaConverters._
import scala.collection.mutable.{HashMap => MutableHashMap}

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


object NormalTreeJoiner extends TreeJoiner {
  override def joinTrees[S,M](functions: TreeJoinerFunctions)(trees: Iterable[UctNode[S, M]]) : TreeNode[S, M]= {
    val castTrees = trees.map(_.asInstanceOf[TreeNode[S, M]]).toVector
    def allTreeStatistics(tree: TreeNode[S, M]) : Iterable[UnmodifiableDecisionStatisticsView] = tree.getStatistics.asScala.map(_.statistics) ++ tree.successors.values.flatMap(allTreeStatistics)
    val globalStats = functions.calculateGlobalStats(castTrees.flatMap(allTreeStatistics))

    def joinNodes(nodes: Seq[TreeNode[S, M]]) : Option[TreeNode[S, M]] =
      if (nodes.isEmpty) None
      else {
        val moveStats = UctNodeStatisticsContainer.joinStats(nodes.map(_.getStatisticsContainer), functions, trees.size)(globalStats)
      val successorsMap = MutableHashMap[M, TreeNode[S,M]]()
        moveStats.getStatistics.asScala.map(_.move).foreach(m => joinNodes(nodes.flatMap(n => n.successors.get(m).toList)).foreach(successorsMap.put(m,_)))

        Some(new TreeNode(nodes.head.getState, moveStats, nodes.head.rootState, nodes.head.rootPossibleMoves, successorsMap))
      }

    joinNodes(castTrees).get
  }
}
