@with_kw struct SimpleLightDark <: POMDP{Int,Int,Float64}
    discount::Float64       = 0.95
    correct_r::Float64      = 100.0
    incorrect_r::Float64    = -100.0
    light_loc::Int          = 10
    radius::Int             = 60
end
discount(p::SimpleLightDark) = p.discount
isterminal(p::SimpleLightDark, s::Number) = !(s in -p.radius:p.radius)

const ACTIONS = [-10, -1, 0, 1, 10]
actions(p::SimpleLightDark) = ACTIONS
n_actions(p::SimpleLightDark) = length(actions(p))
const ACTION_INDS = Dict(a=>i for (i,a) in enumerate(actions(SimpleLightDark())))
actionindex(p::SimpleLightDark, a::Int) = ACTION_INDS[a]

states(p::SimpleLightDark) = -p.radius:p.radius + 1
n_states(p::SimpleLightDark) = length(states(p))
stateindex(p::SimpleLightDark, s::Int) = s+p.radius+1

function transition(p::SimpleLightDark, s::Int, a::Int) 
    if a == 0
        return SparseCat(SVector(p.radius+1), SVector(1.0))
    else
        return SparseCat(SVector(clamp(s+a, -p.radius, p.radius)), SVector(1.0))
    end
end

observation(p::SimpleLightDark, sp) = Normal(sp, abs(sp - p.light_loc) + 0.0001)

function reward(p::SimpleLightDark, s, a)
    if a == 0
        return s == 0 ? p.correct_r : p.incorrect_r
    else
        return -1.0
    end
end

function initialstate(p::SimpleLightDark)
    ps = ones(2*div(p.radius,2)+1)
    ps /= length(ps)
    return SparseCat(div(-p.radius,2):div(p.radius,2), ps)
end

@with_kw struct DSimpleLightDark <: POMDP{Int, Int, Int}
    sld::SimpleLightDark       = SimpleLightDark()
    binsize::Float64           = 1.0
    _d_disc::Vector{SparseCat} = SparseCat[get_discrete_obs_dist(sld, sp, binsize) for sp in -sld.radius:sld.radius]
end

sld(s::SimpleLightDark) = s
sld(s::DSimpleLightDark) = s.sld

function get_discrete_obs_dist(p::SimpleLightDark, sp::Int, binsize::Float64)
    dist = observation(p, sp)
    # discretize observations
    interval_start = 0.0
    discrete_dist = Vector{Float64}(undef, floor(Int, div(p.radius*2, binsize)) + 2)
    @assert div(p.radius*2, binsize) + 2 == length(range(-p.radius, step=binsize, stop=p.radius)) + 1
    for (oi, o) in enumerate(range(-p.radius, step=binsize, stop=p.radius))
        interval_end = cdf(dist, o)
        discrete_dist[oi] = interval_end - interval_start
        interval_start = interval_end
    end
    discrete_dist[end] = 1.0 - interval_start
    return SparseCat(1:length(discrete_dist), discrete_dist)
end

observations(p::DSimpleLightDark) = 1:(div(p.sld.radius * 2, p.binsize) + 2)
n_observations(p::DSimpleLightDark) = length(observations(p))
obsindex(p::DSimpleLightDark, o::Int) = o

actions(p::DSimpleLightDark) = actions(p.sld)
n_actions(p::DSimpleLightDark) = n_actions(p.sld)
actionindex(p::DSimpleLightDark, a::Int) = actionindex(p.sld, a)

states(p::DSimpleLightDark) = states(p.sld)
n_states(p::DSimpleLightDark) = n_states(p.sld)
stateindex(p::DSimpleLightDark, s::Int) = stateindex(p.sld, s)

observation(p::DSimpleLightDark, sp::Int) = p._d_disc[stateindex(p, sp)]
transition(p::DSimpleLightDark, s::Int, a::Int) = transition(p.sld, s, a)
reward(p::DSimpleLightDark, s, a) = reward(p.sld, s, a)
initialstate(p::DSimpleLightDark) = initialstate(p.sld)

discount(p::DSimpleLightDark) = discount(p.sld)
isterminal(p::DSimpleLightDark, s::Number) = isterminal(p.sld, s)