# coding=utf-8
# Copyright 2019 The Interval Bound Propagation Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Evaluates a verifiable model on Mnist or CIFAR-10."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import logging
import interval_bound_propagation as ibp
import tensorflow.compat.v1 as tf


def layers(model_size):
    """Returns the layer specification for a given model name."""
    if model_size == "tiny":
        return (("linear", 100), ("activation", "relu"))
    elif model_size == "small":
        return (
            ("conv2d", (4, 4), 16, "VALID", 2),
            ("activation", "relu"),
            ("conv2d", (4, 4), 32, "VALID", 1),
            ("activation", "relu"),
            ("linear", 100),
            ("activation", "relu"),
        )
    elif model_size == "medium":
        return (
            ("conv2d", (3, 3), 32, "VALID", 1),
            ("activation", "relu"),
            ("conv2d", (4, 4), 32, "VALID", 2),
            ("activation", "relu"),
            ("conv2d", (3, 3), 64, "VALID", 1),
            ("activation", "relu"),
            ("conv2d", (4, 4), 64, "VALID", 2),
            ("activation", "relu"),
            ("linear", 512),
            ("activation", "relu"),
            ("linear", 512),
            ("activation", "relu"),
        )
    elif model_size == "large_200":
        # Some old large checkpoints have 200 hidden neurons in the last linear
        # layer.
        return (
            ("conv2d", (3, 3), 64, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 64, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 2),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 1),
            ("activation", "relu"),
            ("linear", 200),
            ("activation", "relu"),
        )
    elif model_size == "large":
        return (
            ("conv2d", (3, 3), 64, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 64, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 2),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 1),
            ("activation", "relu"),
            ("conv2d", (3, 3), 128, "SAME", 1),
            ("activation", "relu"),
            ("linear", 512),
            ("activation", "relu"),
        )
    else:
        raise ValueError('Unknown model: "{}"'.format(model_size))


def show_metrics(metric_values, bound_method="ibp"):
    if bound_method == "crown-ibp":
        verified_accuracy = metric_values.crown_ibp_verified_accuracy
    else:
        verified_accuracy = metric_values.verified_accuracy
    print(
        "nominal accuracy = {:.2f}%, "
        "verified accuracy = {:.2f}%, "
        "accuracy under PGD attack = {:.2f}%".format(
            metric_values.nominal_accuracy * 100.0,
            verified_accuracy * 100.0,
            metric_values.attack_accuracy * 100.0,
        )
    )


# Test using while loop.
def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack):
    """Returns the test metrics."""
    num_test_batches = len(data_test[0]) // batch_size
    assert (
        len(data_test[0]) % batch_size == 0
    ), "Test data is not a multiple of batch size."

    def cond(i, *unused_args):
        return i < num_test_batches

    def body(i, metrics):
        """Compute the sum of all metrics."""
        test_data = ibp.build_dataset(data_test, batch_size=batch_size, sequential=True)
        predictor(test_data.image, override=True, is_training=False)
        input_interval_bounds = ibp.IntervalBounds(
            tf.maximum(test_data.image - cEPSILON, input_bounds[0]),
            tf.minimum(test_data.image + cEPSILON, input_bounds[1]),
        )
        predictor.propagate_bounds(input_interval_bounds)
        test_specification = ibp.ClassificationSpecification(
            test_data.label, num_classes
        )
        test_attack = attack_builder(
            predictor,
            test_specification,
            cEPSILON,
            input_bounds=input_bounds,
            optimizer_builder=ibp.UnrolledAdam,
        )

        # Use CROWN-IBP bound or IBP bound.
        if cBOUND_METHOD == "crown-ibp":
            test_losses = ibp.crown.Losses(
                predictor,
                test_specification,
                test_attack,
                use_crown_ibp=True,
                crown_bound_schedule=tf.constant(1.0),
            )
        else:
            test_losses = ibp.Losses(predictor, test_specification, test_attack)

        test_losses(test_data.label)
        new_metrics = []
        for m, n in zip(metrics, test_losses.scalar_metrics):
            new_metrics.append(m + n)
        return i + 1, new_metrics

    if cBOUND_METHOD == "crown-ibp":
        metrics = ibp.crown.ScalarMetrics
    else:
        metrics = ibp.ScalarMetrics
    total_count = tf.constant(0, dtype=tf.int32)
    total_metrics = [
        tf.constant(0, dtype=tf.float32) for _ in range(len(metrics._fields))
    ]
    total_count, total_metrics = tf.while_loop(
        cond,
        body,
        loop_vars=[total_count, total_metrics],
        back_prop=False,
        parallel_iterations=1,
    )
    total_count = tf.cast(total_count, tf.float32)
    test_metrics = []
    for m in total_metrics:
        test_metrics.append(m / total_count)
    return metrics(*test_metrics)


def main():
    test_metrics = get_test_metrics(200, ibp.UntargetedPGDAttack)

    # Prepare to load the pretrained-model.
    saver = tf.compat.v1.train.Saver(original_predictor.get_variables())

    # Run everything.
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    with tf.train.SingularMonitoredSession(config=tf_config) as sess:
        logging.info('Restoring from checkpoint "%s".', checkpoint_path)
        saver.restore(sess, checkpoint_path)
        logging.info("Evaluating at epsilon = %f.", cEPSILON)
        metric_values = sess.run(test_metrics)
        show_metrics(metric_values, cBOUND_METHOD)


if __name__ == "__main__":
    print("TensorFlow version {}".format(tf.__version__))
    # print("Sonnet version {}".format(snt.__version__))

    # python sg_eval.py --model_dir './Models/crown-ibp/cifar_2-255_large/' --epsilon 0.00784313725490196 --bound_method=crown-ibp

    cDATASET = "cifar10"
    cMODEL = "large"
    cEPSILON = 0.00784313725490196
    cBOUND_METHOD = "ibp"

    cMODEL_DIR = ".//Models//ibp//cifar_2-255_large_200//"
    checkpoint_path = tf.train.latest_checkpoint(cMODEL_DIR)
    if checkpoint_path is None:
        raise OSError("Cannot find a valid checkpoint in {}.".format(cMODEL_DIR))

    # Dataset.
    input_bounds = (0.0, 1.0)
    num_classes = 10
    if cDATASET == "mnist":
        data_train, data_test = tf.keras.datasets.mnist.load_data()
    else:
        assert cDATASET == "cifar10", 'Unknown dataset "{}"'.format(cDATASET)
        data_train, data_test = tf.keras.datasets.cifar10.load_data()
        data_train = (data_train[0], data_train[1].flatten())
        data_test = (data_test[0], data_test[1].flatten())

    # Base predictor network.
    original_predictor = ibp.DNN(num_classes, layers(cMODEL))
    predictor = original_predictor
    if cDATASET == "cifar10":
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
        predictor = ibp.add_image_normalization(original_predictor, mean, std)
    if cBOUND_METHOD == "crown-ibp":
        predictor = ibp.crown.VerifiableModelWrapper(predictor)
    else:
        predictor = ibp.VerifiableModelWrapper(predictor)

    print(predictor)

    main()