"""Helper methods for building CNN architecture."""

import numpy
import keras
from gewittergefahr.gg_utils import error_checking

KERNEL_INITIALIZER_NAME = 'glorot_uniform'
BIAS_INITIALIZER_NAME = 'zeros'

MAX_POOLING_STRING = 'max'
MEAN_POOLING_STRING = 'mean'
VALID_POOLING_TYPE_STRINGS = [MAX_POOLING_STRING, MEAN_POOLING_STRING]

YES_PADDING_STRING = 'same'
NO_PADDING_STRING = 'valid'
VALID_PADDING_TYPE_STRINGS = [YES_PADDING_STRING, NO_PADDING_STRING]

DEFAULT_L1_WEIGHT = 0.
DEFAULT_L2_WEIGHT = 0.001

DEFAULT_ALPHA_FOR_ELU = 1.
DEFAULT_ALPHA_FOR_RELU = 0.2

ELU_FUNCTION_STRING = 'elu'
RELU_FUNCTION_STRING = 'relu'
SELU_FUNCTION_STRING = 'selu'
SOFTMAX_FUNCTION_STRING = 'softmax'
SIGMOID_FUNCTION_STRING = 'sigmoid'

VALID_CONV_ACTIV_FUNC_STRINGS = [
    ELU_FUNCTION_STRING, RELU_FUNCTION_STRING, SELU_FUNCTION_STRING
]
VALID_ACTIVATION_FUNCTION_STRINGS = VALID_CONV_ACTIV_FUNC_STRINGS + [
    SOFTMAX_FUNCTION_STRING, SIGMOID_FUNCTION
_STRING
]


def _check_convolution_options(
        num_kernel_rows, num_rows_per_stride, padding_type_string,
        num_filters, num_kernel_dimensions, num_kernel_columns=None,
     
   num_columns_per_stride=None, num_kernel_heights=None,
        num_heights_per_stride=None):
    """Checks input args for 1-D, 2-D, or 3-D convolution layer.

    :param num_kernel_rows: Number of rows in kernel.
    :param num_rows_per_stride: Number of rows per stride (number of rows moved
        by the kernel at once).
    :param padding_type_string: Padding type (must be in
        `VALID_PADDING_TYPE_STRINGS`).
    :param num_filters: Number of output filters (channels).
    :param num_kernel_dimensions: Number of dimensions in kernel.
    :param num_kernel_columns: [used only if num_kernel_dimensions > 1]
        Number of columns in kernel.
    :param num_columns_per_stride: [used only if num_kernel_dimensions > 1]
        Number of columns per stride.
    :param num_kernel_heights: [used only if num_kernel_dimensions = 3]
        Number of heights in kernel.
    :param num_heights_per_stride: [used only if num_kernel_dimensions = 3]
        Number of heights per stride.
    :raises: ValueError: if
        `padding_type_string not in VALID_PADDING_TYPE_STRINGS`.
    """

    error_checking.assert_is_integer(num_kernel_rows)
    # error_checking.assert_is_geq(num_kernel_rows, 3)
    error_checking.assert_is_geq(num_kernel_rows, 1)
    error_checking.assert_is_integer(num_rows_per_stride)
    error_checking.assert_is_geq(num_rows_per_stride, 1)
    error_checking.assert_is_leq(num_rows_per_stride, num_kernel_rows)

    error_checking.assert_is_string(padding_type_string)
    if padding_type_string not in VALID_PADDING_TYPE_STRINGS:
        error_string = (
            '\n{0:s}\nValid padding types (listed above) do not include '
            '"{1:s}".'
        ).format(str(VALID_PADDING_TYPE_STRINGS), padding_type_string)

        raise ValueError(error_string)

    error_checking.assert_is_integer(num_filters)
    error_checking.assert_is_geq(num_filters, 1)
    error_checking.assert_is_integer(num_kernel_dimensions)
    error_checking.assert_is_geq(num_kernel_dimensions, 1)
    error_checking.assert_is_leq(num_kernel_dimensions, 3)

    if num_kernel_dimensions >= 2:
        error_checking.assert_is_integer(num_kernel_columns)
        # error_checking.assert_is_geq(num_kernel_columns, 3)
        error_checking.assert_is_geq(num_kernel_columns, 1)
        error_checking.assert_is_integer(num_columns_per_stride)
        error_checking.assert_is_geq(num_columns_per_stride, 1)
        error_checking.assert_is_leq(num_columns_per_stride, num_kernel_columns)

    if num_kernel_dimensions == 3:
        error_checking.assert_is_integer(num_kernel_heights)
        # error_checking.assert_is_geq(num_kernel_heights, 3)
        error_checking.assert_is_geq(num_kernel_heights, 1)
        error_checking.assert_is_integer(num_heights_per_stride)
        error_checking.assert_is_geq(num_heights_per_stride, 1)
        error_checking.assert_is_leq(num_heights_per_stride, num_kernel_heights)


def _check_pooling_options(
        num_rows_in_window, num_rows_per_stride, pooling_type_string,
        num_dimensions, num_columns_in_window=None, num_columns_per_stride=None,
        num_heights_in_window=None, num_heights_per_stride=None):
    """Checks input args for 1-D, 2-D, or 3-D pooling layer.

    :param num_rows_in_window: Number of rows in pooling window.
    :param num_rows_per_stride: Number of rows per stride (number of rows moved
        by the window at once).
    :param pooling_type_string: Pooling type (must be in
        `VALID_POOLING_TYPE_STRINGS`).
    :param num_dimensions: Number of dimensions in pooling window.
    :param num_columns_in_window: [used only if num_dimensions > 1]
        Number of columns in window.
    :param num_columns_per_stride: [used only if num_dimensions > 1]
        Number of columns per stride.
    :param num_heights_in_window: [used only if num_dimensions = 3]
        Number of heights in window.
    :param num_heights_per_stride: [used only if num_dimensions = 3]
        Number of heights per stride.
    :raises: ValueError: if
        `pooling_type_string not in VALID_POOLING_TYPE_STRINGS`.
    """

    error_checking.assert_is_integer(num_rows_in_window)
    error_checking.assert_is_geq(num_rows_in_window, 1)
    error_checking.assert_is_integer(num_rows_per_stride)
    error_checking.assert_is_geq(num_rows_per_stride, 1)
    error_checking.assert_is_leq(num_rows_per_stride, num_rows_in_window)

    error_checking.assert_is_string(pooling_type_string)
    if pooling_type_string not in VALID_POOLING_TYPE_STRINGS:
        error_string = (
            '\n{0:s}\nValid pooling types (listed above) do not include '
            '"{1:s}".'
        ).format(str(VALID_POOLING_TYPE_STRINGS), pooling_type_string)

        raise ValueError(error_string)

    error_checking.assert_is_integer(num_dimensions)
    error_checking.assert_is_geq(num_dimensions, 1)
    error_checking.assert_is_leq(num_dimensions, 3)

    if num_dimensions >= 2:
        error_checking.assert_is_integer(num_columns_in_window)
        error_checking.assert_is_geq(num_columns_in_window, 1)
        error_checking.assert_is_integer(num_columns_per_stride)
        error_checking.assert_is_geq(num_columns_per_stride, 1)
        error_checking.assert_is_leq(num_columns_per_stride,
                                     num_columns_in_window)

    if num_dimensions == 3:
        error_checking.assert_is_integer(num_heights_in_window)
        error_checking.assert_is_geq(num_heights_in_window, 1)
        error_checking.assert_is_integer(num_heights_per_stride)
        error_checking.assert_is_geq(num_heights_per_stride, 1)
        error_checking.assert_is_leq(num_heights_per_stride,
                                     num_heights_in_window)


def check_activation_function(
        activation_function_string, alpha_for_elu=DEFAULT_ALPHA_FOR_ELU,
        alpha_for_relu=DEFAULT_ALPHA_FOR_RELU, for_conv_layer=False):
    """Error-checks activation function.

    :param activation_function_string: Name of activation function.
    :param alpha_for_elu: [used only if activation function is "elu"]
        Slope for negative inputs to activation function.
    :param alpha_for_relu: [used only if activation function is "relu"]
        Slope for negative inputs to activation function.
    :param for_conv_layer: Boolean flag.  If True, activation function is for a
        convolution layer.
    :raises: ValueError: if
        `activation_function_string not in VALID_ACTIVATION_FUNCTION_STRINGS`.
    """

    error_checking.assert_is_boolean(for_conv_layer)
    error_checking.assert_is_string(activation_function_string)

    if for_conv_layer:
        these_valid_strings = VALID_CONV_ACTIV_FUNC_STRINGS
    else:
        these_valid_strings = VALID_ACTIVATION_FUNCTION_STRINGS

    if activation_function_string not in these_valid_strings:
        error_string = (
            '\n{0:s}\nValid activation functions (listed above) do not '
            'include "{1:s}".'
        ).format(str(these_valid_strings), activation_function_string)

      