#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""Created on Thu Nov  1 17:04:52 2018.

@author: smylonas

Modified on Sat Feb 6 14:30:00 2026 by abdollahegazy
to include TF2 compatibility.
"""

import collections
import numpy as np

from tf_slim import arg_scope
from tf_slim import add_arg_scope
from tf_slim.layers import initializers
from tf_slim import layers
from tf_slim import l2_regularizer
from tf_slim.layers import utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import variable_scope
from tensorflow.compat.v1 import GraphKeys # type: ignore

class Block(collections.namedtuple("Block", ["scope", "unit_fn", "args"])):
    """A named tuple describing a ResNet block.

    Its parts are:
      scope: The scope of the `Block`.
      unit_fn: The ResNet unit function which takes as input a `Tensor` and
        returns another `Tensor` with the output of the ResNet unit.
      args: A list of length equal to the number of units in the `Block`. The list
        contains one (depth, depth_bottleneck, stride) tuple for each unit in the
        block to serve as argument to unit_fn.
    """


def subsample(inputs, factor, scope=None):
    """Subsamples the input along the spatial dimensions.

    Args:
      inputs: A `Tensor` of size [batch, height_in, width_in, channels].
      factor: The subsampling factor.
      scope: Optional variable_scope.

    Returns:
      output: A `Tensor`with the input, either intact (if factor == 1) or subsampled (if factor > 1).

    """
    if factor == 1:
        return inputs
    else:
        return layers.max_pool3d(inputs, 1, stride=factor, scope=scope)


def conv3d_same(
    inputs,
    num_outputs,
    kernel_size,
    stride,
    rate=1,
    scope=None,
    outputs_collections=None,
):
    """Strided 3-D convolution with 'SAME' padding.

    When stride > 1, then we do explicit zero-padding, followed by conv2d with
    'VALID' padding.

    Note that

       net = conv2d_same(inputs, num_outputs, 3, stride=stride)

    is equivalent to

       net = tf.contrib.layers.conv2d(inputs, num_outputs, 3, stride=1,
       padding='SAME')
       net = subsample(net, factor=stride)

    whereas

       net = tf.contrib.layers.conv2d(inputs, num_outputs, 3, stride=stride,
       padding='SAME')

    is different when the input's height or width is even, which is why we add the
    current function. For more details, see ResnetUtilsTest.testConv2DSameEven().

    Args:
      inputs: A 4-D tensor of size [batch, height_in, width_in, channels].
      num_outputs: An integer, the number of output filters.
      kernel_size: An int with the kernel_size of the filters.
      stride: An integer, the output stride.
      rate: An integer, rate for atrous convolution.
      scope: Scope.

    Returns:
      output: A 4-D tensor of size [batch, height_out, width_out, channels] with
        the convolution output.

    """
    if stride == 1:
        return layers.conv3d(
            inputs,
            num_outputs,
            kernel_size,
            stride=1,
            rate=rate,
            padding="SAME",
            scope=scope,
            outputs_collections=None,
        )
    else:
        if type(kernel_size) is list:
            kernel = np.array(kernel_size)
            pad_total = kernel - 1
            pad_beg = pad_total // 2
            pad_end = pad_total - pad_beg
            inputs = array_ops.pad(
                inputs,
                [
                    [0, 0],
                    [pad_beg[0], pad_end[0]],
                    [pad_beg[1], pad_end[1]],
                    [pad_beg[2], pad_end[2]],
                    [0, 0],
                ],
            )
        else:
            kernel_size_effective = kernel_size + (kernel_size - 1) * (rate - 1)
            pad_total = kernel_size_effective - 1
            pad_beg = pad_total // 2
            pad_end = pad_total - pad_beg
            inputs = array_ops.pad(
                inputs,
                [
                    [0, 0],
                    [pad_beg, pad_end],
                    [pad_beg, pad_end],
                    [pad_beg, pad_end],
                    [0, 0],
                ],
            )

        return layers.conv3d(
            inputs,
            num_outputs,
            kernel_size,
            stride=stride,
            rate=rate,
            padding="VALID",
            scope=scope,
            outputs_collections=None,
        )


@add_arg_scope
def stack_blocks_dense(net, blocks, output_stride=None, outputs_collections=None):
    """Stacks ResNet `Blocks` and controls output feature density.

    First, this function creates scopes for the ResNet in the form of
    'block_name/unit_1', 'block_name/unit_2', etc.

    Second, this function allows the user to explicitly control the ResNet
    output_stride, which is the ratio of the input to output spatial resolution.
    This is useful for dense prediction tasks such as semantic segmentation or
    object detection.

    Most ResNets consist of 4 ResNet blocks and subsample the activations by a
    factor of 2 when transitioning between consecutive ResNet blocks. This results
    to a nominal ResNet output_stride equal to 8. If we set the output_stride to
    half the nominal network stride (e.g., output_stride=4), then we compute
    responses twice.

    Control of the output feature density is implemented by atrous convolution.

    Args:
      net: A `Tensor` of size [batch, height, width, channels].
      blocks: A list of length equal to the number of ResNet `Blocks`. Each
        element is a ResNet `Block` object describing the units in the `Block`.
      output_stride: If `None`, then the output will be computed at the nominal
        network stride. If output_stride is not `None`, it specifies the requested
        ratio of input to output spatial resolution, which needs to be equal to
        the product of unit strides from the start up to some level of the ResNet.
        For example, if the ResNet employs units with strides 1, 2, 1, 3, 4, 1,
        then valid values for the output_stride are 1, 2, 6, 24 or None (which
        is equivalent to output_stride=24).
      outputs_collections: Collection to add the ResNet block outputs.

    Returns:
      net: Output tensor with stride equal to the specified output_stride.

    Raises:
      ValueError: If the target output_stride is not valid.

    """
    # The current_stride variable keeps track of the effective stride of the
    # activations. This allows us to invoke atrous convolution whenever applying
    # the next residual unit would result in the activations having stride larger
    # than the target output_stride.
    current_stride = 1

    # The atrous convolution rate parameter.
    rate = 1

    for block in blocks:
        with variable_scope.variable_scope(block.scope, "block", [net]) as sc:
            for i, unit in enumerate(block.args):
                if output_stride is not None and current_stride > output_stride:
                    raise ValueError("The target output_stride cannot be reached.")

                with variable_scope.variable_scope("unit_%d" % (i + 1), values=[net]):
                    # If we have reached the target output_stride, then we need to employ
                    # atrous convolution with stride=1 and multiply the atrous rate by the
                    # current unit's stride for use in subsequent layers.
                    if output_stride is not None and current_stride == output_stride:
                        net = block.unit_fn(net, rate=rate, **dict(unit, stride=1))
                        rate *= unit.get("stride", 1)

                    else:
                        net = block.unit_fn(net, rate=1, **unit)
                        current_stride *= unit.get("stride", 1)
            net = utils.collect_named_outputs(outputs_collections, sc.name, net)

    if output_stride is not None and current_stride != output_stride:
        raise ValueError("The target output_stride cannot be reached.")

    return net


def resnet_arg_scope(
    is_training=True,
    weight_decay=0.0001,
    batch_norm_decay=0.997,
    batch_norm_epsilon=0.00001,
    batch_norm_scale=True,
):
    """Defines the default ResNet arg scope.

    Args:
      is_training: Whether or not we are training the parameters in the batch
        normalization layers of the model. (deprecated)
      weight_decay: The weight decay to use for regularizing the model.
      batch_norm_decay: The moving average decay when estimating layer activation
        statistics in batch normalization.
      batch_norm_epsilon: Small constant to prevent division by zero when
        normalizing activations by their variance in batch normalization.
      batch_norm_scale: If True, uses an explicit `gamma` multiplier to scale the
        activations in the batch normalization layer.

    Returns:
      An `arg_scope` to use for the resnet models.

    """
    batch_norm_params = {
        "is_training": is_training,
        "decay": batch_norm_decay,
        "epsilon": batch_norm_epsilon,
        "scale": batch_norm_scale,
        "updates_collections": GraphKeys.UPDATE_OPS,
        #'updates_collections': None,
        #'renorm': True
    }

    with arg_scope(
        [layers.conv3d],
        weights_regularizer=l2_regularizer(weight_decay),
        weights_initializer=initializers.variance_scaling_initializer(),
        activation_fn=nn_ops.relu,
        normalizer_fn=layers.batch_norm,
    ):
        with arg_scope([layers.batch_norm], **batch_norm_params):
            with arg_scope([layers.max_pool3d], padding="SAME") as arg_sc:
                return arg_sc


def resnet_pre_arg_scope(  # is_training=True,
    weight_decay=0.0001,
    batch_norm_decay=0.997,
    batch_norm_epsilon=1e-5,
    batch_norm_scale=True,
):

    batch_norm_params = {
        #'is_training': is_training,
        "decay": batch_norm_decay,
        "epsilon": batch_norm_epsilon,
        "scale": batch_norm_scale,
        # 'activation_fn': nn_ops.relu,
        "variables_collections": ["batch_norm_params"],
        "updates_collections": GraphKeys.UPDATE_OPS,
    }

    with arg_scope(
        [layers.conv3d],
        weights_regularizer=l2_regularizer(weight_decay),
        weights_initializer=initializers.variance_scaling_initializer(),
        activation_fn=None,
    ):
        with arg_scope([layers.batch_norm], **batch_norm_params):
            with arg_scope([layers.max_pool3d], padding="SAME") as arg_sc:
                return arg_sc
