package pl.edu.pw.mini.jk.sg.mixeduct.attacker
import com.google.common.collect.ImmutableList
import java.{util => ju, lang => jl}
import org.apache.commons.math3.distribution.EnumeratedDistribution
import org.apache.commons.math3.random.RandomGenerator
import scala.collection.JavaConverters._
import org.apache.commons.math3.util.Pair

final class DiscountingStrategyJoin[M](
    discountingFactor: Double,
    distributions: ImmutableList[EnumeratedDistribution[ImmutableList[M]]],
    rng: RandomGenerator
) extends AttackerStrategy[M] {
  private val weightedPairs = Stream
    .iterate(1.0, distributions.size)(_ * discountingFactor)
    .reverse
    .view
    .zip(distributions.asScala)
    .map { case (prob, dist) => new Pair(dist, new jl.Double(prob)) }
    .toVector

  override def sample(): ImmutableList[M] = {
    val dist = new EnumeratedDistribution(rng, weightedPairs.asJava)
    dist.sample().sample()
  }
  override def getPmf()
      : ju.List[org.apache.commons.math3.util.Pair[ImmutableList[M], java.lang.Double]] =
    ProbabilisticStrategyAttackerWithHistoryFactory
      .joinDistributions(distributions, discountingFactor, rng)
      .getPmf()
}
