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

from config import CONFIG
import json
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top

import io
import math
import os
import time

from absl import flags
from absl import logging

from easydict import EasyDict
import matplotlib
matplotlib.use('Agg')


FLAGS = flags.FLAGS


def visualize_batch(data, global_step, batch_size, num_steps):
    """Visualizes a batch."""
    frames = data['frames']
    frames_list = tf.unstack(frames, num=num_steps, axis=1)
    frames_summaries = tf.concat(frames_list, axis=2)
    batch_list = tf.split(frames_summaries, batch_size, axis=0)
    batch_summaries = tf.concat(batch_list, axis=1)
    tf.summary.image('train_batch', batch_summaries, step=global_step)


def visualize_nearest_neighbours(model, data, global_step, batch_size,
                                 num_steps, num_frames_per_step, split):
    """Visualize nearest neighbours in embedding space."""
    # Set learning_phase to False to use models in inference mode.
    tf.keras.backend.set_learning_phase(0)

    cnn = model['cnn']
    emb = model['emb']

    if 'tcn' in CONFIG.TRAINING_ALGO:
        cnn_feats = get_cnn_feats(
            cnn, data, training=False, num_steps=2 * num_steps)
        emb_feats = emb(cnn_feats, 2 * num_steps)
        emb_feats = tf.stack(
            tf.split(emb_feats, 2 * num_steps, axis=0)[::2], axis=1)
    else:
        cnn_feats = get_cnn_feats(cnn, data, training=False)
        emb_feats = emb(cnn_feats, num_steps)
        emb_feats = tf.stack(tf.split(emb_feats, num_steps, axis=0), axis=1)

    query_feats = emb_feats[0]

    if CONFIG.OPTICALFLOW:
        frames = data['video_frames']
    else:
        frames = data['frames']
    image_list = tf.unstack(frames, num=batch_size, axis=0)

    if 'tcn' in CONFIG.TRAINING_ALGO:
        im_list = [image_list[0]
                   [num_frames_per_step - 1::num_frames_per_step][::2]]
    else:
        im_list = [image_list[0][num_frames_per_step - 1::num_frames_per_step]]

    sim_matrix = np.zeros(
        (batch_size-1, num_steps, num_steps), dtype=np.float32)

    for i in range(1, batch_size):
        candidate_feats = emb_feats[i]

        if 'tcn' in CONFIG.TRAINING_ALGO:
            img_list = tf.unstack(image_list[i], num=2 * num_steps * num_frames_per_step,
                                  axis=0)[num_frames_per_step - 1::num_frames_per_step][::2]
        else:
            img_list = tf.unstack(image_list[i], num=num_steps * num_frames_per_step,
                                  axis=0)[num_frames_per_step - 1::num_frames_per_step]

        nn_img_list = []

        for j in range(num_steps):
            curr_query_feats = tf.tile(query_feats[j:j+1], [num_steps, 1])
            mean_squared_distance = tf.reduce_mean(
                tf.math.squared_difference(curr_query_feats, candidate_feats), axis=1)
            sim_matrix[i-1, j] = softmax(-1.0 * mean_squared_distance)
            nn_img_list.append(img_list[tf.argmin(mean_squared_distance)])

        nn_img = tf.stack(nn_img_list, axis=0)
        im_list.append(nn_img)

    def vstack(im):
        return tf.concat(tf.unstack(im, num=num_steps), axis=1)

    summary_im = tf.expand_dims(tf.concat([vstack(im) for im in im_list],
                                          axis=0), axis=0)
    tf.summary.image('%s/nn' % split, summary_im, step=global_step)
    # Convert sim_matrix to float32 as summary_image doesn't take float64
    sim_matrix = sim_matrix.astype(np.float32)
    tf.summary.image('%s/similarity_matrix' % split,
                     np.expand_dims(sim_matrix, axis=3), step=global_step)


def softmax(w, t=1.0):
    e = np.exp(np.array(w) / t)
    dist = e / np.sum(e)
    return dist


def random_choice_noreplace(m, n, axis=-1):
    # Generate m random permuations of range (0, n)
    # NumPy version: np.random.rand(m,n).argsort(axis=axis)
    return tf.cast(tf.argsort(tf.random.uniform((m, n)), axis=axis), tf.int64)


def gen_cycles(num_cycles, batch_size, cycle_len):
    """Generate cycles for alignment."""
    random_cycles = random_choice_noreplace(
        num_cycles, batch_size)[:, :cycle_len]
    return random_cycles


def get_warmup_lr(lr, global_step, lr_params):
    """Returns learning rate during warm up phase."""
    if lr_params.NUM_WARMUP_STEPS > 0:
        global_steps_int = tf.cast(global_step, tf.int32)
        warmup_steps_int = tf.constant(
            lr_params.NUM_WARMUP_STEPS, dtype=tf.int32)

        global_steps_float = tf.cast(global_steps_int, tf.float32)
        warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)

        warmup_percent_done = global_steps_float / warmup_steps_float
        warmup_lr = lr_params.INITIAL_LR * warmup_percent_done

        is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)
        lr = (1.0 - is_warmup) * lr + is_warmup * warmup_lr
    return lr


# Minimally adapted from Tensorflow object_detection code.
def manual_stepping(global_step, boundaries, rates):
    boundaries = [0] + boundaries
    num_boundaries = len(boundaries)
    rate_index = tf.reduce_max(
        tf.where(
            tf.greater_equal(global_step, boundaries),
            list(range(num_boundaries)), [0] * num_boundaries))
    return tf.reduce_sum(rates * tf.one_hot(rate_index, depth=num_boundaries))


def get_lr_fn(optimizer_config):
    """Returns function that provides current learning rate based on config.

    NOTE: This returns a function as in Eager we need to call assign to update
    the learning rate.

    Args:
      optimizer_config: EasyDict, contains params required to initialize the
        learning rate and the learning rate decay function.
    Returns:
      lr_fn: function, this can be called to return the current learning rate
        based on the provided config.
    Raises:
      ValueError: in case invalid params have been passed in the config.
    """
    lr_params = optimizer_config.LR
    # pylint: disable=g-long-lambda
    if lr_params.DECAY_TYPE == 'exp_decay':
        def lr_fn(lr, global_step): return tf.train.exponential_decay(
            lr,
            global_step,
            lr_params.EXP_DECAY_STEPS,
            lr_params.EXP_DECAY_RATE,
            staircase=True)()
    elif lr_params.DECAY_TYPE == 'manual':
        lr_step_boundaries = [int(x)
                              for x in lr_params.MANUAL_LR_STEP_BOUNDARIES]

        f = lr_params.MANUAL_LR_DECAY_RATE
        learning_rate_sequence = [(lr_params.INITIAL_LR) * f**p
                                  for p in range(len(lr_step_boundaries) + 1)]
        def lr_fn(lr, global_step): return manual_stepping(
            global_step, lr_step_boundaries, learning_rate_sequence)
    elif lr_params.DECAY_TYPE == 'fixed':
        def lr_fn(lr, global_step): return lr_params.INITIAL_LR
    elif lr_params.DECAY_TYPE == 'poly':
        def lr_fn(lr, global_step): return tf.train.polynomial_decay(
            lr,
            global_step,
            CONFIG.TRAIN.MAX_ITERS,
            end_learning_rate=0.0,
            power=1.0,
            cycle=False)
    else:
        raise ValueError('Learning rate decay type %s not supported. Only support'
                         'the following decay types: fixed, exp_decay, manual,'
                         'and poly.')

    return (lambda lr, global_step: get_warmup_lr(lr_fn(lr, global_step),
                                                  global_step, lr_params))


def get_optimizer(optimizer_config, learning_rate):
    """Returns optimizer based on config and learning rate."""
    if optimizer_config.TYPE == 'AdamOptimizer':
        opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)
    elif optimizer_config.TYPE == 'MomentumOptimizer':
        opt = tf.keras.optimizers.SGD(
            learning_rate=learning_rate, momentum=0.9)
    else:
        raise ValueError('Optimizer %s not supported. Only support the following'
                         'optimizers: AdamOptimizer, MomentumOptimizer .')
    return opt


def get_lr_opt_global_step():
    """Intializes learning rate, optimizer and global step."""
    optimizer = get_optimizer(CONFIG.OPTIMIZER, CONFIG.OPTIMIZER.LR.INITIAL_LR)
    global_step = optimizer.iterations
    learning_rate = optimizer.learning_rate
    return learning_rate, optimizer, global_step


def create_ckpt(logdir, restore=False, **ckpt_objects):
    # Since model is a dict we can insert multiple modular networks in this dict.
    checkpoint = tf.train.Checkpoint(**ckpt_objects)
    ckpt_manager = tf.train.CheckpointManager(
        checkpoint,
        directory=logdir,
        max_to_keep=10,
        keep_checkpoint_every_n_hours=1)
    status = checkpoint.restore(
        ckpt_manager.latest_checkpoint) if restore else -1
    return ckpt_manager, status, checkpoint


def restore_ckpt(logdir, **ckpt_objects):
    """Create and restore checkpoint (if one exists on the path)."""
    # Instantiate checkpoint and restore from any pre-existing checkpoint.
    # Since model is a dict we can insert multiple modular networks in this dict.
    checkpoint = tf.train.Checkpoint(**ckpt_objects)
    ckpt_manager = tf.train.CheckpointManager(
        checkpoint,
        directory=logdir,
        max_to_keep=10,
        keep_checkpoint_every_n_hours=1)
    status = checkpoint.restore(ckpt_manager.latest_checkpoint)
    return ckpt_manager, status, checkpoint


def to_dict(config):
    if isinstance(config, list):
        return [to_dict(c) for c in config]
    elif isinstance(config, EasyDict):
        return dict([(k, to_dict(v)) for k, v in config.items()])
    else:
        return config


def setup_train_dir(logdir, overwrite=False, force_train=True):
    """Setups directory for training."""
    tf.io.gfile.makedirs(logdir)
    config_path = os.path.join(logdir, 'config.json')
    if not os.path.exists(config_path) or overwrite:
        logging.info(
            'Using the existing passed in config as no config.json file exists in '
            '%s', logdir)
        with tf.io.gfile.GFile(config_path, 'w') as config_file:
            config = dict([(k, to_dict(v)) for k, v in CONFIG.items()])
            json.dump(config, config_file, sort_keys=True, indent=4)
    else:
        logging.info(
            'Using config from config.json that exists in %s.', logdir)
        with tf.io.gfile.GFile(config_path, 'r') as config_file:
            config_dict = json.load(config_file)
        CONFIG.update(config_dict)

    train_logs_dir = os.path.join(logdir, 'train.logs')
    if os.path.exists(train_logs_dir) and not force_train:
        raise ValueError('You might be overwriting a directory that already '
                         'has train_logs. Please provide a new logdir name in '
                         'config or pass --force_train while launching script.')
    tf.io.gfile.makedirs(train_logs_dir)


def setup_eval_dir(logdir, config_timeout_seconds=1):
    """Setups directory for evaluation."""
    tf.io.gfile.makedirs(logdir)
    tf.io.gfile.makedirs(os.path.join(logdir, 'eval_logs'))
    config_path = os.path.join(logdir, 'config.json')
    while not tf.io.gfile.exists(config_path):
        logging.info('Waiting for config to exist. Going to sleep '
                     ' %s for  secs.', config_timeout_seconds)
        time.sleep(config_timeout_seconds)

    while True:
        with tf.io.gfile.GFile(config_path, 'r') as config_file:
            config_dict = json.load(config_file)
        if config_dict is None:
            time.sleep(config_timeout_seconds)
        else:
            break
    CONFIG.update(config_dict)


def get_data(iterator):
    """Return a data dict which contains all the requested sequences."""
    data = iterator.get_next()
    return data, data['chosen_steps'], data['seq_lens']


@tf.function
def get_cnn_feats(cnn, data, training, num_steps=None):
    """Passes data through base CNN."""
    if num_steps is None:
        if training:
            num_steps = CONFIG.TRAIN.NUM_FRAMES * CONFIG.DATA.NUM_STEPS
        else:
            num_steps = CONFIG.EVAL.NUM_FRAMES * CONFIG.DATA.NUM_STEPS

    cnn.num_steps = num_steps
    cnn_feats = cnn(data['frames'])
    return cnn_feats


def get_context_steps(step):
    num_steps = CONFIG.DATA.NUM_STEPS
    stride = CONFIG.DATA.FRAME_STRIDE
    # We don't want to see the future.
    steps = np.arange(step - (num_steps - 1) * stride, step + stride, stride)
    return steps


def get_indices(curr_idx, num_steps, seq_len):

    steps = range(curr_idx, curr_idx + num_steps)
    single_steps = np.concatenate([get_context_steps(step) for step in steps])

    single_steps = np.concatenate(np.array(list(map(get_context_steps,
                                                    np.arange(curr_idx, curr_idx + num_steps)))))
    single_steps = np.maximum(0, single_steps)
    single_steps = np.minimum(seq_len, single_steps)
    return single_steps



def get_embeddings_dataset(model, iterator, frames_per_batch,
                           keep_data=False, optical_flow=False, keep_labels=True,
                           max_embs=None, callbacks=[]):
    """Get embeddings from a one epoch iterator."""
    keep_labels = keep_labels and CONFIG.DATA.FRAME_LABELS
    num_frames_per_step = CONFIG.DATA.NUM_STEPS
    cnn = model['cnn']
    emb = model['emb']
    embs_list = []
    labels_list = []
    steps_list = []
    seq_lens_list = []
    names_list = []
    seq_labels_list = []
    if keep_data:
        frames_list = []
    if optical_flow:
        frame_original_list = []

    n = 0

    def cond(n):
        if max_embs is None:
            return True
        else:
            return n < max_embs

    # Make Recurrent Layers stateful, set batch size.
    # We do this as we are embedding the whole sequence and that can take
    # more than one batch to be passed and we don't want to automatically
    # reset hidden states after each batch.
    if CONFIG.MODEL.EMBEDDER_TYPE == 'convgru':
        for gru_layer in emb.gru_layers:
            gru_layer.stateful = True
            gru_layer.input_spec[0].shape = [1, ]

    while cond(n):
        try:
            print(n)
            embs = []
            labels = []
            steps = []
            seq_lens = []
            names = []
            seq_labels = []
            if keep_data:
                frames = []
            if optical_flow:
                frame_original = []

            # Reset GRU states for each video.
            if CONFIG.MODEL.EMBEDDER_TYPE == 'convgru':
                for gru_layer in emb.gru_layers:
                    gru_layer.reset_states()

            data, chosen_steps, seq_len = get_data(iterator)
            seq_len = seq_len.numpy()[0]
            num_batches = int(math.ceil(float(seq_len)/frames_per_batch))
            for i in range(num_batches):
                if (i + 1) * frames_per_batch > seq_len:
                    num_steps = seq_len - i * frames_per_batch
                else:
                    num_steps = frames_per_batch
                curr_idx = i * frames_per_batch

                curr_data = {}
                for k, v in data.items():
                    # Need to do this as some modalities might not exist.
                    if len(v.shape) > 1 and v.shape[1] != 0:
                        idxes = get_indices(curr_idx, num_steps, seq_len)
                        curr_data[k] = tf.gather(v, idxes, axis=1)
                    else:
                        curr_data[k] = v

                cnn_feats = get_cnn_feats(cnn, curr_data,
                                          num_steps=num_frames_per_step * num_steps,
                                          training=False)

                emb_feats = emb(cnn_feats, num_steps)
                logging.debug('On sequence number %d, frames embedded %d', n,
                             curr_idx + num_steps)
                
                # np.save(tf.io.gfile.GFile('/air/team/saman/test_weights_old.npy', 'w'), cnn.weights[0].numpy())
                # np.save(tf.io.gfile.GFile('/air/team/saman/test_batch_old.npy', 'w'), curr_data["frames"])
                # np.save(tf.io.gfile.GFile('/air/team/saman/test_cnn_old.npy', 'w'), cnn_feats.numpy())
                # np.save(tf.io.gfile.GFile('/air/team/saman/test_emb_old.npy', 'w'), emb_feats.numpy())

                embs.append(emb_feats.numpy())
                
            for f in callbacks:
                f(np.concatenate(embs), data, chosen_steps, seq_len)
                
            steps.append(chosen_steps.numpy()[0])
            seq_lens.append(seq_len * [seq_len])
            all_labels = data['frame_labels'].numpy()[0]
            name = data['name'].numpy()[0]
            names.append(seq_len * [name])
            seq_label = data['seq_labels'].numpy()[0]
            seq_labels.append(seq_len * [seq_label])
            labels.append(all_labels)
            embs = np.concatenate(embs, axis=0)
            labels = np.concatenate(labels, axis=0)

            steps = np.concatenate(steps, axis=0)
            seq_lens = np.concatenate(seq_lens, axis=0)
            names = np.concatenate(names, axis=0)
            seq_labels = np.concatenate(seq_labels, axis=0)
            if keep_data:
                frames.append(data['frames'].numpy()[0])
                frames = np.concatenate(frames, axis=0)
            if optical_flow:
                frame_original.append(data['video_frames'].numpy()[0])
                frame_original = np.concatenate(frame_original, axis=0)

            if keep_labels:
                labels = labels[~np.isnan(embs).any(axis=1)]
                assert len(embs) == len(labels)
            seq_labels = seq_labels[~np.isnan(embs).any(axis=1)]

            names = names[~np.isnan(embs).any(axis=1)]
            seq_lens = seq_lens[~np.isnan(embs).any(axis=1)]
            steps = steps[~np.isnan(embs).any(axis=1)]
            if keep_data:
                frames = frames[~np.isnan(embs).any(axis=1)]
            if optical_flow:
                frame_original = frame_original[~np.isnan(embs).any(axis=1)]

            embs = embs[~np.isnan(embs).any(axis=1)]

            assert len(embs) == len(seq_lens)
            assert len(embs) == len(steps)
            assert len(names) == len(steps)

            embs_list.append(embs)
            if keep_labels:
                labels_list.append(labels)
            seq_labels_list.append(seq_labels)
            steps_list.append(steps)
            seq_lens_list.append(seq_lens)
            names_list.append(names)
            if keep_data:
                frames_list.append(frames)
            if optical_flow:
                frame_original_list.append(frame_original)
            n += 1
        except tf.errors.OutOfRangeError:
            logging.info('Finished embedding the dataset.')
            break

    dataset = {'embs': embs_list,
               'seq_lens': seq_lens_list,
               'steps': steps_list,
               'names': names_list,
               'seq_labels': seq_labels_list}
    if keep_data:
        dataset['frames'] = frames_list

    if optical_flow:
        dataset['frames_original'] = frame_original_list

    if keep_labels:
        dataset['labels'] = labels_list

    # Reset statefulness to recurrent layers for other evaluation tasks.
    if CONFIG.MODEL.EMBEDDER_TYPE == 'convgru':
        for gru_layer in emb.gru_layers:
            gru_layer.stateful = False

    return dataset


def gen_plot(x, y):
    """Create a pyplot, save to buffer and return TB compatible image."""
    plt.figure()
    plt.plot(x, y)
    plt.title('Val Accuracy')
    plt.ylim(0, 1)
    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    # Convert PNG buffer to TF image
    image = tf.image.decode_png(buf.getvalue(), channels=4)
    # Add the batch dimension
    image = tf.expand_dims(image, 0)
    return image


class Stopwatch(object):
    """Simple timer for measuring elapsed time."""

    def __init__(self):
        self.reset()

    def elapsed(self):
        return time.time() - self.time

    def done(self, target_interval):
        return self.elapsed() >= target_interval

    def reset(self):
        self.time = time.time()


def set_learning_phase(f):
    """Sets the correct learning phase before calling function f."""
    def wrapper(*args, **kwargs):
        """Calls the function f after setting proper learning phase."""
        if 'training' not in kwargs:
            raise ValueError('Function called with set_learning_phase decorator which'
                             ' does not have training argument.')
        training = kwargs['training']
        if training:
            # Set learning_phase to True to use models in training mode.
            tf.keras.backend.set_learning_phase(1)
        else:
            # Set learning_phase to False to use models in inference mode.
            tf.keras.backend.set_learning_phase(0)
        return f(*args, **kwargs)
    return wrapper


def load_config(config_path):

    config = None
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = json.load(f)

    assert config is not None, "config file is not provided or is corrupted"

    return config


def prepare_gpu(ind=-1):
    ind = int(ind)
    GPUS = tf.config.experimental.list_physical_devices('GPU')
    if GPUS:
        if ind > -1:
            tf.config.experimental.set_visible_devices(GPUS[ind], 'GPU')
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in GPUS:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            logging.info([len(GPUS), "Physical GPUs,", len(logical_gpus),
                         "Logical GPUs"])
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            logging.info(e)
    os.environ["CUDA_VISIBLE_DEVICES"] = str(ind)

[((406, 427), 'matplotlib.use', 'matplotlib.use', (['"""Agg"""'], {}), "('Agg')\n", (420, 427), False, 'import matplotlib\n'), ((591, 632), 'tensorflow.unstack', 'tf.unstack', (['frames'], {'num': 'num_steps', 'axis': '(1)'}), '(frames, num=num_steps, axis=1)\n', (601, 632), True, 'import tensorflow as tf\n'), ((656, 686), 'tensorflow.concat', 'tf.concat', (['frames_list'], {'axis': '(2)'}), '(frames_list, axis=2)\n', (665, 686), True, 'import tensorflow as tf\n'), ((704, 750), 'tensorflow.split', 'tf.split', (['frames_summaries', 'batch_size'], {'axis': '(0)'}), '(frames_summaries, batch_size, axis=0)\n', (712, 750), True, 'import tensorflow as tf\n'), ((773, 802), 'tensorflow.concat', 'tf.concat', (['batch_list'], {'axis': '(1)'}), '(batch_list, axis=1)\n', (782, 802), True, 'import tensorflow as tf\n'), ((807, 873), 'tensorflow.summary.image', 'tf.summary.image', (['"""train_batch"""', 'batch_summaries'], {'step': 'global_step'}), "('train_batch', batch_summaries, step=global_step)\n", (823, 873), True, 'import tensorflow as tf\n'), ((1150, 1188), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', (['(0)'], {}), '(0)\n', (1185, 1188), True, 'import tensorflow as tf\n'), ((1874, 1916), 'tensorflow.unstack', 'tf.unstack', (['frames'], {'num': 'batch_size', 'axis': '(0)'}), '(frames, num=batch_size, axis=0)\n', (1884, 1916), True, 'import tensorflow as tf\n'), ((2169, 2235), 'numpy.zeros', 'np.zeros', (['(batch_size - 1, num_steps, num_steps)'], {'dtype': 'np.float32'}), '((batch_size - 1, num_steps, num_steps), dtype=np.float32)\n', (2177, 2235), True, 'import numpy as np\n'), ((3456, 3519), 'tensorflow.summary.image', 'tf.summary.image', (["('%s/nn' % split)", 'summary_im'], {'step': 'global_step'}), "('%s/nn' % split, summary_im, step=global_step)\n", (3472, 3519), True, 'import tensorflow as tf\n'), ((8624, 8659), 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ([], {}), '(**ckpt_objects)\n', (8643, 8659), True, 'import tensorflow as tf\n'), ((8679, 8788), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (['checkpoint'], {'directory': 'logdir', 'max_to_keep': '(10)', 'keep_checkpoint_every_n_hours': '(1)'}), '(checkpoint, directory=logdir, max_to_keep=10,\n    keep_checkpoint_every_n_hours=1)\n', (8705, 8788), True, 'import tensorflow as tf\n'), ((9241, 9276), 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ([], {}), '(**ckpt_objects)\n', (9260, 9276), True, 'import tensorflow as tf\n'), ((9296, 9405), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (['checkpoint'], {'directory': 'logdir', 'max_to_keep': '(10)', 'keep_checkpoint_every_n_hours': '(1)'}), '(checkpoint, directory=logdir, max_to_keep=10,\n    keep_checkpoint_every_n_hours=1)\n', (9322, 9405), True, 'import tensorflow as tf\n'), ((9891, 9919), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', (['logdir'], {}), '(logdir)\n', (9911, 9919), True, 'import tensorflow as tf\n'), ((9938, 9973), 'os.path.join', 'os.path.join', (['logdir', '"""config.json"""'], {}), "(logdir, 'config.json')\n", (9950, 9973), False, 'import os\n'), ((10640, 10674), 'os.path.join', 'os.path.join', (['logdir', '"""train.logs"""'], {}), "(logdir, 'train.logs')\n", (10652, 10674), False, 'import os\n'), ((10977, 11013), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', (['train_logs_dir'], {}), '(train_logs_dir)\n', (10997, 11013), True, 'import tensorflow as tf\n'), ((11117, 11145), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', (['logdir'], {}), '(logdir)\n', (11137, 11145), True, 'import tensorflow as tf\n'), ((11224, 11259), 'os.path.join', 'os.path.join', (['logdir', '"""config.json"""'], {}), "(logdir, 'config.json')\n", (11236, 11259), False, 'import os\n'), ((11727, 11753), 'config.CONFIG.update', 'CONFIG.update', (['config_dict'], {}), '(config_dict)\n', (11740, 11753), False, 'from config import CONFIG\n'), ((12501, 12566), 'numpy.arange', 'np.arange', (['(step - (num_steps - 1) * stride)', '(step + stride)', 'stride'], {}), '(step - (num_steps - 1) * stride, step + stride, stride)\n', (12510, 12566), True, 'import numpy as np\n'), ((12952, 12979), 'numpy.maximum', 'np.maximum', (['(0)', 'single_steps'], {}), '(0, single_steps)\n', (12962, 12979), True, 'import numpy as np\n'), ((12999, 13032), 'numpy.minimum', 'np.minimum', (['seq_len', 'single_steps'], {}), '(seq_len, single_steps)\n', (13009, 13032), True, 'import numpy as np\n'), ((19706, 19718), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (19716, 19718), True, 'import matplotlib.pyplot as plt\n'), ((19723, 19737), 'matplotlib.pyplot.plot', 'plt.plot', (['x', 'y'], {}), '(x, y)\n', (19731, 19737), True, 'import matplotlib.pyplot as plt\n'), ((19742, 19767), 'matplotlib.pyplot.title', 'plt.title', (['"""Val Accuracy"""'], {}), "('Val Accuracy')\n", (19751, 19767), True, 'import matplotlib.pyplot as plt\n'), ((19772, 19786), 'matplotlib.pyplot.ylim', 'plt.ylim', (['(0)', '(1)'], {}), '(0, 1)\n', (19780, 19786), True, 'import matplotlib.pyplot as plt\n'), ((19791, 19809), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), '()\n', (19807, 19809), True, 'import matplotlib.pyplot as plt\n'), ((19820, 19832), 'io.BytesIO', 'io.BytesIO', ([], {}), '()\n', (19830, 19832), False, 'import io\n'), ((19837, 19867), 'matplotlib.pyplot.savefig', 'plt.savefig', (['buf'], {'format': '"""png"""'}), "(buf, format='png')\n", (19848, 19867), True, 'import matplotlib.pyplot as plt\n'), ((20023, 20047), 'tensorflow.expand_dims', 'tf.expand_dims', (['image', '(0)'], {}), '(image, 0)\n', (20037, 20047), True, 'import tensorflow as tf\n'), ((21221, 21248), 'os.path.exists', 'os.path.exists', (['config_path'], {}), '(config_path)\n', (21235, 21248), False, 'import os\n'), ((21475, 21526), 'tensorflow.config.experimental.list_physical_devices', 'tf.config.experimental.list_physical_devices', (['"""GPU"""'], {}), "('GPU')\n", (21519, 21526), True, 'import tensorflow as tf\n'), ((3172, 3201), 'tensorflow.stack', 'tf.stack', (['nn_img_list'], {'axis': '(0)'}), '(nn_img_list, axis=0)\n', (3180, 3201), True, 'import tensorflow as tf\n'), ((3715, 3749), 'numpy.expand_dims', 'np.expand_dims', (['sim_matrix'], {'axis': '(3)'}), '(sim_matrix, axis=3)\n', (3729, 3749), True, 'import numpy as np\n'), ((3841, 3850), 'numpy.sum', 'np.sum', (['e'], {}), '(e)\n', (3847, 3850), True, 'import numpy as np\n'), ((4484, 4514), 'tensorflow.cast', 'tf.cast', (['global_step', 'tf.int32'], {}), '(global_step, tf.int32)\n', (4491, 4514), True, 'import tensorflow as tf\n'), ((4542, 4597), 'tensorflow.constant', 'tf.constant', (['lr_params.NUM_WARMUP_STEPS'], {'dtype': 'tf.int32'}), '(lr_params.NUM_WARMUP_STEPS, dtype=tf.int32)\n', (4553, 4597), True, 'import tensorflow as tf\n'), ((4641, 4678), 'tensorflow.cast', 'tf.cast', (['global_steps_int', 'tf.float32'], {}), '(global_steps_int, tf.float32)\n', (4648, 4678), True, 'import tensorflow as tf\n'), ((4708, 4745), 'tensorflow.cast', 'tf.cast', (['warmup_steps_int', 'tf.float32'], {}), '(warmup_steps_int, tf.float32)\n', (4715, 4745), True, 'import tensorflow as tf\n'), ((4901, 4957), 'tensorflow.cast', 'tf.cast', (['(global_steps_int < warmup_steps_int)', 'tf.float32'], {}), '(global_steps_int < warmup_steps_int, tf.float32)\n', (4908, 4957), True, 'import tensorflow as tf\n'), ((7775, 7828), 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), '(learning_rate=learning_rate)\n', (7799, 7828), True, 'import tensorflow as tf\n'), ((10035, 10139), 'absl.logging.info', 'logging.info', (['"""Using the existing passed in config as no config.json file exists in %s"""', 'logdir'], {}), "(\n    'Using the existing passed in config as no config.json file exists in %s',\n    logdir)\n", (10047, 10139), False, 'from absl import logging\n'), ((10383, 10455), 'absl.logging.info', 'logging.info', (['"""Using config from config.json that exists in %s."""', 'logdir'], {}), "('Using config from config.json that exists in %s.', logdir)\n", (10395, 10455), False, 'from absl import logging\n'), ((10591, 10617), 'config.CONFIG.update', 'CONFIG.update', (['config_dict'], {}), '(config_dict)\n', (10604, 10617), False, 'from config import CONFIG\n'), ((10682, 10712), 'os.path.exists', 'os.path.exists', (['train_logs_dir'], {}), '(train_logs_dir)\n', (10696, 10712), False, 'import os\n'), ((11171, 11204), 'os.path.join', 'os.path.join', (['logdir', '"""eval_logs"""'], {}), "(logdir, 'eval_logs')\n", (11183, 11204), False, 'import os\n'), ((11274, 11305), 'tensorflow.io.gfile.exists', 'tf.io.gfile.exists', (['config_path'], {}), '(config_path)\n', (11292, 11305), True, 'import tensorflow as tf\n'), ((11315, 11417), 'absl.logging.info', 'logging.info', (['"""Waiting for config to exist. Going to sleep  %s for  secs."""', 'config_timeout_seconds'], {}), "('Waiting for config to exist. Going to sleep  %s for  secs.',\n    config_timeout_seconds)\n", (11327, 11417), False, 'from absl import logging\n'), ((11446, 11480), 'time.sleep', 'time.sleep', (['config_timeout_seconds'], {}), '(config_timeout_seconds)\n', (11456, 11480), False, 'import time\n'), ((20381, 20392), 'time.time', 'time.time', ([], {}), '()\n', (20390, 20392), False, 'import time\n'), ((1669, 1707), 'tensorflow.split', 'tf.split', (['emb_feats', 'num_steps'], {'axis': '(0)'}), '(emb_feats, num_steps, axis=0)\n', (1677, 1707), True, 'import tensorflow as tf\n'), ((2825, 2870), 'tensorflow.tile', 'tf.tile', (['query_feats[j:j + 1]', '[num_steps, 1]'], {}), '(query_feats[j:j + 1], [num_steps, 1])\n', (2832, 2870), True, 'import tensorflow as tf\n'), ((3279, 3308), 'tensorflow.unstack', 'tf.unstack', (['im'], {'num': 'num_steps'}), '(im, num=num_steps)\n', (3289, 3308), True, 'import tensorflow as tf\n'), ((3809, 3820), 'numpy.array', 'np.array', (['w'], {}), '(w)\n', (3817, 3820), True, 'import numpy as np\n'), ((4055, 4080), 'tensorflow.random.uniform', 'tf.random.uniform', (['(m, n)'], {}), '((m, n))\n', (4072, 4080), True, 'import tensorflow as tf\n'), ((5279, 5320), 'tensorflow.greater_equal', 'tf.greater_equal', (['global_step', 'boundaries'], {}), '(global_step, boundaries)\n', (5295, 5320), True, 'import tensorflow as tf\n'), ((5419, 5463), 'tensorflow.one_hot', 'tf.one_hot', (['rate_index'], {'depth': 'num_boundaries'}), '(rate_index, depth=num_boundaries)\n', (5429, 5463), True, 'import tensorflow as tf\n'), ((7898, 7964), 'tensorflow.keras.optimizers.SGD', 'tf.keras.optimizers.SGD', ([], {'learning_rate': 'learning_rate', 'momentum': '(0.9)'}), '(learning_rate=learning_rate, momentum=0.9)\n', (7921, 7964), True, 'import tensorflow as tf\n'), ((9985, 10012), 'os.path.exists', 'os.path.exists', (['config_path'], {}), '(config_path)\n', (9999, 10012), False, 'import os\n'), ((10172, 10207), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['config_path', '"""w"""'], {}), "(config_path, 'w')\n", (10189, 10207), True, 'import tensorflow as tf\n'), ((10308, 10364), 'json.dump', 'json.dump', (['config', 'config_file'], {'sort_keys': '(True)', 'indent': '(4)'}), '(config, config_file, sort_keys=True, indent=4)\n', (10317, 10364), False, 'import json\n'), ((10482, 10517), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['config_path', '"""r"""'], {}), "(config_path, 'r')\n", (10499, 10517), True, 'import tensorflow as tf\n'), ((10560, 10582), 'json.load', 'json.load', (['config_file'], {}), '(config_file)\n', (10569, 10582), False, 'import json\n'), ((11511, 11546), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['config_path', '"""r"""'], {}), "(config_path, 'r')\n", (11528, 11546), True, 'import tensorflow as tf\n'), ((11589, 11611), 'json.load', 'json.load', (['config_file'], {}), '(config_file)\n', (11598, 11611), False, 'import json\n'), ((11656, 11690), 'time.sleep', 'time.sleep', (['config_timeout_seconds'], {}), '(config_timeout_seconds)\n', (11666, 11690), False, 'import time\n'), ((17059, 17087), 'numpy.concatenate', 'np.concatenate', (['embs'], {'axis': '(0)'}), '(embs, axis=0)\n', (17073, 17087), True, 'import numpy as np\n'), ((17109, 17139), 'numpy.concatenate', 'np.concatenate', (['labels'], {'axis': '(0)'}), '(labels, axis=0)\n', (17123, 17139), True, 'import numpy as np\n'), ((17161, 17190), 'numpy.concatenate', 'np.concatenate', (['steps'], {'axis': '(0)'}), '(steps, axis=0)\n', (17175, 17190), True, 'import numpy as np\n'), ((17214, 17246), 'numpy.concatenate', 'np.concatenate', (['seq_lens'], {'axis': '(0)'}), '(seq_lens, axis=0)\n', (17228, 17246), True, 'import numpy as np\n'), ((17267, 17296), 'numpy.concatenate', 'np.concatenate', (['names'], {'axis': '(0)'}), '(names, axis=0)\n', (17281, 17296), True, 'import numpy as np\n'), ((17322, 17356), 'numpy.concatenate', 'np.concatenate', (['seq_labels'], {'axis': '(0)'}), '(seq_labels, axis=0)\n', (17336, 17356), True, 'import numpy as np\n'), ((20228, 20239), 'time.time', 'time.time', ([], {}), '()\n', (20237, 20239), False, 'import time\n'), ((20931, 20969), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', (['(1)'], {}), '(1)\n', (20966, 20969), True, 'import tensorflow as tf\n'), ((21071, 21109), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', (['(0)'], {}), '(0)\n', (21106, 21109), True, 'import tensorflow as tf\n'), ((21308, 21320), 'json.load', 'json.load', (['f'], {}), '(f)\n', (21317, 21320), False, 'import json\n'), ((21573, 21633), 'tensorflow.config.experimental.set_visible_devices', 'tf.config.experimental.set_visible_devices', (['GPUS[ind]', '"""GPU"""'], {}), "(GPUS[ind], 'GPU')\n", (21615, 21633), True, 'import tensorflow as tf\n'), ((21843, 21893), 'tensorflow.config.experimental.list_logical_devices', 'tf.config.experimental.list_logical_devices', (['"""GPU"""'], {}), "('GPU')\n", (21886, 21893), True, 'import tensorflow as tf\n'), ((1466, 1508), 'tensorflow.split', 'tf.split', (['emb_feats', '(2 * num_steps)'], {'axis': '(0)'}), '(emb_feats, 2 * num_steps, axis=0)\n', (1474, 1508), True, 'import tensorflow as tf\n'), ((2581, 2651), 'tensorflow.unstack', 'tf.unstack', (['image_list[i]'], {'num': '(num_steps * num_frames_per_step)', 'axis': '(0)'}), '(image_list[i], num=num_steps * num_frames_per_step, axis=0)\n', (2591, 2651), True, 'import tensorflow as tf\n'), ((2937, 2998), 'tensorflow.math.squared_difference', 'tf.math.squared_difference', (['curr_query_feats', 'candidate_feats'], {}), '(curr_query_feats, candidate_feats)\n', (2963, 2998), True, 'import tensorflow as tf\n'), ((6209, 6325), 'tensorflow.train.exponential_decay', 'tf.train.exponential_decay', (['lr', 'global_step', 'lr_params.EXP_DECAY_STEPS', 'lr_params.EXP_DECAY_RATE'], {'staircase': '(True)'}), '(lr, global_step, lr_params.EXP_DECAY_STEPS,\n    lr_params.EXP_DECAY_RATE, staircase=True)\n', (6235, 6325), True, 'import tensorflow as tf\n'), ((15907, 15994), 'absl.logging.debug', 'logging.debug', (['"""On sequence number %d, frames embedded %d"""', 'n', '(curr_idx + num_steps)'], {}), "('On sequence number %d, frames embedded %d', n, curr_idx +\n    num_steps)\n", (15920, 15994), False, 'from absl import logging\n'), ((17465, 17495), 'numpy.concatenate', 'np.concatenate', (['frames'], {'axis': '(0)'}), '(frames, axis=0)\n', (17479, 17495), True, 'import numpy as np\n'), ((17629, 17667), 'numpy.concatenate', 'np.concatenate', (['frame_original'], {'axis': '(0)'}), '(frame_original, axis=0)\n', (17643, 17667), True, 'import numpy as np\n'), ((18925, 18972), 'absl.logging.info', 'logging.info', (['"""Finished embedding the dataset."""'], {}), "('Finished embedding the dataset.')\n", (18937, 18972), False, 'from absl import logging\n'), ((21764, 21815), 'tensorflow.config.experimental.set_memory_growth', 'tf.config.experimental.set_memory_growth', (['gpu', '(True)'], {}), '(gpu, True)\n', (21804, 21815), True, 'import tensorflow as tf\n'), ((22130, 22145), 'absl.logging.info', 'logging.info', (['e'], {}), '(e)\n', (22142, 22145), False, 'from absl import logging\n'), ((2384, 2458), 'tensorflow.unstack', 'tf.unstack', (['image_list[i]'], {'num': '(2 * num_steps * num_frames_per_step)', 'axis': '(0)'}), '(image_list[i], num=2 * num_steps * num_frames_per_step, axis=0)\n', (2394, 2458), True, 'import tensorflow as tf\n'), ((12887, 12928), 'numpy.arange', 'np.arange', (['curr_idx', '(curr_idx + num_steps)'], {}), '(curr_idx, curr_idx + num_steps)\n', (12896, 12928), True, 'import numpy as np\n'), ((16585, 16605), 'numpy.concatenate', 'np.concatenate', (['embs'], {}), '(embs)\n', (16599, 16605), True, 'import numpy as np\n'), ((3119, 3151), 'tensorflow.argmin', 'tf.argmin', (['mean_squared_distance'], {}), '(mean_squared_distance)\n', (3128, 3151), True, 'import tensorflow as tf\n'), ((7047, 7164), 'tensorflow.train.polynomial_decay', 'tf.train.polynomial_decay', (['lr', 'global_step', 'CONFIG.TRAIN.MAX_ITERS'], {'end_learning_rate': '(0.0)', 'power': '(1.0)', 'cycle': '(False)'}), '(lr, global_step, CONFIG.TRAIN.MAX_ITERS,\n    end_learning_rate=0.0, power=1.0, cycle=False)\n', (7072, 7164), True, 'import tensorflow as tf\n'), ((10279, 10293), 'config.CONFIG.items', 'CONFIG.items', ([], {}), '()\n', (10291, 10293), False, 'from config import CONFIG\n'), ((15539, 15566), 'tensorflow.gather', 'tf.gather', (['v', 'idxes'], {'axis': '(1)'}), '(v, idxes, axis=1)\n', (15548, 15566), True, 'import tensorflow as tf\n'), ((17843, 17857), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (17851, 17857), True, 'import numpy as np\n'), ((17899, 17913), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (17907, 17913), True, 'import numpy as np\n'), ((17960, 17974), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (17968, 17974), True, 'import numpy as np\n'), ((18015, 18029), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (18023, 18029), True, 'import numpy as np\n'), ((18262, 18276), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (18270, 18276), True, 'import numpy as np\n'), ((17730, 17744), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (17738, 17744), True, 'import numpy as np\n'), ((18102, 18116), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (18110, 18116), True, 'import numpy as np\n'), ((18208, 18222), 'numpy.isnan', 'np.isnan', (['embs'], {}), '(embs)\n', (18216, 18222), True, 'import numpy as np\n')]