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

import java.util.{Collection}

import scala.collection.JavaConverters._
import scala.collection.mutable.{HashMap, Queue, Set}

import org.apache.commons.math3.random.RandomGenerator
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.{SelectionFunction, UCTDecisionStatistics}
import pl.edu.pw.mini.jk.sg.uct.movehistory.{MoveHistory, MoveHistoryFactory}
import pl.edu.pw.mini.jk.sg.uct.tree.{PruneTreeFunctions, UctNode, UctMoveStatistics, UctNodeStatisticsContainer, PayoffTransformer}
import scala.runtime.ScalaRunTime
import pl.edu.pw.mini.jk.sg.uct.tree.UnmodifiableDecisionStatisticsView

private[movehistory] final case class NodeKey[S <: AdvanceableState[M], M](state: S, history: MoveHistory[S]) {
  def advance(move: M) = new NodeKey[S, M](state.advance(move).asInstanceOf[S], history.appendMove(state, move))

  override val hashCode = ScalaRunTime._hashCode(this)
}

private[movehistory] final class LimitedNode[S <: AdvanceableState[M], M](val key: NodeKey[S, M], val stat: UctNodeStatisticsContainer[M], val tr: LimitedHistoryTree[S, M]) extends UctNode[S, M] {
  import tr._
  override def getState: S = key.state
  override def getBestMove(rng: RandomGenerator): M = stat.getBestMove(rng)
  def getSuccessorX(move: M): Option[LimitedNode[S, M]] = {
      val ak = key.advance(move)
      nodes.get(ak).map(n => new LimitedNode[S, M](ak, n, tr))
  }

  override def selectOrRandomUnvisidedMove(fun: SelectionFunction, rng: RandomGenerator): M = stat.selectOrRandomUnvisitedMove(fun, rng)
  override def getSuccessor(move: M): Option[UctNode[S, M]] = getSuccessorX(move)
  override def getStatistics: Collection[UctMoveStatistics[M]] = stat.getStatistics
  override def addSuccessor(move: M, state: S, possibleMoves: java.util.List[M]): UctNode[S, M] = {
    val ak = key.advance(move)
    new LimitedNode[S, M](ak, nodes.get(ak) match {
      case Some(n) => n
      case None => {
        val newStats = new UctNodeStatisticsContainer[M](possibleMoves.asScala)
        nodes.put(ak, newStats)
        newStats
      }
    }, tr)
  }
  override def selectMove(fun: SelectionFunction): Option[M] = stat.selectMove(fun)

  override def exportTree = ???

  override def getRealVisitsCount: Int = stat.getRealVisitsCount
  override def getVisitsCount: Int = stat.getVisitsCount
  override def registerVisit(move: M): Unit = stat.registerVisit(move)
  override def registerPayoff(move: M, payoff: Double): Unit = stat.registerPayoff(move, payoff)
  override def moveCount: Int = stat.moveCount

  override def createEmptyTree: UctNode[S, M] = LimitedHistoryTree.emptyTree[S, M](mhf, initialState).getRoot
  override def reduceStatistics(reduceFunction: java.util.function.Function[UCTDecisionStatistics, UCTDecisionStatistics]): UctNode[S, M] = new LimitedHistoryTree[S, M](mhf, initialState, nodes.map { case (k, v) => (k, v.reduceStatistics(reduceFunction)) }).getRoot

  override def reduceStatistics(reduceFunction: UCTDecisionStatistics => UCTDecisionStatistics): UctNode[S, M] = new LimitedHistoryTree[S, M](mhf, initialState, nodes.map { case (k, v) => (k, v.reduceStatistics(reduceFunction)) }).getRoot

  override def pruneTree(keepPredicate: PruneTreeFunctions): UctNode[S, M] = {
    val visitedKeys = Set[NodeKey[S, M]]()
    val nodesToProcess = Queue[LimitedNode[S, M]]()
    val outputNodes = HashMap[NodeKey[S, M], UctNodeStatisticsContainer[M]]()

    visitedKeys += rootKey
    nodesToProcess += new LimitedNode(rootKey, nodes(rootKey), tr)

    while (!nodesToProcess.isEmpty) {
      val n = nodesToProcess.dequeue()
      outputNodes += ((n.key, n.stat.typedClone))

      val nodeStat = keepPredicate.nodeStatistics(n)
      n.getStatistics.asScala.filter(s => keepPredicate.movePredicate(s.statistics, nodeStat)).map(m => n.getSuccessorX(m.move))
        .collect { case Some(o) if !visitedKeys.contains(o.key) => o }
        .foreach(node => {
          visitedKeys += node.key
          nodesToProcess.enqueue(node)
        })
    }

    new LimitedHistoryTree[S, M](mhf, initialState, outputNodes).getRoot
  }

  override def allStatistics : Stream[UnmodifiableDecisionStatisticsView] = tr.nodes.values.toStream.flatMap(_.statistics.toStream.map(_._2.unmodifiableView))

  override def transformStatistics(payoffTransformer: PayoffTransformer) : UctNode[S, M] = {
    val visitedKeys = Set[NodeKey[S, M]]()
    val nodesToProcess = Queue[LimitedNode[S, M]]()
    val outputNodes = HashMap[NodeKey[S, M], UctNodeStatisticsContainer[M]]()
    val globalStat : payoffTransformer.GlobalStat = payoffTransformer.calculateGlobalStat(
      tr.nodes.values.view.flatMap(_.statistics.values.toStream.map(_.unmodifiableView))
    )

    visitedKeys += rootKey
    nodesToProcess += new LimitedNode(rootKey, nodes(rootKey), tr)

    while (!nodesToProcess.isEmpty) {
      val n = nodesToProcess.dequeue()
      val nodeStat = payoffTransformer.calculateNodeStat(n)

      outputNodes += ((n.key, n.stat.transformStatistics(payoffTransformer)(globalStat, nodeStat)))
      n.getStatistics.asScala.map(m => n.getSuccessorX(m.move))
        .collect { case Some(o) if !visitedKeys.contains(o.key) => o }
        .foreach(node => {
          visitedKeys += node.key
          nodesToProcess.enqueue(node)
        })
    }

    new LimitedHistoryTree[S, M](mhf, initialState, outputNodes).getRoot
  }
}

class LimitedHistoryTree[S <: AdvanceableState[M], M] private[movehistory] (
  private[movehistory] val mhf: MoveHistoryFactory,
  private[movehistory] val initialState: S,
  private[movehistory] val nodes: HashMap[NodeKey[S, M], UctNodeStatisticsContainer[M]]
) {

  val rootKey = NodeKey[S, M](initialState, mhf.create())
  def getRoot: UctNode[S, M] = new LimitedNode[S, M](rootKey, nodes(rootKey), this)
}

object LimitedHistoryTree {
  def emptyTree[S <: AdvanceableState[M], M](
    mhf: MoveHistoryFactory,
    initialState: S
  ) = new LimitedHistoryTree[S, M](mhf, initialState, HashMap((NodeKey(initialState, mhf.create()), new UctNodeStatisticsContainer(initialState.getPossibleMoves.asScala))))

}
