using Random
get_image(x::TrainingTest) = x.image
get_action_image(x::TrainingTest) = x.action_image
get_action(x::TrainingTest) = x.action
get_value(x::TrainingTest) = x.value
get_reward(x::TrainingTest) = x.reward
Base.rand(x::TrainingTests) = rand(x.tests)
Base.getindex(x::TrainingTests, idx::Integer) = x.tests[idx]
Base.getindex(x::TrainingTests, idx::Vector{<:Integer}) = x.tests[idx]
Base.length(x::TrainingTests) = length(x.tests)
Base.lastindex(x::TrainingTests) = length(x)
Base.iterate(x::TrainingTests, state=1) = state > length(x) ? nothing : (x[state], state+1)

function initialinferencefn(tests::TrainingTests)
    function fn(n::Network)
        initial_inference(n, tests.images) |> NetworkOutput
    end
end

function recurrentinferencefn(tests::TrainingTests)
    initial_inference_fn = initialinferencefn(tests)
    function fn(n::Network)
        recurrent_inference(
            n, 
            hidden_state(
                initial_inference_fn(n)
            ), 
            tests.action_images
        ) |> NetworkOutput
    end
end

# @todo - Chopping block
#
# function recurrentinferencefn(tests::TrainingTests, initial_output)
#     function fn(n::Network)
#         recurrent_inference(
#             n, 
#             hidden_state(initial_output), 
#             tests.action_images
#         ) |> NetworkOutput
#     end
# end

function getvaluemse(tests::TrainingTests, net_output)
    Flux.mse(
        value(net_output), 
        tests.actual_value
    )
end

function getrewardmse(tests::TrainingTests, net_output)
    Flux.mse(
        reward(net_output), 
        tests.actual_reward
    )
end

function getpolicyfails(tests::Vector{<:TrainingTest}, net_output)
    _policy = policy(net_output)
    greedy_actions = [coord[1] for coord in argmax(_policy, dims=1)[1, :]]
    fails = filter(i->begin
        test = tests[i]
        pred = greedy_actions[i]
        actuals = test.actions
        !in(pred, actuals)
    end, 1:length(tests))
    fails
end

function getpolicyfails(test::TrainingTest, net_output)
    getpolicyfails([test], net_output)
end

function getpolicyacc(tests::TrainingTests, net_output)
    fails = getpolicyfails(tests.tests, net_output)
    total = length(tests.tests)
    acc = PolicyAcc(total, length(fails))
    acc, fails
end