"""
A Julia implementation of the random Blocksworld state generation algorithm presented in Figure 2 of [1].
Note that this implementation fixes a minor bug in the pseudo-code provided by [1].

[1] Slaney, J. and Thiébaux, S., 2001. Blocks World Revisited. Artificial Intelligence, 125(1-2), pp.119-153.
DOI: https://doi.org/10.1016/S0004-3702(00)00079-5
"""

"Number of states with `N` blocks."
function nstates(N::Int)
    N < 1 ? 1 : nstates(N-1) + (N-1) * (nclear(N-1) + nstates(N-1))
end

"Number of states with `N` blocks where a particular block is clear."
function nclear(N::Int)
    N < 1 ? 1 : nstates(N-1) + (N-1) * nclear(N-1)
end

"Direct implementation of the ratio nclear(N)/nstates(N)."
function pground(N::Int) =
    N < 1 ? 1.0 : ((N-1) * pground(N-1) + 1) / ((N-1) * (pground(N-1) + 1) + 1)
end

"Sample a Bernoulli random variable."
bernoulli(p) = rand() < p

"Sample a random Blocksworld state with `N` blocks."
function rand_bw(N::Int)
    n_ungrounded = N # Number of ungrounded towers
    on = zeros(Int, N) # on[i] = j means i is on j, 0=ungrounded, -1=table
    top = collect(1:N) # Top of each block's tower
    # Repeat until all towers are grounded
    while n_ungrounded > 0
        i_tower = rand(findall(==(0), on)) # Select ungrounded tower
        prob_ground = pground(n_ungrounded)
        if bernoulli(prob_ground) # Place tower on table or another tower
            while true # Repeatedly try to ground tower
                n_ungrounded -= 1
                prob_table = 1 / (1 + n_ungrounded * pground(n_ungrounded))
                if bernoulli(prob_table) # Place on table, and break
                    on[i_tower] = -1
                    break
                else # Place on another ungrounded tower, and loop
                    others = filter(1:N) do j
                        j != i_tower && on[j] == 0
                    end
                    j_tower = rand(others)
                    on[i_tower] = top[j_tower]
                    top[j_tower] = top[i_tower]
                    i_tower = j_tower
                end
            end
        else # Place tower below another ungrounded tower
            n_ungrounded -= 1
            others = filter(1:N) do j
                j != i_tower && on[j] == 0
            end
            j_tower = rand(others)
            on[j_tower] = top[i_tower]
            top[i_tower] = top[j_tower]
        end
    end
    # Identities of which block is on which fully define Blocksworld state
    return on
end