Probability distribution for outcome of rolling multiple dice

D20.scala

import cats.Monoid
import cats.implicits._

import scala.collection.immutable.Stream

import com.cra.figaro.language.{ Constant, Element, Apply, Chain, Atomic }
import com.cra.figaro.algorithm.factored.{ VariableElimination }
import com.cra.figaro.algorithm.sampling.Importance
import com.cra.figaro.algorithm.{ ProbQueryAlgorithm }
import com.cra.figaro.library.atomic.discrete.FromRange
import com.cra.figaro.library.collection.{ Container, ContainerElement }

type Algorithm = (Element[_]) => ProbQueryAlgorithm
object Algorithm {
  def importance(element: Element[_]) = Importance(168000, element)
  def variableElimination(element: Element[_]) = VariableElimination(element)
}

def Run[A](alg: ProbQueryAlgorithm): ProbQueryAlgorithm = {
	alg.start()
	alg
}

type Probability[A] = (Double, A)
type Probabilities[A] = Seq[Probability[A]]
def Distribution[A](element: Element[A])(implicit alg: Algorithm): Probabilities[A] = {
  Run(alg(element)).distribution(element)
}

type Sorted[F] = F
def Sorted[A: Ordering](probs: Probabilities[A]): Sorted[Probabilities[A]] =
  probs.toSeq.sortBy(_._2)

def Cumulative[A: Numeric](probs: Sorted[Probabilities[A]]): Probabilities[A] = probs // TODO

def ProbabilityTable[A](probs: Probabilities[A]) = {
  for (
    (odds, value) <- probs
  ) {
    println(f"${odds * 100}%05.2f %%: $value")
  }
  probs
}

type Repeated[A] = ContainerElement[Int, A];
def Repeat[A](times: Int, element: => Element[A]): Repeated[A] = {
  Container(
    Constant(times),
    (_index: Int) => element
  )
}

def Sum[A: Monoid](elements: Repeated[A]): Element[A] = {
  elements.foldLeft(Monoid[A].empty)(Monoid[A].combine)
}

def Maximum[A: Ordering](elements: Repeated[A]): Element[A] =
	elements.reduce((left: A, right: A) => if (Ordering[A].gt(left, right)) left else right)

def Minimum[A: Ordering](elements: Repeated[A]): Element[A] =
	elements.reduce((left: A, right: A) => if (Ordering[A].lt(left, right)) left else right)


object Dice {
  def Die(n: Int): Atomic[Int] = FromRange(1, n+1)
  
  def D20 = Die(20)
  def D12 = Die(12)
  def D10 = Die(10)
}


object Rolls {
  def Roll(times: Int, dice: => Element[Int]): Element[Int] =
    Sum(Repeat(times, dice))

  def WithAdvantage(dice: => Element[Int]) =
    Maximum(Repeat(2, dice))

  def WithDisadvantage(dice: => Element[Int]) =
    Minimum(Repeat(2, dice))

  // AnyDice reference table for cross-checking https://anydice.com/program/1203
  def TwoD20 = Roll(2, Dice.D20)
  def D20WithAdvantage = WithAdvantage(Dice.D20)
  def D20WithDisadvantage = WithDisadvantage(Dice.D20)
}

ProbabilityTable(
  Cumulative(
    Sorted(
      Distribution(
        Rolls.D20WithAdvantage
      )(Algorithm.importance)
    )
  )
)