import tensorflow as tf

import numpy as np

# pylint: disable=no-name-in-module

from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file

from lungs import i3d

def assign(global_vars, model_path):

    reader = tf.pywrap_tensorflow.NewCheckpointReader(model_path)

    for var_3d in global_vars:

        if 'Logits' not in var_3d.op.name:

            var_2d_name = var_3d.op.name.replace('RGB/inception_i3d', 'InceptionV1')

            var_2d_name = var_2d_name.replace('Conv3d', 'Conv2d')

            var_2d_name = var_2d_name.replace('conv_3d/w', 'weights')

            var_2d_name = var_2d_name.replace('conv_3d/b', 'biases')

            var_2d_name = var_2d_name.replace('batch_norm', 'BatchNorm')

            var_value = reader.get_tensor(var_2d_name)

            if len(var_3d.get_shape()) - 1 == len(var_value.shape):

                n = var_value.shape[0]

                inflated = np.tile(np.expand_dims(var_value / n, 0), [n,1,1,1,1])

                var_3d.assign(tf.convert_to_tensor(inflated))

            else:

                var_3d.assign(tf.convert_to_tensor(var_value.reshape(var_3d.get_shape())))

def inflate_inception_v1_checkpoint_to_i3d(ckpt_2d, ckpt_3d, num_slices=145, volume_size=224, num_classes=2):

    '''

    Bootstrap the filters from a pre-trained two-dimensional 

    Inception-v1 checkpoint into a three-dimensional I3D checkpoint.

    Usage example:

    inflate_inception_v1_checkpoint_to_i3d('path/to/inception_v1.ckpt', 'path/to/new/i3d_model.ckpt')

    '''

    rgb_input = tf.placeholder(tf.float32,

        shape=(1, num_slices, volume_size, volume_size, 3))

    with tf.variable_scope('RGB'):

        i3d.InceptionI3d(num_classes, spatial_squeeze=True, final_endpoint='Logits')\

                        (rgb_input, is_training=False)

    init_op = tf.compat.v1.global_variables_initializer()

    with tf.Session() as sess:

        sess.run(init_op)

        assign(tf.global_variables(), ckpt_2d)

        tf.train.Saver().save(sess, ckpt_3d)

def inspect_checkpoint(ckpt):

    print_tensors_in_checkpoint_file(ckpt, all_tensors=False, tensor_name='')