"""
Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
"""
import argparse
import data
import json
import numpy as np
import os
import subprocess
import sys
import tensorflow as tf

from collections import deque
from datetime import datetime

# Configuration
root_dir = '/tmpscratch/amotta/klinger/2019-07-23-inh-exc-relu-rnn-with-pruning'

run_id = datetime.now().strftime('%Y%m%dT%H%M%S')
output_dir = os.path.join(root_dir, run_id)
info_file = os.path.join(output_dir, 'info.json')

train_dir = os.path.join(output_dir, 'train')
valid_dir = os.path.join(output_dir, 'valid')

# Save summaries every n steps
summary_period_steps = 100
ckpt_period_secs = 1800
log_period_steps = 1000

# RNN
num_inputs = 1
num_neurons = 2000
inh_neuron_frac = 0.1
num_outputs = data.num_classes
sample_len = 100

# Training, validation, test splits
data_splits = [0.5, 0.25, 0.25]
learning_rate = 0.0001
grad_clip_norm = 1.0
batch_size = 128

# Validation config
valid_batch_size = 128
valid_period_steps = 10

# Early stopping / pruning
valid_patience = 20000
valid_last_min_loss = np.inf
valid_last_min_step = np.nan

class RunningMedian(deque):
    def is_full(self):
        return len(self) == self.maxlen

    def update(self, loss):
        self.append(loss)
        return np.median(self) if self.is_full() else np.nan

# NOTE(amotta): When it comes to early termination, the "effective validation
# loss" at time T is the median validation loss over the last 100 validations.
valid_last_losses = RunningMedian([], 100)

num_inh_neurons = int(num_neurons * inh_neuron_frac)
num_exc_neurons = num_neurons - num_inh_neurons


# Process command line arguments
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('--seed', default=0, type=int,
                        help='Seed for random number generation')
args = arg_parser.parse_args()

config = {k: v for k, v in locals().items()
          if not k.startswith('__') and
          type(v) in [int, float, str, list]}


# Save JSON file with info
this_file = os.path.abspath(__file__)
diff = subprocess.run(['git', 'diff', '--submodule=diff', '--no-color'],
                      check=True, stdout=subprocess.PIPE)
diff = diff.stdout.decode().strip()

commit = subprocess.run(['git', 'rev-parse', 'HEAD'],
                        check=True, stdout=subprocess.PIPE)
commit = commit.stdout.decode().strip()

info = {'file': this_file, 'commit': commit, 'diff': diff,
        'args': vars(args), 'config': config}

os.mkdir(output_dir)
with open(info_file, 'w') as f:
    json.dump(info, f, indent="\t")


# Input pipeline
def build_dataset(gen):
    # NOTE(amotta): Tensorflow's `Dataset.from_generator` method doesn't accept
    # generators as input. Instead, it wants a callable that then returns a
    # generator. Hence the lambda business.
    tf_dset = tf.data.Dataset.from_generator(lambda: gen,
                                             (tf.float32, tf.int32),
                                             ([sample_len, num_inputs], []))

    # NOTE(amotta): Make sure that (TC) inputs are non-negative and remain
    # within a fixed range. Also, replicate class label along time axis.
    tf_dset = tf_dset.map(lambda inputs, label: (tf.clip_by_value(inputs, -2, +2) + 2,
                                                 tf.fill(inputs.shape[:-1], label)))

    return tf_dset


generators = data.generators(sample_len, data_splits)
tf_dsets = tuple(map(build_dataset, generators))
tf_train_dset = tf_dsets[0].batch(batch_size)
tf_valid_dset = tf_dsets[1].batch(valid_batch_size)

# NOTE(amotta): The trick to online validation during training is to create a
# more general / abstract data iterator, which can be initialized from both the
# training and the validation datasets. We just need to manually call the
# respective initializers before doing training / validation.
tf_iter = tf.data.Iterator.from_structure(tf_train_dset.output_types,
                                          tf_train_dset.output_shapes)
tf_iter_init_train = tf_iter.make_initializer(tf_train_dset)
tf_iter_init_valid = tf_iter.make_initializer(tf_valid_dset)


# Network
np.random.seed(args.seed)
tf.set_random_seed(args.seed)

def trunc_normal(shape, trunc=2):
    out = np.random.randn(*shape)

    while True:
        mask = np.abs(out) > trunc
        if not np.any(mask): return out
        out[mask] = np.random.randn(np.sum(mask))


def build_rnn_weights(n_exc, n_inh, std=1.0):
    # NOTE(amotta): We build the initial weight matrix for the RNN as follows:
    # 1. Sample weights from the standard normal distribution truncated at two
    #    standard deviations (to prevent extreme values).
    # 2. Take absolute value to convert to truncated half-normal distribution.
    #    We want a non-negative weight matrix because this makes the
    #    implementation of the non-negative weight constraint much easier.
    # 3. Make sure that the signed weight matrix will have zero mean and the
    #    specified standard deviation.
    n_all = n_exc + n_inh
    w = np.abs(trunc_normal((n_all, n_all)))
    m = np.ones((n_all, n_all), dtype=np.bool)

    s = np.arange(n_all) < n_exc
    s = s.astype(np.float32) * 2 - 1

    # Set diagonal to zero (in-place)
    np.fill_diagonal(w, 0)
    np.fill_diagonal(m, False)

    w_exc = w[:n_exc, :]
    w_inh = w[n_exc:, :]

    w_inh *= w_exc.sum(axis=0) / w_inh.sum(axis=0) # Fix mean
    w *= std / np.mean(w * w, axis=0) # Fix standard deviation

    return w, m, s


# NOTE(amotta): He initialization
rnn_std = np.sqrt(2 / (num_neurons + num_inputs))

with tf.variable_scope('input'):
    tf_input_weights = rnn_std * np.abs(trunc_normal((num_inputs, num_neurons)))
    tf_input_weight_mean = np.mean(tf_input_weights)

    tf_input_weights = tf.Variable(tf_input_weights,
                                   dtype=tf.float32, name='weights',
                                   constraint=lambda t: tf.maximum(0.0, t))

with tf.variable_scope('rnn'):
    tf_rnn_weights, tf_rnn_weight_mask, tf_rnn_signs = \
        build_rnn_weights(num_exc_neurons, num_inh_neurons, rnn_std)

    tf_rnn_weight_mask = tf.Variable(tf_rnn_weight_mask,
                                     dtype=tf.bool, name='mask',
                                     trainable=False)
    tf_rnn_signs = tf.constant(tf_rnn_signs,
                               dtype=tf.float32, name='exc_mask')

    tf_rnn_weight_mask_float = tf.cast(tf_rnn_weight_mask, tf.float32)
    tf_rnn_weight_mask_count = tf.reduce_sum(tf_rnn_weight_mask_float)
    tf_rnn_weight_mask_apply = lambda t: tf.multiply(t, tf_rnn_weight_mask_float)

    tf_rnn_weight_constraint = lambda t: \
        tf_rnn_weight_mask_apply(tf.maximum(0.0, t))
    tf_rnn_weights = tf.Variable(tf_rnn_weights,
                                 dtype=tf.float32, name='weights',
                                 constraint=tf_rnn_weight_constraint)

    @tf.custom_gradient
    def freeze_masked_weights(tf_weights):
        # NOTE(amotta): We want TensorFlow to consider masked-out weights as
        # constant zeros. The simplest, most intuitive, and probably also most
        # performant way to implement this is via this magic identity function.
        #   In the forward run we just return the weight matrix passed as
        # input. The Jacobian of this "identity function" is, however, a bit
        # special: Instead of all-ones, as would be the case for the true
        # identity function, we return zeros for the masked-out entries of the
        # weight matrix.
        return tf_weights, tf_rnn_weight_mask_apply

    tf_rnn_weights_frozen = freeze_masked_weights(tf_rnn_weights)

    # NOTE(amotta): Make inhibitory neurons inhibitory!
    tf_rnn_weights_signed = tf.multiply(tf_rnn_weights_frozen,
                                        tf.reshape(tf_rnn_signs, (-1, 1)),
                                        name='weights_signed')

    # NOTE(amotta): Because the inputs are non-negative and the mean input
    # weight is positive, each neuron is positively biased. Initially, we want
    # the mean logits of the neurons to be zero. So, we're counter-balancing
    # said bias here.
    tf_rnn_biases = -2.0 * tf_input_weight_mean
    tf_rnn_biases = tf.fill((1, num_neurons), np.float32(tf_rnn_biases))
    tf_rnn_biases = tf.Variable(tf_rnn_biases, dtype=tf.float32, name='biases')


tf_inputs, tf_labels = tf_iter.get_next()
tf_input_seq = tf.unstack(tf_inputs, axis=1)

tf_rnn_logits_seq = []
tf_rnn_relus_seq = []
cur_tf_rnn_state = tf.zeros([1, num_neurons])
for cur_idx, cur_tf_input in enumerate(tf_input_seq):
    cur_tf_rnn_state = tf.matmul(cur_tf_rnn_state, tf_rnn_weights_signed) \
                     + tf.matmul(cur_tf_input, tf_input_weights) \
                     + tf_rnn_biases
    tf_rnn_logits_seq.append(cur_tf_rnn_state)

    cur_tf_rnn_state = tf.nn.relu(cur_tf_rnn_state)
    cur_tf_rnn_state = tf.minimum(2.0, cur_tf_rnn_state)
    tf_rnn_relus_seq.append(cur_tf_rnn_state)

# NOTE(amotta): We're using the first N excitatory neurons as output /
# projection neurons based on which the loss is calculated.
tf_rnn_logits = tf.stack(tf_rnn_logits_seq, axis=1)
tf_output_logits = tf_rnn_logits[:, :, :num_outputs]


# Loss calculation
tf_losses = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf_labels,
                                                           logits=tf_output_logits)
tf_losses = tf.reduce_mean(tf_losses, axis=1)
tf_loss = tf.reduce_mean(tf_losses)


# Accuracy
tf_pred_class_ids = tf.argmax(tf_output_logits, axis=-1, output_type=tf.int32)
tf_accuracy = tf.equal(tf_pred_class_ids, tf_labels)
tf_accuracy = tf.reduce_mean(tf.cast(tf_accuracy, tf.float32))


# Optimization
tf_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
tf_grad_var_pairs = tf_optimizer.compute_gradients(tf_loss)
tf_grads, tf_grad_vars = zip(*tf_grad_var_pairs)
tf_grad_norm = tf.global_norm(tf_grads)

# Gradient clipping
if not grad_clip_norm is None:
    tf_grads, _ = tf.clip_by_global_norm(tf_grads, grad_clip_norm,
                                         use_norm=tf_grad_norm)
    tf_grad_var_pairs = zip(tf_grads, tf_grad_vars)

tf_global_step = tf.Variable(0, trainable=False, name='global_step')
tf_train_op = tf_optimizer.apply_gradients(tf_grad_var_pairs, global_step=tf_global_step)


# Summaries
tf.summary.histogram('inputs', tf_inputs[0, :, :])
tf.summary.histogram('logits/rnn', tf_rnn_logits[0, :, :])
tf.summary.histogram('logits/outputs', tf_output_logits[0, :, :])

# NOTE(amotta): There variables are only set after calling the
# compute_gradients / apply_gradients methods of the Optimizer class...
tf_optimizer_vars = tf_optimizer.variables()
tf_trainable_vars = tf.trainable_variables()

for cur_tf_var in tf_optimizer_vars:
    cur_name = cur_tf_var.name.replace(':', '_')
    cur_name = 'optimizer/{}'.format(cur_name)
    tf.summary.histogram(cur_name, cur_tf_var)

for cur_tf_var in tf_trainable_vars:
    cur_name = cur_tf_var.name.replace(':', '_')
    cur_name = 'vars/{}'.format(cur_name)
    tf.summary.histogram(cur_name, cur_tf_var)

tf.summary.scalar('eval/loss', tf_loss)
tf.summary.scalar('eval/accuracy', tf_accuracy)
tf.summary.scalar('misc/grad_norm', tf_grad_norm)
tf.summary.scalar('misc/weight_count', tf_rnn_weight_mask_count)

# Separate summaries for validation data
# NOTE(amotta): It's important that we manually override the default
# collection. Otherwise, the MonitoredTrainingSession is likely to calculate
# and save these values based on the training data.
tf_valid_loss = tf.summary.scalar('eval/loss_valid', tf_loss, collections=[])
tf_valid_accuracy = tf.summary.scalar('eval/accuracy_valid', tf_accuracy, collections=[])
tf_valid_summaries = tf.summary.merge([tf_valid_loss, tf_valid_accuracy], collections=[])


## Pruning of synapses
def build_pruning_op(tf_weights, tf_mask, percentile):
    q = 1 - percentile / 100

    tf_abs_weights = tf.abs(tf_weights)
    tf_min_weight = tf.boolean_mask(tf_abs_weights, tf_mask)
    tf_numel = tf.size(tf_min_weight)

    tf_min_idx = q * tf.cast(tf_numel, tf.float32)
    tf_min_idx = tf.cast(tf.floor(tf_min_idx), tf.int32)

    tf_min_weight, _ = tf.nn.top_k(tf_min_weight, tf_numel)
    tf_min_weight = tf.gather(tf_min_weight, tf_min_idx)

    # Update mask and weights
    tf_mask_new = tf_abs_weights > tf_min_weight
    tf_weights_new = tf.where(tf_mask_new, tf_weights,
                              tf.zeros_like(tf_weights))

    return tf.group(tf.assign(tf_weights, tf_weights_new),
                    tf.assign(tf_mask, tf_mask_new))

tf_pruning_op = build_pruning_op(tf_rnn_weights,
                                 tf_rnn_weight_mask,
                                 10)


tf_init_local_vars = tf.local_variables_initializer()
tf_init_global_vars = tf.global_variables_initializer()

tf_save_vars = tf_trainable_vars + [tf_global_step, tf_rnn_weight_mask]
tf_saver = tf.train.Saver(var_list=tf_save_vars,
                          save_relative_paths=True,
                          max_to_keep=None)

# NOTE(amotta): TensorFlow's MonitoredTrainingSession evalutes some magic
# operatations during calls to the run method. These operations fail because
# none of the data Iterator's initializers have been called yet. But to
# initialize the Iterator, we need to call the run method...
#   This issue has been reported on GitHub and here's the proposed workaround:
# https://github.com/tensorflow/tensorflow/issues/12859#issuecomment-327783827
tf_local_init = tf.group(tf_init_local_vars, tf_iter_init_train)
tf_scaffold = tf.train.Scaffold(local_init_op=tf_local_init, saver=tf_saver)


# Training
with tf.train.MonitoredTrainingSession(
        scaffold=tf_scaffold,
        checkpoint_dir=train_dir,
        log_step_count_steps=log_period_steps,
        save_checkpoint_secs=ckpt_period_secs,
        save_summaries_steps=summary_period_steps) as sess:
    # Initialization
    sess.run(tf_init_global_vars)
    cur_valid_writer = tf.summary.FileWriter(valid_dir, sess.graph)

    # NOTE(amotta): The MonitoredTrainingSession claims to be a drop-in
    # replacement for a Session object created with tf.Session(). But that's
    # not true! The save method of a tf.train.Saver complains that the
    # monitored session is not a session. So, let's extract the **true**
    # session, here.
    #   Obviously, this is super-duper brittle and might break when upgrading
    # from TensorFlow 1.12 to some newer version. You've been warned!
    cur_tf_sess = sess._sess._sess._sess._sess

    while True:
        _, cur_step = sess.run((tf_train_op, tf_global_step))

        # Validation
        if cur_step % valid_period_steps == 0:
            sess.run(tf_iter_init_valid)
            cur_valid_loss, cur_valid_summaries = sess.run((tf_loss, tf_valid_summaries))
            cur_valid_writer.add_summary(cur_valid_summaries, global_step=cur_step)
            sess.run(tf_iter_init_train)

            # Calculate running median and decide how to continue
            cur_running_valid_loss = valid_last_losses.update(cur_valid_loss)

            if cur_running_valid_loss < valid_last_min_loss:
                valid_last_min_loss = cur_running_valid_loss
                valid_last_min_step = cur_step

            if cur_step - valid_last_min_step >= valid_patience:
                print('Step {}: Pruning network'.format(cur_step))
                cur_ckpt_name = 'model_pre-pruning.ckpt'.format(cur_step)
                cur_ckpt_path = os.path.join(train_dir, cur_ckpt_name)
                tf_saver.save(cur_tf_sess, cur_ckpt_path, global_step=tf_global_step)

                # NOTE(amotta): If we reach this point, the validation loss
                # hasn't improved for quite a while. Typically, we would now
                # stop the training session. Instead, we now look at the weight
                # matrix and force all small weights to zero by turning them
                # off the mask.
                sess.run(tf_pruning_op)

                # Reset validation counter
                valid_last_losses.clear()
                valid_last_min_loss = np.inf
                valid_last_min_step = cur_step

