import tensorflow as tf
from tensorflow.contrib.layers.python.layers.initializers import (
    variance_scaling_initializer,
)
import tensorflow.contrib.slim as slim

from src.ops import tf_pad


def get_image_encoder(model_type='resnet'):
    """
    Retrieves encoder fn for image and 3D
    """
    models = {
        'resnet': encoder_resnet
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown image encoder:', model_type)
        exit(1)


def get_temporal_encoder(model_type='AZ_FC2GN'):
    models = {
        'AZ_FC2GN': az_fc2_groupnorm,
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown temporal encoder:', model_type)
        exit(1)


def get_prediction_model(model_type='AZ_FC2GN'):
    models = {
        'AZ_FC2GN': az_fc2_groupnorm,
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown prediction model:', model_type)
        exit(1)


# Functions for image encoder.

def encoder_resnet(x, is_training=True, weight_decay=0.001, reuse=False):
    """
    Resnet v2-50
    Assumes input is [batch, height_in, width_in, channels]!!
    Input:
    - x: N x H x W x 3
    - weight_decay: float
    - reuse: bool->True if test

    Outputs:
    - cam: N x 3
    - Pose vector: N x 72
    - Shape vector: N x 10
    - variables: tf variables
    """
    from tensorflow.contrib.slim.python.slim.nets import resnet_v2
    with tf.name_scope('Encoder_resnet', values=[x]):
        with slim.arg_scope(
                resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
            net, end_points = resnet_v2.resnet_v2_50(
                x,
                num_classes=None,
                is_training=is_training,
                reuse=reuse,
                scope='resnet_v2_50')
            net = tf.squeeze(net, axis=[1, 2])
    variables_scope = 'resnet_v2_50'
    return net, variables_scope


def encoder_fc3_dropout(x,
                        num_output=85,
                        is_training=True,
                        reuse=False,
                        name='3D_module'):
    """
    3D inference module. 3 MLP layers (last is the output)
    With dropout  on first 2.
    Input:
    - x: N x [|img_feat|, |3D_param|]
    - reuse: bool

    Outputs:
    - 3D params: N x num_output
      if orthogonal:
           either 85: (3 + 24*3 + 10) or 109 (3 + 24*4 + 10) for factored
           axis-angle representation
      if perspective:
          86: (f, tx, ty, tz) + 24*3 + 10, or 110 for factored axis-angle.
    - variables: tf variables
    """
    with tf.variable_scope(name, reuse=reuse) as scope:
        net = slim.fully_connected(x, 1024, scope='fc1')
        net = slim.dropout(net, 0.5, is_training=is_training, scope='dropout1')
        net = slim.fully_connected(net, 1024, scope='fc2')
        net = slim.dropout(net, 0.5, is_training=is_training, scope='dropout2')
        small_xavier = variance_scaling_initializer(
            factor=.01, mode='FAN_AVG', uniform=True)
        net = slim.fully_connected(
            net,
            num_output,
            activation_fn=None,
            weights_initializer=small_xavier,
            scope='fc3')

    variables = tf.contrib.framework.get_variables(scope)
    return net, variables


# Functions for f_{movie strip} and f_{AR}.

def az_fc2_groupnorm(net, num_conv_layers, prefix, reuse=None):
    """
    Causal architecture for movie strip encoder and autoregressive prediction.

    Each block has 2 convs:
    norm --> relu --> conv --> norm --> relu --> conv --> add.
    Uses full convolution.
    """
    for i in range(num_conv_layers):
        net = az_fc_causal_block2(
            kernel_width=3,
            name='block_{}'.format(i),
            net_input=net,
            num_filter=2048,
            pad_edges=True,
            prefix=prefix,
            reuse=reuse,
        )
    return net


def az_fc_causal_block2(net_input, num_filter, kernel_width, prefix='', name='',
                        pad_edges=False, reuse=None):
    """
    Causal res block:
    BN -> Relu -> Weight -> BN -> Relu -> Weight -> Add

    Causal is implemented by padding the left size by receptive field - 1 and
    using 'VALID' mode. The output indices thus only look at the input indices
    that are less than or equal.
    """
    pad_mode = 'EDGE' if pad_edges else 'ZERO'
    num_groups = 32

    # NTC -> NT1C
    net_input_expand = tf.expand_dims(net_input, axis=2)
    # group-norm
    net_norm = tf.contrib.layers.group_norm(
        net_input_expand,
        channels_axis=-1,
        reduction_axes=(-3, -2),
        scope=(prefix + 'AZ_FC_causal_block_preact_gn1' + name),
        reuse=reuse,
        groups=num_groups,
    )
    # relu
    net_relu = tf.nn.relu(net_norm)
    # weight
    net_relu_padded = tf_pad(
        tensor=net_relu,
        paddings=[[0, 0], [kernel_width - 1, 0], [0, 0], [0, 0]],
        mode=pad_mode
    )
    net_conv1 = tf.contrib.layers.conv2d(
        inputs=net_relu_padded,
        num_outputs=num_filter,
        kernel_size=[kernel_width, 1],
        stride=1,
        padding='VALID',
        data_format='NHWC',  # was previously 'NCHW',
        rate=1,
        activation_fn=None,
        scope=prefix + 'AZ_FC_causal_block2_conv1' + name,
        reuse=reuse,
    )
    # group-norm
    net_norm2 = tf.contrib.layers.group_norm(
        net_conv1,
        channels_axis=-1,
        reduction_axes=(-3, -2),
        reuse=reuse,
        groups=num_groups,
        scope=(prefix + 'AZ_FC_causal_block_preact_gn2' + name),
    )
    # relu
    net_relu2 = tf.nn.relu(net_norm2)
    net_relu_padded2 = tf_pad(
        tensor=net_relu2,
        paddings=[[0, 0], [kernel_width - 1, 0], [0, 0], [0, 0]],
        mode=pad_mode
    )
    # Initalization
    small_xavier = variance_scaling_initializer(
        factor=.001,
        mode='FAN_AVG',
        uniform=True
    )
    net_final = tf.contrib.layers.conv2d(
        inputs=net_relu_padded2,
        num_outputs=num_filter,
        kernel_size=[kernel_width, 1],
        stride=1,
        padding='VALID',
        data_format='NHWC',
        rate=1,
        activation_fn=None,
        weights_initializer=small_xavier,
        scope=prefix + 'AZ_FC_causal_block2_conv2' + name,
        reuse=reuse,
    )
    # NT1C -> NTC
    net_final = tf.squeeze(net_final, axis=2)
    # skip connection
    residual = tf.add(net_final, net_input)
    return residual


# Functions for f_{3D}.

def batch_pred_omega(input_features, batch_size, is_training, num_output,
                     omega_mean, sequence_length, scope, predict_delta_keys=(),
                     use_delta_from_pred=False, use_optcam=False):
    """
    Given B x T x * inputs, computes IEF on them by batching them
    as BT x *.

    if use_optcam is True, only outputs 72 or 82 dims.
    and appends fixed camera [1,0,0]
    """
    # run in batch
    # omega_mean comes in as shape: BT x 85
    input_features_reshape = tf.reshape(input_features,
                                        (batch_size * sequence_length, -1))
    omega_pred, delta_predictions = call_hmr_ief(
        phi=input_features_reshape,
        omega_start=omega_mean,
        scope=scope,
        num_output=num_output,
        is_training=is_training,
        predict_delta_keys=predict_delta_keys,
        use_delta_from_pred=use_delta_from_pred,
        use_optcam=use_optcam,
    )
    omega_pred = tf.reshape(
        omega_pred,
        (batch_size, sequence_length, num_output)
    )
    new_delta_predictions = {}
    for delta_t, prediction in delta_predictions.items():
        new_delta_predictions[delta_t] = tf.reshape(
            prediction,
            (batch_size, sequence_length, num_output)
        )
    return omega_pred, new_delta_predictions


def call_hmr_ief(phi, omega_start, scope, num_output=85, num_stage=3,
                 is_training=True, predict_delta_keys=(),
                 use_delta_from_pred=False, use_optcam=True):
    """
    Wrapper for doing HMR-style IEF.

    If predict_delta, then also makes num_delta_t predictions forward and
    backward in time, with each step of delta_t.

    Args:
        phi (Bx2048): Image features.
        omega_start (Bx85): Starting Omega as input to first IEF.
        scope (str): Name of scope for reuse.
        num_output (int): Size of output.
        num_stage (int): Number of iterations for IEF.
        is_training (bool): If False, don't apply dropout.
        predict_delta_keys (iterable): List of keys for delta_t.
        use_delta_from_pred (bool): If True, initializes delta prediction from
            current frame prediction.
        use_optcam (bool): If True, uses [1, 0, 0] for cam.

    Returns:
        Final theta (Bx{num_output})
        Deltas predictions (List of outputs)
    """
    theta_here = hmr_ief(
        phi=phi,
        omega_start=omega_start,
        scope=scope,
        num_output=num_output,
        num_stage=num_stage,
        is_training=is_training
    )

    # Delta only needs to do cam/pose, no shape!
    if use_optcam:
        num_output_delta = 72
    else:
        num_output_delta = 3 + 72

    deltas_predictions = {}
    for delta_t in predict_delta_keys:
        if delta_t == 0:
            # This should just be the normal IEF.
            continue
        elif delta_t > 0:
            scope_delta = scope + '_future{}'.format(delta_t)
        elif delta_t < 0:
            scope_delta = scope + '_past{}'.format(abs(delta_t))

        omega_start_delta = theta_here if use_delta_from_pred else omega_start
        # append this later.
        beta = omega_start_delta[:, -10:]

        if use_optcam:
            # trim the first 3D camera + last shpae
            omega_start_delta = omega_start_delta[:, 3:3 + num_output_delta]
        else:
            omega_start_delta = omega_start_delta[:, :num_output_delta]

        delta_pred = hmr_ief(
            phi=phi,
            omega_start=omega_start_delta,
            scope=scope_delta,
            num_output=num_output_delta,
            num_stage=num_stage,
            is_training=is_training
        )
        if use_optcam:
            # Add camera + shape
            scale = tf.ones([delta_pred.shape[0], 1])
            trans = tf.zeros([delta_pred.shape[0], 2])
            delta_pred = tf.concat([scale, trans, delta_pred, beta], 1)
        else:
            delta_pred = tf.concat([delta_pred[:, :75], beta], 1)

        deltas_predictions[delta_t] = delta_pred

    return theta_here, deltas_predictions


def hmr_ief(phi, omega_start, scope, num_output=85, num_stage=3,
            is_training=True):
    """
    Runs HMR-style IEF.

    Args:
        phi (Bx2048): Image features.
        omega_start (Bx85): Starting Omega as input to first IEF.
        scope (str): Name of scope for reuse.
        num_output (int): Size of output.
        num_stage (int): Number of iterations for IEF.
        is_training (bool): If False, don't apply dropout.

    Returns:
        Final theta (Bx{num_output})
    """
    with tf.variable_scope(scope):
        theta_prev = omega_start
        theta_here = None

        for _ in range(num_stage):
            # ---- Compute outputs
            state = tf.concat([phi, theta_prev], 1)
            delta_theta, _ = encoder_fc3_dropout(
                state,
                is_training=is_training,
                num_output=num_output,
                reuse=tf.AUTO_REUSE
            )
            # Compute new theta
            theta_here = theta_prev + delta_theta

            # Finally update to end iteration.
            theta_prev = theta_here

    return theta_here

import tensorflow as tf
from tensorflow.contrib.layers.python.layers.initializers import (
    variance_scaling_initializer,
)
import tensorflow.contrib.slim as slim

from src.ops import tf_pad


def get_image_encoder(model_type='resnet'):
    """
    Retrieves encoder fn for image and 3D
    """
    models = {
        'resnet': encoder_resnet
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown image encoder:', model_type)
        exit(1)


def get_temporal_encoder(model_type='AZ_FC2GN'):
    models = {
        'AZ_FC2GN': az_fc2_groupnorm,
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown temporal encoder:', model_type)
        exit(1)


def get_prediction_model(model_type='AZ_FC2GN'):
    models = {
        'AZ_FC2GN': az_fc2_groupnorm,
    }
    if model_type in models.keys():
        return models[model_type]
    else:
        print('Unknown prediction model:', model_type)
        exit(1)


# Functions for image encoder.

def encoder_resnet(x, is_training=True, weight_decay=0.001, reuse=False):
    """
    Resnet v2-50
    Assumes input is [batch, height_in, width_in, channels]!!
    Input:
    - x: N x H x W x 3
    - weight_decay: float
    - reuse: bool->True if test

    Outputs:
    - cam: N x 3
    - Pose vector: N x 72
    - Shape vector: N x 10
    - variables: tf variables
    """
    from tensorflow.contrib.slim.python.slim.nets import resnet_v2
    with tf.name_scope('Encoder_resnet', values=[x]):
        with slim.arg_scope(
                resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
            net, end_points = resnet_v2.resnet_v2_50(
                x,
                num_classes=None,
                is_training=is_training,
                reuse=reuse,
                scope='resnet_v2_50')
            net = tf.squeeze(net, axis=[1, 2])
    variables_scope = 'resnet_v2_50'
    return net, variables_scope


def encoder_fc3_dropout(x,
                        num_output=85,
                        is_training=True,
                        reuse=False,
                        name='3D_module'):
    """
    3D inference module. 3 MLP layers (last is the output)
    With dropout  on first 2.
    Input:
    - x: N x [|img_feat|, |3D_param|]
    - reuse: bool

    Outputs:
    - 3D params: N x num_output
      if orthogonal:
           either 85: (3 + 24*3 + 10) or 109 (3 + 24*4 + 10) for factored
           axis-angle representation
      if perspective:
          86: (f, tx, ty, tz) + 24*3 + 10, or 110 for factored axis-angle.
    - variables: tf variables
    """
    with tf.variable_scope(name, reuse=reuse) as scope:
        net = slim.fully_connected(x, 1024, scope='fc1')
        net = slim.dropout(net, 0.5, is_training=is_training, scope='dropout1')
        net = slim.fully_connected(net, 1024, scope='fc2')
        net = slim.dropout(net, 0.5, is_training=is_training, scope='dropout2')
        small_xavier = variance_scaling_initializer(
            factor=.01, mode='FAN_AVG', uniform=True)
        net = slim.fully_connected(
            net,
            num_output,
            activation_fn=None,
            weights_initializer=small_xavier,
            scope='fc3')

    variables = tf.contrib.framework.get_variables(scope)
    return net, variables


# Functions for f_{movie strip} and f_{AR}.

def az_fc2_groupnorm(net, num_conv_layers, prefix, reuse=None):
    """
    Causal architecture for movie strip encoder and autoregressive prediction.

    Each block has 2 convs:
    norm --> relu --> conv --> norm --> relu --> conv --> add.
    Uses full convolution.
    """
    for i in range(num_conv_layers):
        net = az_fc_causal_block2(
            kernel_width=3,
            name='block_{}'.format(i),
            net_input=net,
            num_filter=2048,
            pad_edges=True,
            prefix=prefix,
            reuse=reuse,
        )
    return net


def az_fc_causal_block2(net_input, num_filter, kernel_width, prefix='', name='',
                        pad_edges=False, reuse=None):
    """
    Causal res block:
    BN -> Relu -> Weight -> BN -> Relu -> Weight -> Add

    Causal is implemented by padding the left size by receptive field - 1 and
    using 'VALID' mode. The output indices thus only look at the input indices
    that are less than or equal.
    """
    pad_mode = 'EDGE' if pad_edges else 'ZERO'
    num_groups = 32

    # NTC -> NT1C
    net_input_expand = tf.expand_dims(net_input, axis=2)
    # group-norm
    net_norm = tf.contrib.layers.group_norm(
        net_input_expand,
        channels_axis=-1,
        reduction_axes=(-3, -2),
        scope=(prefix + 'AZ_FC_causal_block_preact_gn1' + name),
        reuse=reuse,
        groups=num_groups,
    )
    # relu
    net_relu = tf.nn.relu(net_norm)
    # weight
    net_relu_padded = tf_pad(
        tensor=net_relu,
        paddings=[[0, 0], [kernel_width - 1, 0], [0, 0], [0, 0]],
        mode=pad_mode
    )
    net_conv1 = tf.contrib.layers.conv2d(
        inputs=net_relu_padded,
        num_outputs=num_filter,
        kernel_size=[kernel_width, 1],
        stride=1,
        padding='VALID',
        data_format='NHWC',  # was previously 'NCHW',
        rate=1,
        activation_fn=None,
        scope=prefix + 'AZ_FC_causal_block2_conv1' + name,
        reuse=reuse,
    )
    # group-norm
    net_norm2 = tf.contrib.layers.group_norm(
        net_conv1,
        channels_axis=-1,
        reduction_axes=(-3, -2),
        reuse=reuse,
        groups=num_groups,
        scope=(prefix + 'AZ_FC_causal_block_preact_gn2' + name),
    )
    # relu
    net_relu2 = tf.nn.relu(net_norm2)
    net_relu_padded2 = tf_pad(
        tensor=net_relu2,
        paddings=[[0, 0], [kernel_width - 1, 0], [0, 0], [0, 0]],
        mode=pad_mode
    )
    # Initalization
    small_xavier = variance_scaling_initializer(
        factor=.001,
        mode='FAN_AVG',
        uniform=True
    )
    net_final = tf.contrib.layers.conv2d(
        inputs=net_relu_padded2,
        num_outputs=num_filter,
        kernel_size=[kernel_width, 1],
        stride=1,
        padding='VALID',
        data_format='NHWC',
        rate=1,
        activation_fn=None,
        weights_initializer=small_xavier,
        scope=prefix + 'AZ_FC_causal_block2_conv2' + name,
        reuse=reuse,
    )
    # NT1C -> NTC
    net_final = tf.squeeze(net_final, axis=2)
    # skip connection
    residual = tf.add(net_final, net_input)
    return residual


# Functions for f_{3D}.

def batch_pred_omega(input_features, batch_size, is_training, num_output,
                     omega_mean, sequence_length, scope, predict_delta_keys=(),
                     use_delta_from_pred=False, use_optcam=False):
    """
    Given B x T x * inputs, computes IEF on them by batching them
    as BT x *.

    if use_optcam is True, only outputs 72 or 82 dims.
    and appends fixed camera [1,0,0]
    """
    # run in batch
    # omega_mean comes in as shape: BT x 85
    input_features_reshape = tf.reshape(input_features,
                                        (batch_size * sequence_length, -1))
    omega_pred, delta_predictions = call_hmr_ief(
        phi=input_features_reshape,
        omega_start=omega_mean,
        scope=scope,
        num_output=num_output,
        is_training=is_training,
        predict_delta_keys=predict_delta_keys,
        use_delta_from_pred=use_delta_from_pred,
        use_optcam=use_optcam,
    )
    omega_pred = tf.reshape(
        omega_pred,
        (batch_size, sequence_length, num_output)
    )
    new_delta_predictions = {}
    for delta_t, prediction in delta_predictions.items():
        new_delta_predictions[delta_t] = tf.reshape(
            prediction,
            (batch_size, sequence_length, num_output)
        )
    return omega_pred, new_delta_predictions


def call_hmr_ief(phi, omega_start, scope, num_output=85, num_stage=3,
                 is_training=True, predict_delta_keys=(),
                 use_delta_from_pred=False, use_optcam=True):
    """
    Wrapper for doing HMR-style IEF.

    If predict_delta, then also makes num_delta_t predictions forward and
    backward in time, with each step of delta_t.

    Args:
        phi (Bx2048): Image features.
        omega_start (Bx85): Starting Omega as input to first IEF.
        scope (str): Name of scope for reuse.
        num_output (int): Size of output.
        num_stage (int): Number of iterations for IEF.
        is_training (bool): If False, don't apply dropout.
        predict_delta_keys (iterable): List of keys for delta_t.
        use_delta_from_pred (bool): If True, initializes delta prediction from
            current frame prediction.
        use_optcam (bool): If True, uses [1, 0, 0] for cam.

    Returns:
        Final theta (Bx{num_output})
        Deltas predictions (List of outputs)
    """
    theta_here = hmr_ief(
        phi=phi,
        omega_start=omega_start,
        scope=scope,
        num_output=num_output,
        num_stage=num_stage,
        is_training=is_training
    )

    # Delta only needs to do cam/pose, no shape!
    if use_optcam:
        num_output_delta = 72
    else:
        num_output_delta = 3 + 72

    deltas_predictions = {}
    for delta_t in predict_delta_keys:
        if delta_t == 0:
            # This should just be the normal IEF.
            continue
        elif delta_t > 0:
            scope_delta = scope + '_future{}'.format(delta_t)
        elif delta_t < 0:
            scope_delta = scope + '_past{}'.format(abs(delta_t))

        omega_start_delta = theta_here if use_delta_from_pred else omega_start
        # append this later.
        beta = omega_start_delta[:, -10:]

        if use_optcam:
            # trim the first 3D camera + last shpae
            omega_start_delta = omega_start_delta[:, 3:3 + num_output_delta]
        else:
            omega_start_delta = omega_start_delta[:, :num_output_delta]

        delta_pred = hmr_ief(
            phi=phi,
            omega_start=omega_start_delta,
            scope=scope_delta,
            num_output=num_output_delta,
            num_stage=num_stage,
            is_training=is_training
        )
        if use_optcam:
            # Add camera + shape
            scale = tf.ones([delta_pred.shape[0], 1])
            trans = tf.zeros([delta_pred.shape[0], 2])
            delta_pred = tf.concat([scale, trans, delta_pred, beta], 1)
        else:
            delta_pred = tf.concat([delta_pred[:, :75], beta], 1)

        deltas_predictions[delta_t] = delta_pred

    return theta_here, deltas_predictions


def hmr_ief(phi, omega_start, scope, num_output=85, num_stage=3,
            is_training=True):
    """
    Runs HMR-style IEF.

    Args:
        phi (Bx2048): Image features.
        omega_start (Bx85): Starting Omega as input to first IEF.
        scope (str): Name of scope for reuse.
        num_output (int): Size of output.
        num_stage (int): Number of iterations for IEF.
        is_training (bool): If False, don't apply dropout.

    Returns:
        Final theta (Bx{num_output})
    """
    with tf.variable_scope(scope):
        theta_prev = omega_start
        theta_here = None

        for _ in range(num_stage):
            # ---- Compute outputs
            state = tf.concat([phi, theta_prev], 1)
            delta_theta, _ = encoder_fc3_dropout(
                state,
                is_training=is_training,
                num_output=num_output,
                reuse=tf.AUTO_REUSE
            )
            # Compute new theta
            theta_here = theta_prev + delta_theta

            # Finally update to end iteration.
            theta_prev = theta_here

    return theta_here

Retrieves encoder fn for image and 3D # Functions for image encoder. Resnet v2-50     Assumes input is [batch, height_in, width_in, channels]!!     Input:     - x: N x H x W x 3     - weight_decay: float     - reuse: bool->True if test      Outputs:     - cam: N x 3     - Pose vector: N x 72     - Shape vector: N x 10     - variables: tf variables 3D inference module. 3 MLP layers (last is the output)     With dropout  on first 2.     Input:     - x: N x [|img_feat|, |3D_param|]     - reuse: bool      Outputs:     - 3D params: N x num_output       if orthogonal:            either 85: (3 + 24*3 + 10) or 109 (3 + 24*4 + 10) for factored            axis-angle representation       if perspective:           86: (f, tx, ty, tz) + 24*3 + 10, or 110 for factored axis-angle.     - variables: tf variables # Functions for f_{movie strip} and f_{AR}. Causal architecture for movie strip encoder and autoregressive prediction.      Each block has 2 convs:     norm --> relu --> conv --> norm --> relu --> conv --> add.     Uses full convolution. Causal res block:     BN -> Relu -> Weight -> BN -> Relu -> Weight -> Add      Causal is implemented by padding the left size by receptive field - 1 and     using 'VALID' mode. The output indices thus only look at the input indices     that are less than or equal. # NTC -> NT1C # group-norm # relu # weight # was previously 'NCHW', # group-norm # relu # Initalization # NT1C -> NTC # skip connection # Functions for f_{3D}. Given B x T x * inputs, computes IEF on them by batching them     as BT x *.      if use_optcam is True, only outputs 72 or 82 dims.     and appends fixed camera [1,0,0] # run in batch # omega_mean comes in as shape: BT x 85 Wrapper for doing HMR-style IEF.      If predict_delta, then also makes num_delta_t predictions forward and     backward in time, with each step of delta_t.      Args:         phi (Bx2048): Image features.         omega_start (Bx85): Starting Omega as input to first IEF.         scope (str): Name of scope for reuse.         num_output (int): Size of output.         num_stage (int): Number of iterations for IEF.         is_training (bool): If False, don't apply dropout.         predict_delta_keys (iterable): List of keys for delta_t.         use_delta_from_pred (bool): If True, initializes delta prediction from             current frame prediction.         use_optcam (bool): If True, uses [1, 0, 0] for cam.      Returns:         Final theta (Bx{num_output})         Deltas predictions (List of outputs) # Delta only needs to do cam/pose, no shape! # This should just be the normal IEF. # append this later. # trim the first 3D camera + last shpae # Add camera + shape Runs HMR-style IEF.      Args:         phi (Bx2048): Image features.         omega_start (Bx85): Starting Omega as input to first IEF.         scope (str): Name of scope for reuse.         num_output (int): Size of output.         num_stage (int): Number of iterations for IEF.         is_training (bool): If False, don't apply dropout.      Returns:         Final theta (Bx{num_output}) # ---- Compute outputs # Compute new theta # Finally update to end iteration.