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

import pl.edu.pw.mini.jk.sg.uct.tree.{PruneTreeFunctions, UctMoveStatistics}

import scala.collection.JavaConverters._
import scala.collection.mutable.HashMap

import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.UCTDecisionStatistics
import pl.edu.pw.mini.jk.sg.uct.tree.{UctNode, UctNodeStandardStats, UctNodeStatisticsContainer, PayoffTransformer}
import pl.edu.pw.mini.jk.sg.uct.tree.export.ExportUctTree
import _root_.pl.edu.pw.mini.jk.sg.uct.tree.export.ExportStatistics
import pl.edu.pw.mini.jk.sg.uct.tree.UnmodifiableDecisionStatisticsView

final class TreeNode[S, M] private[normal] (
  val state: S,
  statisticsContainer: UctNodeStatisticsContainer[M],
  val rootState: S,
  val rootPossibleMoves: Seq[M],
  val successors: HashMap[M, TreeNode[S, M]]
) extends UctNodeStandardStats[S, M](statisticsContainer) {

  private def this(state: S, possibleMoves: Seq[M], rootState: S, rootPossibleMoves: Seq[M]) =
    this(
      state,
      new UctNodeStatisticsContainer[M](possibleMoves),
      rootState, rootPossibleMoves, HashMap[M, TreeNode[S, M]]()
    )

  override def addSuccessor(move: M, state: S, possibleMoves: java.util.List[M]): UctNode[S, M] = {
    if (successors.contains(move)) throw new IllegalStateException("Node arledy containst successor for this move")
    else {
      val node = new TreeNode(state, possibleMoves.asScala, rootState, rootPossibleMoves)
      successors += ((move, node))
      node
    }
  }

  def getStatisticsContainer = statisticsContainer
  override def getState(): S = state

  override def getSuccessor(move: M): Option[UctNode[S, M]] = successors.get(move)
  override def createEmptyTree() = new TreeNode[S, M](rootState, rootPossibleMoves, rootState, rootPossibleMoves)

  override def pruneTree(pruneFunctions: PruneTreeFunctions): TreeNode[S, M] = {
    val stats = pruneFunctions.nodeStatistics(this)
    val movesToKeep = statisticsContainer.statistics.filter { case (_, st) => pruneFunctions.movePredicate(st, stats) }.map(_._1).toVector
    new TreeNode[S, M](state, statisticsContainer.typedClone, rootState, rootPossibleMoves,
      successors.filter { case (move, _) => movesToKeep.contains(move) }
        .map { case (move, succ) => (move, succ.pruneTree(pruneFunctions)) })
  }
  override def reduceStatistics(reduction: java.util.function.Function[UCTDecisionStatistics, UCTDecisionStatistics]): TreeNode[S, M] = new TreeNode[S, M](
    state,
    statisticsContainer.reduceStatistics(reduction),
    rootState, rootPossibleMoves,
    successors.map { case (move, node) => (move, node.reduceStatistics(reduction)) }
  )

  override def reduceStatistics(reduction: UCTDecisionStatistics => UCTDecisionStatistics): TreeNode[S, M] = new TreeNode[S, M](
    state,
    statisticsContainer.reduceStatistics(reduction),
    rootState, rootPossibleMoves,
    successors.map { case (move, node) => (move, node.reduceStatistics(reduction)) }
  )

  override def exportTree: ExportUctTree[S,M] = ExportUctTree(
    state, successors.map{case (m, succ) =>
      (m, (ExportStatistics.fromUDS(statisticsContainer.statistics(m)), succ.exportTree))}.toMap
  )

  override def allStatistics: Stream[UnmodifiableDecisionStatisticsView] = {
    statisticsContainer.statistics.values.toStream.map(_.unmodifiableView)++
    successors.values.flatMap(_.allStatistics)
  }

  private def  transformStatisticsWithStats(payoffTransformer: PayoffTransformer)(globalStat: payoffTransformer.GlobalStat) : TreeNode[S, M] = {
    val nodeStat = payoffTransformer.calculateNodeStat(this)
    new TreeNode(
      state,
      statisticsContainer.transformStatistics(payoffTransformer)(globalStat, nodeStat),
      rootState,
      rootPossibleMoves,
      HashMap(successors.mapValues(_.transformStatisticsWithStats(payoffTransformer)(globalStat)).toSeq :_*)
    )
  }

  override def transformStatistics(payoffTransformer: PayoffTransformer) : TreeNode[S, M] =
    transformStatisticsWithStats(payoffTransformer)(payoffTransformer.calculateGlobalStat(allStatistics))
}

object TreeNode {
  def fromInitialState[S <: AdvanceableState[M], M](initial: S): TreeNode[S, M] =
    new TreeNode(initial, initial.getPossibleMoves.asScala, initial, initial.getPossibleMoves.asScala)
}
