import logging, os
logging.disable(logging.WARNING)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import tensorflow as tf
from basic_ops import *


"""This script defines non-attention same-, up-, down- modules.
Note that pre-activation is used for residual-like blocks.
Note that the residual block could be used for downsampling.
"""


def res_block(inputs, output_filters, training, dimension, name):
    """Standard residual block with pre-activation.

    Args:
        inputs: a Tensor with shape [batch, (d,) h, w, channels]
        output_filters: an integer
        training: a boolean for batch normalization and dropout
        dimension: a string, dimension of inputs/outputs -- 2D, 3D
        name: a string
        
    Returns:
        A Tensor of shape [batch, (_d,) _h, _w, output_filters]
    """
    if dimension == '2D':
        convolution = convolution_2D
        kernel_size = 3
    elif dimension == '3D':
        convolution = convolution_3D
        kernel_size = 3
    else:
        raise ValueError("Dimension (%s) must be 2D or 3D." % (dimension))

    with tf.variable_scope(name):
        if inputs.shape[-1] == output_filters:
            shortcut = inputs
            inputs = batch_norm(inputs, training, 'batch_norm_1')
            inputs = relu(inputs, 'relu_1')
        else:
            inputs = batch_norm(inputs, training, 'batch_norm_1')
            inputs = relu(inputs, 'relu_1')
            shortcut = convolution(inputs, output_filters, 1, 1, False, 'projection_shortcut')
        inputs = convolution(inputs, output_filters, kernel_size, 1, False, 'convolution_1')
        inputs = batch_norm(inputs, training, 'batch_norm_2')
        inputs = relu(inputs, 'relu_2')
        inputs = convolution(inputs, output_filters, kernel_size, 1, False, 'convolution_2')
        return tf.add(shortcut, inputs)


def down_res_block(inputs, output_filters, training, dimension, name):
    """Standard residual block with pre-activation for downsampling."""
    if dimension == '2D':
        convolution = convolution_2D
        projection_shortcut = convolution_2D
    elif dimension == '3D':
        convolution = convolution_3D
        projection_shortcut = convolution_3D
    else:
        raise ValueError("Dimension (%s) must be 2D or 3D." % (dimension))

    with tf.variable_scope(name):
        # The projection_shortcut should come after the first batch norm and ReLU.
        inputs = batch_norm(inputs, training, 'batch_norm_1')
        inputs = relu(inputs, 'relu_1')
        shortcut = projection_shortcut(inputs, output_filters, 1, 2, False, 'projection_shortcut')
        inputs = convolution(inputs, output_filters, 2, 2, False, 'convolution_1')
        inputs = batch_norm(inputs, training, 'batch_norm_2')
        inputs = relu(inputs, 'relu_2')
        inputs = convolution(inputs, output_filters, 3, 1, False, 'convolution_2')
        return tf.add(shortcut, inputs)

def down_convolution(inputs, output_filters, training, dimension, name):
    """Use a single stride 2 convolution for downsampling."""
    if dimension == '2D':
        convolution = convolution_2D
        pool = tf.layers.max_pooling2d
    elif dimension == '3D':
        convolution = convolution_3D
        pool = tf.layers.max_pooling3d
    else:
        raise ValueError("Dimension (%s) must be 2D or 3D." % (dimension))

    with tf.variable_scope(name):
        inputs = convolution(inputs, output_filters, 2, 2, True, 'convolution')
        return inputs

def up_transposed_convolution(inputs, output_filters, training, dimension, name):
    """Use a single stride 2 transposed convolution for upsampling."""
    if dimension == '2D':
        transposed_convolution = transposed_convolution_2D
    elif dimension == '3D':
        transposed_convolution = transposed_convolution_3D
    else:
        raise ValueError("Dimension (%s) must be 2D or 3D." % (dimension))

    with tf.variable_scope(name):
        inputs = transposed_convolution(inputs, output_filters, 2, 2, True, 'transposed_convolution')
        return inputs