"""
Variables from a posterior distribution (typically generated by MCMC).

`keys` and `vars` are sorted by `keys`. Access outside this library
should always use keys.
"""
@auto_hash_equals immutable PosteriorDraws
    len::Int
    keys::Vector{Symbol}
    vars::Vector{Any}
    function PosteriorDraws(len::Int, keys::Vector{Symbol}, vars::Vector)
        length(keys) == length(vars) ||
            error(ArgumentError("Non-conforming keys and vars."))
        all(length(var) == len for var in vars) ||
            error("Non-conformable lengths")
        vars = Any[isa(v, Vector) ? v : convert(Vector, v) for v in vars]
        p = sortperm(keys)
        sorted_keys = keys[p]
        all(a < b for (a,b) in partition(sorted_keys, 2, 1)) ||
            error(ArgumentError("Duplicate or non-sorted keys."))
        new(len, sorted_keys, vars[p])
    end
end

# function PosteriorDraws(len::Int, keys, vars)
#     PosteriorDraws(len, convert(Vector{Symbol}, keys), convert(Vector{Any}, vars))
# end

PosteriorDraws(len::Int) = PosteriorDraws(len, Symbol[], Any[])

function PosteriorDraws(keys::Vector{Symbol}, vars::Vector)
    length(vars) ≥ 1 ||
        error(ArgumentError("Need at least one variable to determine length."))
    PosteriorDraws(length(vars[1]), keys, vars)
end

function PosteriorDraws{T}(key_var_pairs::Pair{Symbol, T}...)
    PosteriorDraws(first.([key_var_pairs...]), last.([key_var_pairs...]))
end

function show(io::IO, pd::PosteriorDraws)
    println(io, "PosteriorDraws with $(pd.len) observations")
    for (key,var) in zip(pd.keys, pd.vars)
        println(io, "    $(key) => $(eltype(var))")
    end
end

keys(pd::PosteriorDraws) = pd.keys

size(pd::PosteriorDraws) = (length(pd.keys), pd.len)
size(pd::PosteriorDraws, index) = size(pd)[index]

"Lookup `key` in `pd`, returning the corresponding indexes."
function _key2index(pd::PosteriorDraws, key::Symbol)
    index = searchsortedfirst(pd.keys, key)
    index > length(pd.keys) && error(ArgumentError("Variable $(key) not found"))
    index
end

function _key2index(pd::PosteriorDraws, keys::Vector{Symbol})
    [_key2index(pd, key) for key in keys]
end

_key2index(pd::PosteriorDraws, ::Colon) = indices(pd.keys,1)

function getindex(pd::PosteriorDraws, keys)
    keyinds = _key2index(pd, keys)
    if isa(keyinds, Int)
        pd.vars[keyinds]
    else
        PosteriorDraws(pd.len, pd.keys[keyinds], pd.vars[keyinds])
    end
end

function getindex(pd::PosteriorDraws, keys, drawinds)
    keyinds = _key2index(pd, keys)
    if isa(keyinds, Int)
        pd.vars[keyinds][drawinds]
    else
        PosteriorDraws(pd.keys[keyinds],
                       [var[drawinds] for var in pd.vars[keyinds]])
    end
end

ndims(pd::PosteriorDraws) = 2
