abstract type ContextDist end
abstract type RewardDist end
abstract type BanditPolicy end


mutable struct Trajectory
    t::Int64
    X::Matrix{Float64}
    a::Vector{Int64}
    r::Vector{Float64}
    total_regret::Vector{Float64}
end


function get_X(traj::Trajectory, a::Int64, view::Bool=false)
    if view
        return traj.X[traj.a .== a, :]
    else
        X = @view traj.X[traj.a .== a, :]
        return X
    end
end


function get_r(traj::Trajectory, a::Int64, view::Bool=false)
    if view
        return traj.r[traj.a .== a]
    else
        r = @view traj.r[traj.a .== a]
        return r
    end
end


function Trajectory(d::Int64)
    X = Matrix{Float64}(undef, 0, d)
    a = Vector{Float64}(undef, 0)
    r = Vector{Float64}(undef, 0)
    total_regret = Vector{Float64}(undef, 0)
    return Trajectory(0, X, a, r, total_regret)
end


function add_memory!(traj::Trajectory, T::Int64)
    X1 = Matrix{Float64}(undef, T, size(traj.X, 2))
    traj.X = vcat(traj.X, X1)
    push!(traj.a, zeros(Int64, T)...)
    push!(traj.r, zeros(Float64, T)...)
    push!(traj.total_regret, zeros(Float64, T)...)
end


function extend!(traj::Trajectory, x::Vector{Float64}, a::Int64, r::Float64, total_regret::Float64)
    traj.t += 1
    traj.X[traj.t, :] = x
    traj.a[traj.t] = a
    traj.r[traj.t] = r
    traj.total_regret[traj.t] = total_regret
end


# function extend!(
#     traj::Trajectory,
#     X::Matrix{Float64},
#     a::Vector{Int64},
#     r::Vector{Float64},
#     total_regret::Vector{Float64}
# )
#     n = length(r)
#     traj.X[(traj.t + 1):(traj.t + n), :] = X
#     traj.a[(traj.t + 1):(traj.t + n)] = a
#     traj.r[(traj.t + 1):(traj.t + n)] = r
#     traj.total_regret[(traj.t + 1):(traj.t + n)] = total_regret
#     traj.t += n
# end


function simulate!(
    T::Int64, 
    traj::Trajectory, 
    cdist::ContextDist, 
    rdist::RewardDist, 
    policy::BanditPolicy;
    verbose::Bool = true
)
    add_memory!(traj, T)
    start = traj.t + 1
    finish = traj.t + T
    X = sample(T, cdist)

    for t in start:finish
        if verbose print("\rTime step $t/$finish") end
        a = choose(X[t - start + 1, :], policy)
        r = sample(X[t - start + 1, :], a, rdist)
        regret = expected_regret(X[t - start + 1, :], a, rdist)
        total_regret = t == 1 ? regret : regret + traj.total_regret[t - 1]
        extend!(traj, X[t - start + 1, :], a, r, total_regret)
        update!(policy, X[t - start + 1, :], a, r)
    end
end