# MIT License
#
# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2020
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
# Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from __future__ import absolute_import, division, print_function, unicode_literals

import logging
import unittest

import numpy as np

from art.attacks.evasion.wasserstein import Wasserstein
from art.estimators.estimator import BaseEstimator, LossGradientsMixin
from art.estimators.classification.classifier import ClassifierMixin
from art.utils import get_labels_np_array

from tests.utils import TestBase
from tests.utils import get_image_classifier_tf
from tests.attacks.utils import backend_test_classifier_type_check_fail

logger = logging.getLogger(__name__)


class TestWasserstein(TestBase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()

        cls.n_train = 10
        cls.n_test = 10
        cls.x_train_mnist = cls.x_train_mnist[0 : cls.n_train]
        cls.y_train_mnist = cls.y_train_mnist[0 : cls.n_train]
        cls.x_test_mnist = cls.x_test_mnist[0 : cls.n_test]
        cls.y_test_mnist = cls.y_test_mnist[0 : cls.n_test]

    # def test_keras_mnist(self):
    #     classifier = get_image_classifier_kr()
    #
    #     scores = classifier._model.evaluate(self.x_train_mnist, self.y_train_mnist)
    #     logger.info("[Keras, MNIST] Accuracy on training set: %.2f%%", scores[1] * 100)
    #     scores = classifier._model.evaluate(self.x_test_mnist, self.y_test_mnist)
    #     logger.info("[Keras, MNIST] Accuracy on test set: %.2f%%", scores[1] * 100)
    #
    #     self._test_backend_mnist(
    #         classifier, self.x_train_mnist, self.y_train_mnist, self.x_test_mnist, self.y_test_mnist
    #     )

    def test_tensorflow_mnist(self):
        classifier, sess = get_image_classifier_tf()

        scores = get_labels_np_array(classifier.predict(self.x_train_mnist))
        acc = np.sum(np.argmax(scores, axis=1) == np.argmax(self.y_train_mnist, axis=1)) / self.y_train_mnist.shape[0]
        logger.info("[TF, MNIST] Accuracy on training set: %.2f%%", acc * 100)

        scores = get_labels_np_array(classifier.predict(self.x_test_mnist))
        acc = np.sum(np.argmax(scores, axis=1) == np.argmax(self.y_test_mnist, axis=1)) / self.y_test_mnist.shape[0]
        logger.info("[TF, MNIST] Accuracy on test set: %.2f%%", acc * 100)

        self._test_backend_mnist(
            classifier, self.x_train_mnist, self.y_train_mnist, self.x_test_mnist, self.y_test_mnist
        )

    # def test_pytorch_mnist(self):
    #     x_train_mnist = np.swapaxes(self.x_train_mnist, 1, 3).astype(np.float32)
    #     x_test_mnist = np.swapaxes(self.x_test_mnist, 1, 3).astype(np.float32)
    #     classifier = get_image_classifier_pt()
    #
    #     scores = get_labels_np_array(classifier.predict(x_train_mnist))
    #     acc = np.sum(np.argmax(scores, axis=1) == np.argmax(self.y_train_mnist, axis=1)) / self.y_train_mnist.shape[0]
    #     logger.info("[PyTorch, MNIST] Accuracy on training set: %.2f%%", acc * 100)
    #
    #     scores = get_labels_np_array(classifier.predict(x_test_mnist))
    #     acc = np.sum(np.argmax(scores, axis=1) == np.argmax(self.y_test_mnist, axis=1)) / self.y_test_mnist.shape[0]
    #     logger.info("[PyTorch, MNIST] Accuracy on test set: %.2f%%", acc * 100)
    #
    #     self._test_backend_mnist(classifier, x_train_mnist, self.y_train_mnist, x_test_mnist, self.y_test_mnist)

    def _test_backend_mnist(self, classifier, x_train, y_train, x_test, y_test):

        base_success_rate = 0.1
        num_iter = 5
        regularization = 100
        batch_size = 5
        eps = 0.3

        # Test Wasserstein with wasserstein ball and wasserstein norm
        attack = Wasserstein(
            classifier,
            regularization=regularization,
            max_iter=num_iter,
            conjugate_sinkhorn_max_iter=num_iter,
            projected_sinkhorn_max_iter=num_iter,
            norm="wasserstein",
            ball="wasserstein",
            targeted=False,
            p=2,
            eps_iter=2,
            eps_factor=1.05,
            eps=eps,
            eps_step=0.1,
            kernel_size=5,
            batch_size=batch_size,
        )

        x_train_adv = attack.generate(x_train)
        x_test_adv = attack.generate(x_test)

        self.assertFalse((x_train_adv == x_train).all())
        self.assertFalse((x_test_adv == x_test).all())

        train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)

        train_success_rate = (
            np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1)) / y_train.shape[0]
        )
        self.assertGreaterEqual(train_success_rate, base_success_rate)

        test_success_rate = (
            np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        )
        self.assertGreaterEqual(test_success_rate, base_success_rate)

        # # Test Wasserstein with wasserstein ball and l2 norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="2",
        #     ball="wasserstein",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, base_success_rate)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with wasserstein ball and inf norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="inf",
        #     ball="wasserstein",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, base_success_rate)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with wasserstein ball and l1 norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="1",
        #     ball="wasserstein",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, base_success_rate)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with l2 ball and wasserstein norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="wasserstein",
        #     ball="2",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.05,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, 0)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with l1 ball and wasserstein norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="wasserstein",
        #     ball="1",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, 0)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with inf ball and Wasserstein norm
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="wasserstein",
        #     ball="inf",
        #     targeted=False,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, base_success_rate)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertGreaterEqual(test_success_rate, base_success_rate)
        #
        # # Test Wasserstein with targeted attack
        # master_seed(1234)
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="wasserstein",
        #     ball="wasserstein",
        #     targeted=True,
        #     p=2,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=5,
        #     batch_size=batch_size,
        # )
        #
        # train_y_rand = random_targets(y_train, nb_classes=10)
        # test_y_rand = random_targets(y_test, nb_classes=10)
        #
        # x_train_adv = attack.generate(x_train, train_y_rand)
        # x_test_adv = attack.generate(x_test, test_y_rand)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) == np.argmax(train_y_rand, axis=1)) / y_train.shape[0]
        # )
        # self.assertGreaterEqual(train_success_rate, base_success_rate)
        #
        # test_success_rate = np.sum(np.argmax(test_y_pred, axis=1) == np.argmax(test_y_rand, axis=1)) / y_test.shape[0]
        # self.assertGreaterEqual(test_success_rate, 0)
        #
        # # Test Wasserstein with p-wasserstein=1 and kernel_size=3
        # attack = Wasserstein(
        #     classifier,
        #     regularization=regularization,
        #     max_iter=num_iter,
        #     conjugate_sinkhorn_max_iter=num_iter,
        #     projected_sinkhorn_max_iter=num_iter,
        #     norm="wasserstein",
        #     ball="wasserstein",
        #     targeted=False,
        #     p=1,
        #     eps_iter=2,
        #     eps_factor=1.05,
        #     eps=eps,
        #     eps_step=0.1,
        #     kernel_size=3,
        #     batch_size=batch_size,
        # )
        #
        # x_train_adv = attack.generate(x_train)
        # x_test_adv = attack.generate(x_test)
        #
        # train_y_pred = get_labels_np_array(classifier.predict(x_train_adv)).astype(float)
        # test_y_pred = get_labels_np_array(classifier.predict(x_test_adv)).astype(float)
        #
        # train_success_rate = (
        #     np.sum(np.argmax(train_y_pred, axis=1) != np.argmax(classifier.predict(x_train), axis=1))
        #     / y_train.shape[0]
        # )
        # self.assertTrue(train_success_rate >= base_success_rate)
        #
        # test_success_rate = (
        #     np.sum(np.argmax(test_y_pred, axis=1) != np.argmax(classifier.predict(x_test), axis=1)) / y_test.shape[0]
        # )
        # self.assertTrue(test_success_rate >= base_success_rate)

    def test_classifier_type_check_fail(self):
        backend_test_classifier_type_check_fail(Wasserstein, (BaseEstimator, LossGradientsMixin, ClassifierMixin))

    def test_unsquared_images(self):
        from art.estimators.estimator import (
            BaseEstimator,
            LossGradientsMixin,
            NeuralNetworkMixin,
        )

        from art.estimators.classification.classifier import (
            ClassGradientsMixin,
            ClassifierMixin,
        )

        class DummyClassifier(
            ClassGradientsMixin, ClassifierMixin, NeuralNetworkMixin, LossGradientsMixin, BaseEstimator
        ):
            def __init__(self):
                super(DummyClassifier, self).__init__()
                self._nb_classes = 10
                self._channels_first = True

            def class_gradient(self):
                return None

            def fit(self):
                pass

            def loss_gradient(self, x, y):
                return np.random.normal(size=(1, 3, 33, 32))

            def predict(self, x, batch_size=1):
                return np.array([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0]])

            def get_activations(self):
                return None

            def save(self):
                pass

            def loss(self, x, y, **kwargs):
                pass

            def set_learning_phase(self):
                pass

        classifier = DummyClassifier()
        attack = Wasserstein(
            classifier,
            regularization=1,
            kernel_size=3,
            max_iter=1,
            conjugate_sinkhorn_max_iter=10,
            projected_sinkhorn_max_iter=10,
        )

        x = np.random.normal(size=(1, 3, 33, 32))
        x_adv = attack.generate(x)

        self.assertTrue(x_adv.shape == x.shape)


if __name__ == "__main__":
    unittest.main()
