using MIPVerify
using Gurobi
using MAT

param_dict = Base.download("https://github.com/vtjeng/MIPVerify_data/raw/master/weights/mnist/n2.mat") |> matread

fc1_weight = param_dict["fc1/weight"]

fc1_bias = param_dict["fc1/bias"]

fc1_manual = Linear(fc1_weight, squeeze(fc1_bias, 1))

fc1 = get_matrix_params(param_dict, "fc1", (784, 24))

fc1_manual == fc1

fc2 = get_matrix_params(param_dict, "fc2", (24, 24))

fc3 = get_matrix_params(param_dict, "fc3", (24, 24))

logits = get_matrix_params(param_dict, "logits", (24, 10))

nn = Sequential([
        Flatten(4),
        fc1,
        ReLU(),
        fc2,
        ReLU(),
        fc3,
        ReLU(),
        logits
    ], "MNIST.n2")

mnist = read_datasets("MNIST")
MIPVerify.frac_correct(nn, mnist.test, 10000)

sample_image = MIPVerify.get_image(mnist.test.images, 1);

MIPVerify.find_adversarial_example(nn, sample_image, 4, GurobiSolver())