package pl.edu.pw.mini.jk.sg.mixeduct.util

import scala.annotation.tailrec
import scala.collection.JavaConverters._
import scala.collection.immutable.HashMap
import com.google.common.collect.ImmutableList 
import org.apache.commons.math3.random.RandomGenerator
import pl.edu.pw.mini.jk.sg.game.common.{DefenderVisibleState, DefenderMove}
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableState
import pl.edu.pw.mini.jk.sg.uct.tree.UctNode
import pl.edu.pw.mini.jk.sg.mixeduct.transfer.probabilityfunctions.NodeProbabilityWeightFunction
import pl.edu.pw.mini.jk.sg.game.common.DefenderMove
import scala.Vector
import pl.edu.pw.mini.jk.sg.game.common.advanceable.AdvanceableDefenderState
import pl.edu.pw.mini.jk.sg.uct.tree.UnmodifiableDecisionStatisticsView

object TreeUtils {

  @tailrec
  def advanceStateByMoves[M](state: AdvanceableState[M], moves: Seq[M]) : AdvanceableState[M] = {
    if(moves.length == 0)
      state
    else
      advanceStateByMoves(state.advance(moves.head), moves.tail)
  }

  def getBestPathInTree[S <: AdvanceableState[M], M](root: UctNode[S, M], timeSteps: Int, rng: RandomGenerator) : List[M] = {
    val bestMove = root.getBestMove(rng);

    root.getSuccessor(bestMove) match {
      case  Some(n) if root.getState.getTimeStep < timeSteps => bestMove :: getBestPathInTree(n, timeSteps, rng)
      case _ => List()
    }
  }

  def getBestPathInTreeExtended[S <: AdvanceableState[M], M](root: UctNode[S, M], length: Int, rng: RandomGenerator) = {
    val path = getBestPathInTree(root, length, rng)
    @tailrec
    def randomPath(state: AdvanceableState[M], length: Int, currentPath: Vector[M]) : Vector[M] = {
      if(state.getTimeStep >= length)
        currentPath
      else {
        val move = state.getRandomMove(rng)
        randomPath(state.advance(move).asInstanceOf[AdvanceableState[M]], length, currentPath:+move)
      }
    }
    path ++ (randomPath(advanceStateByMoves(root.getState, path), length, Vector())).toList
  }

  def getBestPathInTreeExtendedJava[S <: AdvanceableState[M], M <: DefenderMove](root: UctNode[S, M], length: Int, rng: RandomGenerator) = ImmutableList.copyOf(getBestPathInTreeExtended(root, length, rng).asJava.asInstanceOf[java.lang.Iterable[M]])

  def toImmutableList[T](col: Iterable[T]) : ImmutableList[T] =
    ImmutableList.copyOf(col.asJava)

  @tailrec def finishPath[S <: DefenderVisibleState with AdvanceableState[M], M](node: S, gameLength: Int, path: Vector[M]) : Vector[M] =
    if(node.getTimeStep == gameLength) path
    else {
      val move = node.getPossibleMoves().get(0)
      val nextState = node.advance(move).asInstanceOf[S]
      finishPath(nextState, gameLength, path :+ move)
    }

  def moveSequenceWeights[S <: DefenderVisibleState with AdvanceableState[M], M](root: UctNode[S, M], gameLength: Int,
    weightFunction: NodeProbabilityWeightFunction, filterCoefficient: Double) : java.util.Map[ImmutableList[M], java.lang.Double] = {
    val globStat = weightFunction.calculateGlobalStat(root.allStatistics)
    def sequenceWalker(node: UctNode[S, M], currentSequence: Vector[M], weight: Double) : HashMap[ImmutableList[M], java.lang.Double] =
      if(node.getState.getTimeStep() == gameLength)
        return HashMap[ImmutableList[M], java.lang.Double] ((toImmutableList(currentSequence), weight))
      else if (node.getVisitsCount > 0) {
        val maxWeight = node.getStatistics.asScala.map(m => weightFunction.probabilityWeight(m.statistics, node, globStat)).max
        val ret = node.getStatistics.asScala.map( s => {
          lazy val w = weightFunction.probabilityWeight(s.statistics, node, globStat)
          node.getSuccessor(s.move) match {
            case Some(n) if w >= maxWeight*filterCoefficient =>
              Some(sequenceWalker(n, currentSequence:+s.move, if(weightFunction.leafOnlyWeight) w else weight*w))
            case _ => None
          }
        }).collect{case Some(x) => x}.foldLeft(HashMap.empty[ImmutableList[M], java.lang.Double])(_++_)
        if(ret.isEmpty) {
          println("Empty seq!")
        }
        //        Jak często to jest puste???
        ret
      } else {
        val path = finishPath(node.getState, gameLength, currentSequence)
        HashMap(toImmutableList(path) ->  weight)
      }

    sequenceWalker(root, Vector(), 1).asJava
  }
}
