using MIPVerify
using HiGHS
using JuMP
using Images

mnist = MIPVerify.read_datasets("MNIST")
n1 = MIPVerify.get_example_network_params("MNIST.n1")
sample_image = MIPVerify.get_image(mnist.test.images, 1);

function view_diff(diff::Array{<:Real, 2})
    n = 1001
    colormap("RdBu", n)[ceil.(Int, (diff .+ 1) ./ 2 .* n)]
end;

d1 = MIPVerify.find_adversarial_example(
    n1, 
    sample_image, 
    10, 
    HiGHS.Optimizer,
    Dict(),
    pp = MIPVerify.LInfNormBoundedPerturbationFamily(0.05),
    norm_order = Inf,
    tightening_algorithm = lp,
)

m = d1[:Model];

num_variables(m)

num_constraints(m, VariableRef, MOI.GreaterThan{Float64})

solve_time(m)

objective_bound(m)

objective_value(m)

d1[:PerturbationFamily]

d1[:TargetIndexes]

d2 = MIPVerify.find_adversarial_example(
    n1, 
    sample_image, 
    8,
    HiGHS.Optimizer,
    Dict("time_limit" => 5.0),
    pp = MIPVerify.LInfNormBoundedPerturbationFamily(0.05),
    norm_order = Inf,
    tightening_algorithm = lp,
    invert_target_selection = true,
)
d2[:TargetIndexes]

d1[:SolveStatus]

d2[:SolveStatus]

perturbation = JuMP.value.(d1[:Perturbation])

view_diff(perturbation[1, :, :, 1])

perturbed_input = JuMP.value.(d1[:PerturbedInput])

colorview(Gray, perturbed_input[1, :, :, 1])

all(perturbed_input == sample_image + perturbation)

output = JuMP.value.(d1[:Output])

perturbed_input |> n1