module DiceRolling export d4, d6, d8, d10, d12, d20, distribution using FFTW using OffsetArrays using RecipesBase abstract type AbstractDie end """ Die{F} <: AbstractDie A single die with `F` faces """ struct Die{F} <: AbstractDie end """ DieRolls{F} <: AbstractDie Represents an `F`-sided die rolled multiple times. """ struct DieRolls{F} <: AbstractDie rolls::Int die::Die{F} function DieRolls(rolls::Integer, die::Die{F}) where F if F < 0 return new{-F}(-rolls, Die{-F}()) else return new{F}(rolls, die) end end end DieRolls(rolls::Integer, die::Integer) = DieRolls(rolls, Die{die}()) """ CompoundDieRolls <: AbstractDie A combination of dice rolls and optional bonus. """ struct CompoundDieRolls <: AbstractDie die::Vector{DieRolls} bonus::Int function CompoundDieRolls(d::CompoundDieRolls) return new(copy(d.die), d.bonus) end function CompoundDieRolls(dice, bonus::Int = 0) bonus += mapreduce(diebonus, +, dice) rolls = Dict{Int,Int}() for d in dice r = get!(rolls, diebase(d), 0) rolls[diebase(d)] += dierolls(d) end dice′ = [DieRolls(r, d) for (d, r) in rolls] return new(dice′, bonus) end end CompoundDieRolls(d::Die, bonus::Int = 0) = CompoundDieRolls([1d], bonus) CompoundDieRolls(d::DieRolls, bonus::Int = 0) = CompoundDieRolls([d], bonus) # Accessors diebase(d::Die{F}) where {F} = F diebase(d::DieRolls) = diebase(d.die) dierolls(d::Die) = 1 dierolls(d::DieRolls) = d.rolls diebonus(::Die) = 0 diebonus(::DieRolls) = 0 diebonus(d::CompoundDieRolls) = d.bonus # The fundamental D&D dice const d4 = Die{4}() const d6 = Die{6}() const d8 = Die{8}() const d10 = Die{10}() const d12 = Die{12}() const d20 = Die{20}() # Define math over dice objects Base.:-(d::Die) = DieRolls(-1, d) Base.:-(d::DieRolls) = DieRolls(-d.rolls, d.die) Base.:-(d::CompoundDieRolls) = CompoundDieRolls(map(-, d.die), -d.bonus) Base.:+(d::Die) = DieRolls(1, d) Base.:+(d::DieRolls) = d Base.:+(d::CompoundDieRolls) = d Base.:*(r::Integer, d::Die) = DieRolls(Int(r), d) Base.:*(r::Integer, d::DieRolls) = DieRolls(Int(r * d.rolls), d.die) Base.:*(r::Integer, d::CompoundDieRolls) = CompoundDieRolls(map(d -> r * d, d.die), r * d.bonus) Base.:+(b::Integer, d::AbstractDie) = d + b Base.:-(b::Integer, d::AbstractDie) = -d + b Base.:+(d::AbstractDie, b::Integer) = CompoundDieRolls([d], b) Base.:-(d::AbstractDie, b::Integer) = CompoundDieRolls([d], -b) Base.:+(d1::AbstractDie, d2::AbstractDie) = CompoundDieRolls(d1) + CompoundDieRolls(d2) Base.:-(d1::AbstractDie, d2::AbstractDie) = CompoundDieRolls(d1) - CompoundDieRolls(d2) Base.:+(d::CompoundDieRolls, b::Integer) = CompoundDieRolls(d.die, d.bonus + b) Base.:-(d::CompoundDieRolls, b::Integer) = CompoundDieRolls(d.die, d.bonus - b) Base.:+(d1::CompoundDieRolls, d2::CompoundDieRolls) = CompoundDieRolls(vcat(d1.die, d2.die), d1.bonus + d2.bonus) Base.:-(d1::CompoundDieRolls, d2::CompoundDieRolls) = d1 + -d2 Base.minimum(d::Die{F}) where {F} = F < 0 ? -F : 1 Base.maximum(d::Die{F}) where {F} = F < 0 ? -1 : F Base.minimum(d::DieRolls) = d.rolls < 0 ? d.rolls * maximum(d.die) : d.rolls Base.maximum(d::DieRolls) = d.rolls > 0 ? d.rolls * maximum(d.die) : d.rolls Base.minimum(d::CompoundDieRolls) = mapreduce(minimum, +, d.die) + d.bonus Base.maximum(d::CompoundDieRolls) = mapreduce(maximum, +, d.die) + d.bonus Base.extrema(d::AbstractDie) = (minimum(d), maximum(d)) # Pretty-printing of dice rolling objects function Base.show(io::IO, d::Die{F}) where {F} if F > 0 print(io, "d", F) else print(io, "-1d", -F) end end Base.show(io::IO, d::DieRolls) = print(io, d.rolls, d.die) function Base.show(io::IO, d::CompoundDieRolls) for ii in eachindex(d.die) d′ = d.die[ii] if ii == firstindex(d.die) print(io, d′) else if dierolls(d′) < 0 print(io, " - ", -d′) else print(io, " + ", d′) end end end if !iszero(d.bonus) if d.bonus < 0 print(io, " - ", -d.bonus) else print(io, " + ", d.bonus) end end end distribution(d::Die, maxN = nothing) = distribution(1*d + 0) distribution(d::DieRolls, maxN = nothing) = distribution(d + 0) function distribution(d::CompoundDieRolls) # Deal with the bonus purely in terms of axis indices if d.bonus != 0 bonus = d.bonus P = distribution(d - bonus) return OffsetVector(parent(P), axes(P, 1) .+ bonus) end # Calculate the "negative" distribution if it is biased negative if abs(minimum(d)) > abs(maximum(d)) P = distribution(-d) ax = axes(P, 1) return OffsetVector(reverse!(parent(P)), -last(ax):-first(ax)) end maxN = maximum(d) + 1 # if the distribution crosses the zero line, make the "negative frequencies" # disjoint from the positive portion len = minimum(d) < 0 ? 2maxN + 1 : maxN + 1 P = Vector{Float64}(undef, len) T′ = ones(ComplexF64, (len >> 1) + 1) # real-only half-length FFT for d′ in d.die rolls = dierolls(d′) rolls == 0 && continue # skip if contributing no rolls fill!(P, 0.0) F = diebase(d′) if rolls < 0 # by periodicity, negative indices are second half of vector; # see `fftshift()` and `fftfreq()`. P[end-F+1:end] .= 1 / F else # positive rolls P[2:F+1] .= 1 / F end # convolve with cumulative distribution T′ .*= rfft(P) .^ abs(rolls) end P .= irfft(T′, len) if minimum(d) < 0 return OffsetVector(fftshift(P), -maxN:maxN) else return OffsetVector(P, 0:maxN) end end @recipe function f(dice::AbstractDie; xshift = 0.0) l, h = extrema(dice) P = distribution(dice) seriestype := :sticks xlims --> tuple(l ≥ 0 ? 0 : -Inf, h ≤ 0 ? 0 : Inf) # avoid suppressing zero label --> sprint(show, dice) yguide --> "probability" xguide --> "total of rolls" return ((l:h) .+ xshift, parent(P[l:h])) end end