# Copyright 2018 Google LLC
#
# 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
#
#     https://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.

r"""Run experiments with NNGP Kernel.

Usage:

python run_experiments.py \
      --num_train=100 \
      --num_eval=1000 \
      --hparams='nonlinearity=relu,depth=10,weight_var=1.79,bias_var=0.83' \
      --n_gauss=501 --n_var=501 --n_corr=500 --max_gauss=10

"""

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

import csv
import os.path
import time

import numpy as np
import tensorflow as tf

import gpr
import load_dataset
import nngp

tf.logging.set_verbosity(tf.logging.INFO)

flags = tf.app.flags
FLAGS = flags.FLAGS

flags.DEFINE_string('hparams', '',
                    'Comma separated list of name=value hyperparameter pairs to'
                    'override the default setting.')
flags.DEFINE_string('experiment_dir', '/tmp/nngp',
                    'Directory to put the experiment results.')
flags.DEFINE_string('grid_path', './grid_data',
                    'Directory to put or find the training data.')
flags.DEFINE_integer('num_train', 1000, 'Number of training data.')
flags.DEFINE_integer('num_eval', 1000,
                     'Number of evaluation data. Use 10_000 for full eval')
flags.DEFINE_integer('seed', 1234, 'Random number seed for data shuffling')
flags.DEFINE_boolean('save_kernel', False, 'Save Kernel do disk')
flags.DEFINE_string('dataset', 'mnist',
                    'Which dataset to use ["mnist"]')
flags.DEFINE_boolean('use_fixed_point_norm', False,
                     'Normalize input variance to fixed point variance')

flags.DEFINE_integer('n_gauss', 501,
                     'Number of gaussian integration grid. Choose odd integer.')
flags.DEFINE_integer('n_var', 501,
                     'Number of variance grid points.')
flags.DEFINE_integer('n_corr', 500,
                     'Number of correlation grid points.')
flags.DEFINE_integer('max_var', 100,
                     'Max value for variance grid.')
flags.DEFINE_integer('max_gauss', 10,
                     'Range for gaussian integration.')


def set_default_hparams():
    return tf.contrib.training.HParams(
        nonlinearity='tanh', weight_var=1.3, bias_var=0.2, depth=2)


def do_eval(sess, model, x_data, y_data, save_pred=False):
    """Run evaluation."""

    gp_prediction, stability_eps = model.predict(x_data, sess)

    pred_1 = np.argmax(gp_prediction, axis=1)
    accuracy = np.sum(pred_1 == np.argmax(y_data, axis=1)) / float(len(y_data))
    mse = np.mean(np.mean((gp_prediction - y_data) ** 2, axis=1))
    pred_norm = np.mean(np.linalg.norm(gp_prediction, axis=1))
    tf.logging.info('Accuracy: %.4f' % accuracy)
    tf.logging.info('MSE: %.8f' % mse)

    if save_pred:
        with tf.gfile.Open(
                os.path.join(FLAGS.experiment_dir, 'gp_prediction_stats.npy'),
                'w') as f:
            np.save(f, gp_prediction)

    return accuracy, mse, pred_norm, stability_eps


def run_nngp_eval(hparams, run_dir):
    """Runs experiments."""

    tf.gfile.MakeDirs(run_dir)
    # Write hparams to experiment directory.
    with tf.gfile.GFile(run_dir + '/hparams', mode='w') as f:
        f.write(hparams.to_proto().SerializeToString())

    tf.logging.info('Starting job.')
    tf.logging.info('Hyperparameters')
    tf.logging.info('---------------------')
    tf.logging.info(hparams)
    tf.logging.info('---------------------')
    tf.logging.info('Loading data')

    # Get the sets of images and labels for training, validation, and
    # # test on dataset.
    if FLAGS.dataset == 'mnist':
        (train_image, train_label, valid_image, valid_label, test_image,
         test_label) = load_dataset.load_mnist(
            num_train=FLAGS.num_train,
            mean_subtraction=True,
            random_roated_labels=False)
    else:
        raise NotImplementedError

    tf.logging.info('Building Model')

    if hparams.nonlinearity == 'tanh':
        nonlin_fn = tf.tanh
    elif hparams.nonlinearity == 'relu':
        nonlin_fn = tf.nn.relu
    else:
        raise NotImplementedError

    with tf.Session() as sess:
        # Construct NNGP kernel
        nngp_kernel = nngp.NNGPKernel(
            depth=hparams.depth,
            weight_var=hparams.weight_var,
            bias_var=hparams.bias_var,
            nonlin_fn=nonlin_fn,
            grid_path=FLAGS.grid_path,
            n_gauss=FLAGS.n_gauss,
            n_var=FLAGS.n_var,
            n_corr=FLAGS.n_corr,
            max_gauss=FLAGS.max_gauss,
            max_var=FLAGS.max_var,
            use_fixed_point_norm=FLAGS.use_fixed_point_norm)

        # Construct Gaussian Process Regression model
        model = gpr.GaussianProcessRegression(
            train_image, train_label, kern=nngp_kernel)

        start_time = time.time()
        tf.logging.info('Training')

        # For large number of training points, we do not evaluate on full set to
        # save on training evaluation time.
        if FLAGS.num_train <= 5000:
            acc_train, mse_train, norm_train, final_eps = do_eval(
                sess, model, train_image[:FLAGS.num_eval],
                train_label[:FLAGS.num_eval])
            tf.logging.info('Evaluation of training set (%d examples) took '
                            '%.3f secs' % (
                                min(FLAGS.num_train, FLAGS.num_eval),
                                time.time() - start_time))
        else:
            acc_train, mse_train, norm_train, final_eps = do_eval(
                sess, model, train_image[:1000], train_label[:1000])
            tf.logging.info('Evaluation of training set (%d examples) took '
                            '%.3f secs' % (1000, time.time() - start_time))

        start_time = time.time()
        tf.logging.info('Validation')
        acc_valid, mse_valid, norm_valid, _ = do_eval(
            sess, model, valid_image[:FLAGS.num_eval],
            valid_label[:FLAGS.num_eval])
        tf.logging.info('Evaluation of valid set (%d examples) took %.3f secs' % (
            FLAGS.num_eval, time.time() - start_time))

        start_time = time.time()
        tf.logging.info('Test')
        acc_test, mse_test, norm_test, _ = do_eval(
            sess,
            model,
            test_image[:FLAGS.num_eval],
            test_label[:FLAGS.num_eval],
            save_pred=False)
        tf.logging.info('Evaluation of test set (%d examples) took %.3f secs' % (
            FLAGS.num_eval, time.time() - start_time))

    metrics = {
        'train_acc': float(acc_train),
        'train_mse': float(mse_train),
        'train_norm': float(norm_train),
        'valid_acc': float(acc_valid),
        'valid_mse': float(mse_valid),
        'valid_norm': float(norm_valid),
        'test_acc': float(acc_test),
        'test_mse': float(mse_test),
        'test_norm': float(norm_test),
        'stability_eps': float(final_eps),
    }

    record_results = [
        FLAGS.num_train, hparams.nonlinearity, hparams.weight_var,
        hparams.bias_var, hparams.depth, acc_train, acc_valid, acc_test,
        mse_train, mse_valid, mse_test, final_eps
    ]
    if nngp_kernel.use_fixed_point_norm:
        metrics['var_fixed_point'] = float(nngp_kernel.var_fixed_point_np[0])
        record_results.append(nngp_kernel.var_fixed_point_np[0])

    # Store data
    result_file = os.path.join(run_dir, 'results.csv')
    with tf.gfile.Open(result_file, 'a') as f:
        filewriter = csv.writer(f)
        filewriter.writerow(record_results)

    return metrics


def main(argv):
    del argv  # Unused
    hparams = set_default_hparams().parse(FLAGS.hparams)
    run_nngp_eval(hparams, FLAGS.experiment_dir)


if __name__ == '__main__':
    tf.app.run(main)
