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