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

import pl.edu.pw.mini.sg.mcts.tree.StatContainer


private[stat] final class Statistics[Q](
  var stat: Q,
  var visitsCount: Int
)

final class MapContainer[M,Q] private(
  map: Map[M, Statistics[Q]]
) extends StatContainer[M,Q] {
  def registerPlayout(result: Double,move: M)(implicit sc: StatCalculator{type StatStorage = Q}): Unit = {
    val s = map(move)
    s.visitsCount = s.visitsCount+1
    s.stat = sc.registerPlayout(s.stat, result)
  }

  def unvisitedMoves: Vector[M] = map.view.collect{case (move, s) if s.visitsCount == 0 => move}.toVector
  def visitCounter: Int = map.view.map(_._2.visitsCount).sum

  override def estimates(implicit sc: StatCalculator{type StatStorage = Q}) : Seq[MoveEstimate[M]] = map.view.map{case (move, stat) => MoveEstimate(move, sc.estimate(stat.stat, stat.visitsCount), stat.visitsCount)}.foldLeft(List[MoveEstimate[M]]())((list, el) => el :: list)

  override def isUnvisitedMovesEmpty: Boolean = !map.values.exists(_.visitsCount == 0)
  override def isEstimatesEmpty: Boolean = map.isEmpty
}

object MapContainer {
  def create[M,Q](moves: Seq[M], initialQ: Q) = new MapContainer[M,Q](moves.map((_, new Statistics[Q](initialQ,0))).toMap)
}
