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

import com.google.common.collect.ImmutableList
import org.apache.commons.math3.distribution.EnumeratedDistribution
import org.apache.commons.math3.random.{ RandomGenerator }
import org.apache.commons.math3.util.Pair
import org.log4s._
import pl.edu.pw.mini.jk.sg.mixeduct.IterationResult
import pl.edu.pw.mini.jk.sg.uct.{ SelectionFunction, UCTDecisionStatistics }
import pl.edu.pw.mini.jk.sg.uct.tree.UnmodifiableDecisionStatisticsView
import collection.JavaConverters._


case class SelectionStats(avResult: Double, visits: Int) {
  def registerPlayout(result: Double) = SelectionStats((avResult*visits+result)/(visits+1), visits+1)
}

object SelectionStats {
  val empty=SelectionStats(0,0)

  import language.implicitConversions
  implicit def toDecisionStats(ss: SelectionStats) : UnmodifiableDecisionStatisticsView =
    new UnmodifiableDecisionStatisticsView(new UCTDecisionStatistics(ss.visits, ss.visits, ss.avResult))
}


class QualityWeightedHistoryAttacker[M](strategyPoolSize: Int, strategySampleSize: Int, selectionFunction: SelectionFunction, initialStrategy: List[(ImmutableList[M], Double)], rng: RandomGenerator) extends AttackerWithHistory[M] {
  private val strategies = collection.mutable.Map.empty[ImmutableList[M], SelectionStats]
  private val logger = getLogger

  def getInstance(): AttackerStrategy[M] = {
    if (strategies.size > 0) {
      val visitsSum = strategies.values.map(_.visits).sum
      def selectionZeroSupport(stats: SelectionStats) = if(stats.visits == 0) selectionFunction(SelectionStats.empty.registerPlayout(0),visitsSum) else selectionFunction(stats, visitsSum)
      val selFunList = strategies.mapValues(selectionZeroSupport)
      val minSelFun = selFunList.values.min
      val maxSelFun = selFunList.values.max
      val positivizationFactor = Math.min(minSelFun, 0)
      val probabilityList = {
        if(maxSelFun > 0)
          selFunList.map((el) => new Pair[ImmutableList[M], java.lang.Double](el._1,el._2-positivizationFactor)).toList.asJava
        else
          selFunList.map((el) => new Pair[ImmutableList[M], java.lang.Double](el._1,1)).toList.asJava
      }
      val dist = new EnumeratedDistribution[ImmutableList[M]](rng, probabilityList)
      UniformStrategy.fromRepeats(dist.sample(strategySampleSize, Array[ImmutableList[M]]()), rng)
    } else {
      UniformStrategy.fromPairs(initialStrategy, rng)
    }
  }

  def registerIterationResult(iteration: Int,
    iterationResult: IterationResult[_, M, _, _, _],
    resultAttackerStrategy: AttackerStrategy[M]): Unit = {
    val payoff = iterationResult.givenPayoff.defenderPayoff
    strategies ++= resultAttackerStrategy.getPmf.asScala
      .map(x => (x.getFirst, strategies.get(x.getFirst).getOrElse(SelectionStats.empty).registerPlayout(payoff)))
    val visitsSum = strategies.values.map(_.visits).sum
    strategies ++= iterationResult.attackerStrategy.strategyAsPairs.asScala.map(p => (p.getFirst, strategies.get(p.getFirst).getOrElse(SelectionStats.empty)))
    val toDelete = strategies.toSeq.sortBy { case (_, s) =>  selectionFunction(s, visitsSum)}.take(strategies.size-strategyPoolSize)
    logger.debug("N of ent: "+strategies.size)
    logger.debug("Deleting history entries: "+toDelete)
    strategies --= toDelete.map(_._1)
  }
}


