package pl.edu.pw.mini.sg.mcts.tree.simple

import pl.edu.pw.mini.sg.mcts.tree.stat.{MapContainer, StatCalculator}
import pl.edu.pw.mini.sg.mcts.tree.{ MctsNode, StatContainer }
import scala.collection.mutable

final class SimpleNode[M, S, Q](
  override val state: S,
  successors: mutable.Map[M, mutable.Map[S, SimpleNode[M,S,Q]]],
  statisticsContainer: StatContainer[M, Q],
  statCalculator: StatCalculator{type StatStorage = Q}
) extends MctsNode[M, S] {
  def addSuccessor(move: M,newState: S,possibleMoves: Vector[M]): pl.edu.pw.mini.sg.mcts.tree.MctsNode[M,S] = {
    val stateMap = successors.getOrElseUpdate(move, mutable.Map())
    val newNode = new SimpleNode[M, S, Q](newState, mutable.Map(), MapContainer.create[M,Q](possibleMoves, statCalculator.emptyStat), statCalculator)
    stateMap.put(newState, newNode) match {
      case None => newNode
      case _ => throw new IllegalStateException("Can't add existing node");
    }
  }

  implicit val sc = statCalculator

  def estimates: Seq[pl.edu.pw.mini.sg.mcts.tree.stat.MoveEstimate[M]] = statisticsContainer.estimates
  def registerPlayout(result: Double,move: M): Unit = statisticsContainer.registerPlayout(result, move)
  def visitCounter: Int = statisticsContainer.visitCounter

  def successor(move: M, state: S): Option[SimpleNode[M,S, Q]] =
    successors.get(move).flatMap(_.get(state))
  def unvisitedMoves: Seq[M] = statisticsContainer.unvisitedMoves
  def isUnvisitedMovesEmpty: Boolean = statisticsContainer.isUnvisitedMovesEmpty
  def isEstimatesEmpty: Boolean = statisticsContainer.isEstimatesEmpty
}
