
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)

[((22, 0, 22, 21), 'matplotlib.use', 'matplotlib.use', ({(22, 15, 22, 20): '"""Agg"""'}, {}), "('Agg')", False, 'import matplotlib\n'), ((31, 18, 31, 59), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((32, 23, 32, 53), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((33, 17, 33, 63), 'tensorflow.split', 'tf.split', (), '', True, 'import tensorflow as tf\n'), ((34, 22, 34, 51), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((35, 4, 35, 70), 'tensorflow.summary.image', 'tf.summary.image', (), '', True, 'import tensorflow as tf\n'), ((42, 4, 42, 42), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', ({(42, 40, 42, 41): '(0)'}, {}), '(0)', True, 'import tensorflow as tf\n'), ((64, 17, 64, 59), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((72, 17, 73, 63), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((102, 4, 102, 67), 'tensorflow.summary.image', 'tf.summary.image', (), '', True, 'import tensorflow as tf\n'), ((232, 17, 232, 52), 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ({}, {}), '(**ckpt_objects)', True, 'import tensorflow as tf\n'), ((233, 19, 237, 40), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (), '', True, 'import tensorflow as tf\n'), ((247, 17, 247, 52), 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ({}, {}), '(**ckpt_objects)', True, 'import tensorflow as tf\n'), ((248, 19, 252, 40), 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', (), '', True, 'import tensorflow as tf\n'), ((268, 4, 268, 32), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', ({(268, 25, 268, 31): 'logdir'}, {}), '(logdir)', True, 'import tensorflow as tf\n'), ((269, 18, 269, 53), 'os.path.join', 'os.path.join', ({(269, 31, 269, 37): 'logdir', (269, 39, 269, 52): '"""config.json"""'}, {}), "(logdir, 'config.json')", False, 'import os\n'), ((284, 21, 284, 55), 'os.path.join', 'os.path.join', ({(284, 34, 284, 40): 'logdir', (284, 42, 284, 54): '"""train.logs"""'}, {}), "(logdir, 'train.logs')", False, 'import os\n'), ((289, 4, 289, 40), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', ({(289, 25, 289, 39): 'train_logs_dir'}, {}), '(train_logs_dir)', True, 'import tensorflow as tf\n'), ((294, 4, 294, 32), 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', ({(294, 25, 294, 31): 'logdir'}, {}), '(logdir)', True, 'import tensorflow as tf\n'), ((296, 18, 296, 53), 'os.path.join', 'os.path.join', ({(296, 31, 296, 37): 'logdir', (296, 39, 296, 52): '"""config.json"""'}, {}), "(logdir, 'config.json')", False, 'import os\n'), ((309, 4, 309, 30), 'config.CONFIG.update', 'CONFIG.update', ({(309, 18, 309, 29): 'config_dict'}, {}), '(config_dict)', False, 'from config import CONFIG\n'), ((336, 12, 336, 77), 'numpy.arange', 'np.arange', ({(336, 22, 336, 53): 'step - (num_steps - 1) * stride', (336, 55, 336, 68): 'step + stride', (336, 70, 336, 76): 'stride'}, {}), '(step - (num_steps - 1) * stride, step + stride, stride)', True, 'import numpy as np\n'), ((347, 19, 347, 46), 'numpy.maximum', 'np.maximum', ({(347, 30, 347, 31): '0', (347, 33, 347, 45): 'single_steps'}, {}), '(0, single_steps)', True, 'import numpy as np\n'), ((348, 19, 348, 52), 'numpy.minimum', 'np.minimum', ({(348, 30, 348, 37): 'seq_len', (348, 39, 348, 51): 'single_steps'}, {}), '(seq_len, single_steps)', True, 'import numpy as np\n'), ((526, 4, 526, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((527, 4, 527, 18), 'matplotlib.pyplot.plot', 'plt.plot', ({(527, 13, 527, 14): 'x', (527, 16, 527, 17): 'y'}, {}), '(x, y)', True, 'import matplotlib.pyplot as plt\n'), ((528, 4, 528, 29), 'matplotlib.pyplot.title', 'plt.title', ({(528, 14, 528, 28): '"""Val Accuracy"""'}, {}), "('Val Accuracy')", True, 'import matplotlib.pyplot as plt\n'), ((529, 4, 529, 18), 'matplotlib.pyplot.ylim', 'plt.ylim', ({(529, 13, 529, 14): '(0)', (529, 16, 529, 17): '(1)'}, {}), '(0, 1)', True, 'import matplotlib.pyplot as plt\n'), ((530, 4, 530, 22), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((531, 10, 531, 22), 'io.BytesIO', 'io.BytesIO', ({}, {}), '()', False, 'import io\n'), ((532, 4, 532, 34), 'matplotlib.pyplot.savefig', 'plt.savefig', (), '', True, 'import matplotlib.pyplot as plt\n'), ((537, 12, 537, 36), 'tensorflow.expand_dims', 'tf.expand_dims', ({(537, 27, 537, 32): 'image', (537, 34, 537, 35): '0'}, {}), '(image, 0)', True, 'import tensorflow as tf\n'), ((578, 7, 578, 34), 'os.path.exists', 'os.path.exists', ({(578, 22, 578, 33): 'config_path'}, {}), '(config_path)', False, 'import os\n'), ((589, 11, 589, 62), 'tensorflow.config.experimental.list_physical_devices', 'tf.config.experimental.list_physical_devices', ({(589, 56, 589, 61): '"""GPU"""'}, {}), "('GPU')", True, 'import tensorflow as tf\n'), ((94, 17, 94, 46), 'tensorflow.stack', 'tf.stack', (), '', True, 'import tensorflow as tf\n'), ((106, 21, 106, 55), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((111, 15, 111, 24), 'numpy.sum', 'np.sum', ({(111, 22, 111, 23): 'e'}, {}), '(e)', True, 'import numpy as np\n'), ((131, 27, 131, 57), 'tensorflow.cast', 'tf.cast', ({(131, 35, 131, 46): 'global_step', (131, 48, 131, 56): 'tf.int32'}, {}), '(global_step, tf.int32)', True, 'import tensorflow as tf\n'), ((132, 27, 133, 55), 'tensorflow.constant', 'tf.constant', (), '', True, 'import tensorflow as tf\n'), ((135, 29, 135, 66), 'tensorflow.cast', 'tf.cast', ({(135, 37, 135, 53): 'global_steps_int', (135, 55, 135, 65): 'tf.float32'}, {}), '(global_steps_int, tf.float32)', True, 'import tensorflow as tf\n'), ((136, 29, 136, 66), 'tensorflow.cast', 'tf.cast', ({(136, 37, 136, 53): 'warmup_steps_int', (136, 55, 136, 65): 'tf.float32'}, {}), '(warmup_steps_int, tf.float32)', True, 'import tensorflow as tf\n'), ((141, 20, 141, 76), 'tensorflow.cast', 'tf.cast', ({(141, 28, 141, 63): 'global_steps_int < warmup_steps_int', (141, 65, 141, 75): 'tf.float32'}, {}), '(global_steps_int < warmup_steps_int, tf.float32)', True, 'import tensorflow as tf\n'), ((212, 14, 212, 67), 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (), '', True, 'import tensorflow as tf\n'), ((271, 8, 273, 25), 'absl.logging.info', 'logging.info', ({(272, 12, 273, 16): '"""Using the existing passed in config as no config.json file exists in %s"""', (273, 18, 273, 24): 'logdir'}, {}), "(\n    'Using the existing passed in config as no config.json file exists in %s',\n    logdir)", False, 'from absl import logging\n'), ((278, 8, 279, 71), 'absl.logging.info', 'logging.info', ({(279, 12, 279, 62): '"""Using config from config.json that exists in %s."""', (279, 64, 279, 70): 'logdir'}, {}), "('Using config from config.json that exists in %s.', logdir)", False, 'from absl import logging\n'), ((282, 8, 282, 34), 'config.CONFIG.update', 'CONFIG.update', ({(282, 22, 282, 33): 'config_dict'}, {}), '(config_dict)', False, 'from config import CONFIG\n'), ((285, 7, 285, 37), 'os.path.exists', 'os.path.exists', ({(285, 22, 285, 36): 'train_logs_dir'}, {}), '(train_logs_dir)', False, 'import os\n'), ((295, 25, 295, 58), 'os.path.join', 'os.path.join', ({(295, 38, 295, 44): 'logdir', (295, 46, 295, 57): '"""eval_logs"""'}, {}), "(logdir, 'eval_logs')", False, 'import os\n'), ((297, 14, 297, 45), 'tensorflow.io.gfile.exists', 'tf.io.gfile.exists', ({(297, 33, 297, 44): 'config_path'}, {}), '(config_path)', True, 'import tensorflow as tf\n'), ((298, 8, 299, 62), 'absl.logging.info', 'logging.info', ({(298, 21, 299, 37): '"""Waiting for config to exist. Going to sleep  %s for  secs."""', (299, 39, 299, 61): 'config_timeout_seconds'}, {}), "('Waiting for config to exist. Going to sleep  %s for  secs.',\n    config_timeout_seconds)", False, 'from absl import logging\n'), ((300, 8, 300, 42), 'time.sleep', 'time.sleep', ({(300, 19, 300, 41): 'config_timeout_seconds'}, {}), '(config_timeout_seconds)', False, 'import time\n'), ((554, 20, 554, 31), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((56, 29, 56, 67), 'tensorflow.split', 'tf.split', (), '', True, 'import tensorflow as tf\n'), ((88, 31, 88, 74), 'tensorflow.tile', 'tf.tile', ({(88, 39, 88, 57): 'query_feats[j:j + 1]', (88, 59, 88, 73): '[num_steps, 1]'}, {}), '(query_feats[j:j + 1], [num_steps, 1])', True, 'import tensorflow as tf\n'), ((98, 25, 98, 54), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((110, 15, 110, 26), 'numpy.array', 'np.array', ({(110, 24, 110, 25): 'w'}, {}), '(w)', True, 'import numpy as np\n'), ((118, 30, 118, 55), 'tensorflow.random.uniform', 'tf.random.uniform', ({(118, 48, 118, 54): '(m, n)'}, {}), '((m, n))', True, 'import tensorflow as tf\n'), ((152, 12, 152, 53), 'tensorflow.greater_equal', 'tf.greater_equal', ({(152, 29, 152, 40): 'global_step', (152, 42, 152, 52): 'boundaries'}, {}), '(global_step, boundaries)', True, 'import tensorflow as tf\n'), ((154, 33, 154, 77), 'tensorflow.one_hot', 'tf.one_hot', (), '', True, 'import tensorflow as tf\n'), ((214, 14, 215, 54), 'tensorflow.keras.optimizers.SGD', 'tf.keras.optimizers.SGD', (), '', True, 'import tensorflow as tf\n'), ((270, 11, 270, 38), 'os.path.exists', 'os.path.exists', ({(270, 26, 270, 37): 'config_path'}, {}), '(config_path)', False, 'import os\n'), ((274, 13, 274, 48), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', ({(274, 31, 274, 42): 'config_path', (274, 44, 274, 47): '"""w"""'}, {}), "(config_path, 'w')", True, 'import tensorflow as tf\n'), ((276, 12, 276, 68), 'json.dump', 'json.dump', (), '', False, 'import json\n'), ((280, 13, 280, 48), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', ({(280, 31, 280, 42): 'config_path', (280, 44, 280, 47): '"""r"""'}, {}), "(config_path, 'r')", True, 'import tensorflow as tf\n'), ((281, 26, 281, 48), 'json.load', 'json.load', ({(281, 36, 281, 47): 'config_file'}, {}), '(config_file)', False, 'import json\n'), ((303, 13, 303, 48), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', ({(303, 31, 303, 42): 'config_path', (303, 44, 303, 47): '"""r"""'}, {}), "(config_path, 'r')", True, 'import tensorflow as tf\n'), ((304, 26, 304, 48), 'json.load', 'json.load', ({(304, 36, 304, 47): 'config_file'}, {}), '(config_file)', False, 'import json\n'), ((306, 12, 306, 46), 'time.sleep', 'time.sleep', ({(306, 23, 306, 45): 'config_timeout_seconds'}, {}), '(config_timeout_seconds)', False, 'import time\n'), ((453, 19, 453, 47), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((454, 21, 454, 51), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((456, 20, 456, 49), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((457, 23, 457, 55), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((458, 20, 458, 49), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((459, 25, 459, 59), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((548, 15, 548, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((567, 12, 567, 50), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', ({(567, 48, 567, 49): '(1)'}, {}), '(1)', True, 'import tensorflow as tf\n'), ((570, 12, 570, 50), 'tensorflow.keras.backend.set_learning_phase', 'tf.keras.backend.set_learning_phase', ({(570, 48, 570, 49): '(0)'}, {}), '(0)', True, 'import tensorflow as tf\n'), ((580, 21, 580, 33), 'json.load', 'json.load', ({(580, 31, 580, 32): 'f'}, {}), '(f)', False, 'import json\n'), ((592, 12, 592, 72), 'tensorflow.config.experimental.set_visible_devices', 'tf.config.experimental.set_visible_devices', ({(592, 55, 592, 64): 'GPUS[ind]', (592, 66, 592, 71): '"""GPU"""'}, {}), "(GPUS[ind], 'GPU')", True, 'import tensorflow as tf\n'), ((597, 27, 597, 77), 'tensorflow.config.experimental.list_logical_devices', 'tf.config.experimental.list_logical_devices', ({(597, 71, 597, 76): '"""GPU"""'}, {}), "('GPU')", True, 'import tensorflow as tf\n'), ((52, 12, 52, 54), 'tensorflow.split', 'tf.split', (), '', True, 'import tensorflow as tf\n'), ((82, 23, 83, 41), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((90, 16, 90, 77), 'tensorflow.math.squared_difference', 'tf.math.squared_difference', ({(90, 43, 90, 59): 'curr_query_feats', (90, 61, 90, 76): 'candidate_feats'}, {}), '(curr_query_feats, candidate_feats)', True, 'import tensorflow as tf\n'), ((175, 43, 180, 27), 'tensorflow.train.exponential_decay', 'tf.train.exponential_decay', (), '', True, 'import tensorflow as tf\n'), ((432, 16, 433, 50), 'absl.logging.debug', 'logging.debug', ({(432, 30, 432, 73): '"""On sequence number %d, frames embedded %d"""', (432, 75, 432, 76): 'n', (433, 29, 433, 49): '(curr_idx + num_steps)'}, {}), "('On sequence number %d, frames embedded %d', n, curr_idx +\n    num_steps)", False, 'from absl import logging\n'), ((462, 25, 462, 55), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((465, 33, 465, 71), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((499, 12, 499, 59), 'absl.logging.info', 'logging.info', ({(499, 25, 499, 58): '"""Finished embedding the dataset."""'}, {}), "('Finished embedding the dataset.')", False, 'from absl import logging\n'), ((596, 16, 596, 67), 'tensorflow.config.experimental.set_memory_growth', 'tf.config.experimental.set_memory_growth', ({(596, 57, 596, 60): 'gpu', (596, 62, 596, 66): '(True)'}, {}), '(gpu, True)', True, 'import tensorflow as tf\n'), ((602, 12, 602, 27), 'absl.logging.info', 'logging.info', ({(602, 25, 602, 26): 'e'}, {}), '(e)', False, 'from absl import logging\n'), ((79, 23, 80, 41), 'tensorflow.unstack', 'tf.unstack', (), '', True, 'import tensorflow as tf\n'), ((346, 52, 346, 93), 'numpy.arange', 'np.arange', ({(346, 62, 346, 70): 'curr_idx', (346, 72, 346, 92): 'curr_idx + num_steps'}, {}), '(curr_idx, curr_idx + num_steps)', True, 'import numpy as np\n'), ((443, 18, 443, 38), 'numpy.concatenate', 'np.concatenate', ({(443, 33, 443, 37): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((92, 40, 92, 72), 'tensorflow.argmin', 'tf.argmin', ({(92, 50, 92, 71): 'mean_squared_distance'}, {}), '(mean_squared_distance)', True, 'import tensorflow as tf\n'), ((193, 43, 199, 24), 'tensorflow.train.polynomial_decay', 'tf.train.polynomial_decay', (), '', True, 'import tensorflow as tf\n'), ((275, 55, 275, 69), 'config.CONFIG.items', 'CONFIG.items', ({}, {}), '()', False, 'from config import CONFIG\n'), ((423, 39, 423, 66), 'tensorflow.gather', 'tf.gather', (), '', True, 'import tensorflow as tf\n'), ((470, 37, 470, 51), 'numpy.isnan', 'np.isnan', ({(470, 46, 470, 50): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((472, 27, 472, 41), 'numpy.isnan', 'np.isnan', ({(472, 36, 472, 40): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((473, 33, 473, 47), 'numpy.isnan', 'np.isnan', ({(473, 42, 473, 46): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((474, 27, 474, 41), 'numpy.isnan', 'np.isnan', ({(474, 36, 474, 40): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((480, 25, 480, 39), 'numpy.isnan', 'np.isnan', ({(480, 34, 480, 38): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((468, 33, 468, 47), 'numpy.isnan', 'np.isnan', ({(468, 42, 468, 46): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((476, 33, 476, 47), 'numpy.isnan', 'np.isnan', ({(476, 42, 476, 46): 'embs'}, {}), '(embs)', True, 'import numpy as np\n'), ((478, 49, 478, 63), 'numpy.isnan', 'np.isnan', ({(478, 58, 478, 62): 'embs'}, {}), '(embs)', True, 'import numpy as np\n')]