mutable struct KNNUCB <: BanditPolicy
    A::Int64
    t::Int64
    traj::Trajectory
    phi::Function
    theta::Float64
end


function KNNUCB(A::Int64, traj::Trajectory, phi::Function, theta::Float64)
    return KNNUCB(A, 0, traj, phi, theta)
end


function sort_nearest(pol::KNNUCB, x::Vector{Float64}, a::Int64)
    dist = Euclidean()
    x1 = reshape(x, (1, :))
    X = get_X(pol.traj, a, false)
    r = get_r(pol.traj, a, false)
    pwdist = pairwise(dist, X, x1; dims = 1)
    pwdist = reshape(pwdist, (:))
    inds = sortperm(pwdist)
    return pwdist[inds], r[inds]
end


function compute_upper_bound(pol::KNNUCB, x::Vector{Float64}, a::Int64)
    t = pol.t
    distances, r_nearest = sort_nearest(pol, x, a)
    uvals = Vector{Float64}(undef, length(r_nearest))
    for k = 1:length(r_nearest)
        radius = distances[k]
        uvals[k] = sqrt(pol.theta * log(t) / k) + pol.phi(t) * radius
    end
    min_uncertainty, min_k = findmin(uvals)
    return mean(r_nearest[1:min_k]) + min_uncertainty
end


function choose(x::Vector{Float64}, pol::KNNUCB)
    if pol.t + 1 <= pol.A
        return pol.t + 1
    end
    u = [compute_upper_bound(pol, x, a) for a in 1:pol.A]
    return findmax(u)[2]
end


function update!(pol::KNNUCB, x::Vector{Float64}, a::Int64, r::Float64)
    pol.t += 1
    if pol.t != pol.traj.t
        throw(ErrorException("Policy out of sync with trajectory"))
    end
end
