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

import org.apache.commons.math3.random.RandomGenerator
import pl.edu.pw.mini.jk.sg.uct.{ SelectionFunction, UCTDecisionStatistics }

import scala.collection.JavaConverters._
import scala.collection.mutable.HashMap
import scala.collection.immutable.{HashMap => IHashMap}

final class UctNodeStatisticsContainer[M](
  val allPossibleMoves : Vector[M],
  val statistics : HashMap[M, UCTDecisionStatistics],
  private var visitsCount: Int,
  private var visitedMoves: Int,
  private var bestMove: Option[M]
){
  def this(possibleMoves: Seq[M]) =
    this(possibleMoves.toVector, HashMap(), 0, 0, None)

  def this(allPossibleMoves: Vector[M], statistics : HashMap[M, UCTDecisionStatistics]) = this(allPossibleMoves, statistics,
    statistics.values.map(_.visitsCount).sum, statistics.values.count(_.visitsCount > 0), None)

//  lazy val statisticsL = statistics.view.map{case(m, s) => UctMoveStatistics.apply(m,new UnmodifiableDecisionStatisticsView(s))}.toVector
  def moveCount = allPossibleMoves.size
  def getStatistics(): java.util.Collection[UctMoveStatistics[M]] = {
    statistics.view.map{case (m,s) => UctMoveStatistics.apply(m, new UnmodifiableDecisionStatisticsView(s))}.toVector.asJava
  }

  private def getStatistics(move: M) : UCTDecisionStatistics =
    if(statistics.contains(move))
      statistics(move)
    else {
      val s = new UCTDecisionStatistics(0,0,0)
      statistics += ((move, s))
      s
    }


  def getBestMove(rng: RandomGenerator) : M = {
    if(visitedMoves < allPossibleMoves.size) {
      allPossibleMoves(rng.nextInt(allPossibleMoves.size))
    } else {
      bestMove match {
        case Some(m) => m
        case None => {
          val bMove = statistics.toSeq.maxBy(_._2.averagePayoff)._1
          bestMove = Some(bMove)
          bMove
        }
      }
    }
  }

  def getVisitsCount : Int = visitsCount

  //TODO: źle
  def selectOrRandomUnvisitedMove(sf: SelectionFunction, rng: RandomGenerator) =
    selectMove(sf).getOrElse(allPossibleMoves.filter(x => !statistics.contains(x)).drop(rng.nextInt(moveCount - visitedMoves)).head)

  def selectMove(sf: SelectionFunction) : Option[M] =
    if(visitedMoves < moveCount)
      None
    else
      Some(statistics.maxBy{case (k,v) => sf.apply(v, visitsCount)}._1)


  def registerVisit(move: M): Unit = {
    visitsCount = visitsCount + 1
    if(getStatistics(move).getVisitsCount == 0) { visitedMoves = visitedMoves + 1 };
    getStatistics(move).increaseVisitCounter()
    bestMove = None
  }

  def getRealVisitsCount : Int = statistics.values.map(_.getRealVisitsCount).sum

  def registerPayoff(move: M, payoff: Double): Unit = {
    getStatistics(move).registerPayoff(payoff)
    bestMove = None
  }

  def typedClone = new UctNodeStatisticsContainer[M](allPossibleMoves,statistics.map{case (m, s) => (m, s.typedClone)}, visitsCount, visitedMoves, bestMove)

  def reduceStatistics(function: UCTDecisionStatistics => UCTDecisionStatistics) =
    new UctNodeStatisticsContainer(allPossibleMoves, statistics.map{case (m, s) => (m, function(s))})

  def reduceStatistics(function: java.util.function.Function[UCTDecisionStatistics, UCTDecisionStatistics]) : UctNodeStatisticsContainer[M] = reduceStatistics(function.apply(_))

  def transformStatistics(
    transformer: PayoffTransformer
  )(
    globalStat: transformer.GlobalStat,
    nodeStat: transformer.NodeStat
  ) : UctNodeStatisticsContainer[M] = new UctNodeStatisticsContainer[M](
    allPossibleMoves,
    HashMap(statistics.mapValues(s => transformer.transformStatistics(s.unmodifiableView, nodeStat, globalStat)).toSeq: _*)
  )
}


object UctNodeStatisticsContainer {
  def mapFromMoveList[M](possibleMoves: java.util.List[M]) : HashMap[M, UCTDecisionStatistics] = mapFromMoveList(possibleMoves.asScala)

  def mapFromMoveList[M](possibleMoves: Seq[M]): HashMap[M, UCTDecisionStatistics] =  {
    val builder = HashMap.newBuilder[M, UCTDecisionStatistics]
    builder.sizeHint(possibleMoves.size)
    //builder++=(possibleMoves.view.map((_, new UCTDecisionStatistics(0,0,0))))
    possibleMoves.foreach(x => builder += ((x, new UCTDecisionStatistics(0, 0, 0))))
    builder.result()
  }

  def joinStats[M](seq: Seq[UctNodeStatisticsContainer[M]], joiner : TreeJoinerFunctions, treeCount: Int)(globalStats: joiner.GS) = {

    val contextStats = joiner.calculateContextStats(seq.flatMap(_.statistics).map(_._2))
    val moveMap = seq.map(x => IHashMap(x.statistics.view.map { case (m,s) => (m, Vector(s.unmodifiableView)) }.toSeq : _*))
      .reduce((m1, m2) => m1.merged(m2){case ((m1, v1), (m2, v2)) => (m1, v1++v2) })
      .map { case (m, sv) => (m, joiner.joinStats(sv, seq.length, contextStats, globalStats)) }
    new UctNodeStatisticsContainer[M](seq(0).allPossibleMoves, HashMap(moveMap.toSeq : _*))
  }
}
