def test_mnist_certified_loss(art_warning, fix_get_mnist_data):
    """
    Check the certified losses with interval_loss_cce, max_logit_loss, and make sure that we give a lower
    bound compared to PGD.
    """
    bound = 0.05
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    ptc = get_image_classifier_pt(from_logits=True, use_maxpool=False)
    import torch.optim as optim

    optimizer = optim.Adam(ptc.model.parameters(), lr=1e-4)
    zonotope_model = PytorchDeepZ(
        model=ptc.model.to(device),
        optimizer=optimizer,
        clip_values=(0, 1),
        loss=torch.nn.CrossEntropyLoss(),
        input_shape=(1, 28, 28),
        nb_classes=10,
    )

    pgd_params = {"eps": 0.25, "eps_step": 0.05, "max_iter": 20, "num_random_init": 0, "batch_size": 64}

    trainer = AdversarialTrainerCertifiedPytorch(zonotope_model, pgd_params=pgd_params, batch_size=10, bound=bound)

    np.random.seed(123)
    random.seed(123)
    torch.manual_seed(123)
    concrete_correct = 0
    samples_x, samples_eps = trainer.predict_zonotopes(np.copy(fix_get_mnist_data[0]), bound)
    try:
        certified_loss = 0
        for i, (x, y) in enumerate(zip(fix_get_mnist_data[0], fix_get_mnist_data[1])):
            x = x.astype("float32")
            pred_sample = np.copy(x)
            pred_sample = np.expand_dims(pred_sample, axis=0)
            zonotope_model.model.set_forward_mode("concrete")
            prediction = trainer._classifier.model.forward(torch.from_numpy(pred_sample.astype("float32")).to(device))
            if np.argmax(prediction.cpu().detach().numpy()) == y:
                concrete_correct += 1

            eps_bound = np.eye(784) * bound
            zonotope_model.model.set_forward_mode("abstract")
            data_sample_processed, eps_bound = zonotope_model.pre_process(cent=x, eps=eps_bound)
            data_sample_processed = np.expand_dims(data_sample_processed, axis=0)
            bias, eps_bound = trainer._classifier.model.forward(eps=eps_bound, cent=data_sample_processed)
            bias = torch.unsqueeze(bias, dim=0)

            assert np.allclose(samples_x[i], bias.cpu().detach().numpy())
            assert np.allclose(samples_eps[i], eps_bound.cpu().detach().numpy())

            certified_loss += trainer._classifier.interval_loss_cce(
                prediction=torch.cat((bias, eps_bound)),
                target=torch.from_numpy(np.expand_dims(y, axis=0)).to(device),
            )

        assert round(float(certified_loss.cpu().detach().numpy()), 4) == 39.7759

        samples_certified = 0

        certified_loss = 0
        for x, y in zip(fix_get_mnist_data[0], fix_get_mnist_data[1]):
            x = x.astype("float32")
            pred_sample = np.copy(x)
            pred_sample = np.expand_dims(pred_sample, axis=0)
            zonotope_model.model.set_forward_mode("concrete")
            prediction = trainer._classifier.model.forward(torch.from_numpy(pred_sample.astype("float32")).to(device))
            if np.argmax(prediction.cpu().detach().numpy()) == y:
                concrete_correct += 1

            eps_bound = np.eye(784) * bound
            zonotope_model.model.set_forward_mode("abstract")
            data_sample_processed, eps_bound = zonotope_model.pre_process(cent=np.copy(x), eps=eps_bound)
            data_sample_processed = np.expand_dims(data_sample_processed, axis=0)
            bias, eps_bound = trainer._classifier.model.forward(eps=eps_bound, cent=data_sample_processed)
            bias = torch.unsqueeze(bias, dim=0)

            sample_loss = trainer._classifier.max_logit_loss(
                prediction=torch.cat((bias, eps_bound)),
                target=torch.from_numpy(np.expand_dims(y, axis=0)).to(device),
            )

            data_sample_processed, eps_bound = zonotope_model.pre_process(cent=np.copy(x), eps=np.eye(784) * bound)
            data_sample_processed = np.expand_dims(data_sample_processed, axis=0)
            is_certified = zonotope_model.certify(
                eps=eps_bound, cent=data_sample_processed, prediction=np.argmax(prediction.cpu().detach().numpy())
            )

            if (
                is_certified and int(np.argmax(prediction.cpu().detach().numpy())) == y
            ):  # if it is classified correctly AND is certified
                assert sample_loss <= 0.0  # ...the max_logit_loss gives negatve.
                samples_certified += 1

            certified_loss += sample_loss

        assert float(certified_loss.cpu().detach().numpy()) == pytest.approx(-309.2724, abs=0.001)
        assert samples_certified == 94

        # empirically check that PGD does not give a lower acc
        zonotope_model.model.set_forward_mode("concrete")
        attack = ProjectedGradientDescent(
            estimator=trainer._classifier,
            eps=bound,
            eps_step=bound / 20,
            max_iter=30,
            num_random_init=0,
        )
        i_batch = attack.generate(fix_get_mnist_data[0], y=fix_get_mnist_data[1])
        preds = trainer._classifier.model.forward(i_batch)
        preds = preds.cpu().detach().numpy()
        acc = np.sum(np.argmax(preds, axis=1) == fix_get_mnist_data[1]) / len(fix_get_mnist_data[1])
        assert acc * 100 >= samples_certified

    except ARTTestException as e:
        art_warning(e)

    Check the certified losses with interval_loss_cce, max_logit_loss, and make sure that we give a lower
    bound compared to PGD.
    
tests/defences/trainer/test_certified_adversarial_trainer.py