from typing import Optional, Tuple, Union

import tensorflow as tf

from tfcv.modeling.layers import normalization


__all__ = ["conv2d", "separable_conv2d"]


def norm_act(
    x,
    norm: Union[str, None],
    weight_decay: float,
    *,
    norm_weight_zero_init: bool = False,
    activation: Union[str, None],
):
    if norm is not None:
        x = normalization(
            x, norm, weight_decay, norm_weight_zero_init=norm_weight_zero_init
        )
    if activation is not None:
        x = tf.keras.layers.Activation(activation)(x)
    return x


def build_conv(
    x,
    conv: tf.keras.layers.Layer,
    weight_decay: float,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
):
    if pre_activation:
        assert activation is not None, "Activation must be enabled for pre-activation!"
        assert norm is not None, "Normalization must be enabled for pre-activation!"
        x = norm_act(
            x,
            norm,
            weight_decay,
            norm_weight_zero_init=norm_weight_zero_init,
            activation=activation,
        )
        x = conv(x)
    else:
        x = conv(x)
        x = norm_act(
            x,
            norm,
            weight_decay,
            norm_weight_zero_init=norm_weight_zero_init,
            activation=activation,
        )
    return x


def conv2d(
    x,
    filters: int,
    kernel_size: Union[int, Tuple[int, int]],
    weight_decay: float,
    *,
    strides: Union[int, Tuple[int, int]] = (1, 1),
    padding: str = "same",
    dilation: Union[int, Tuple[int, int]] = (1, 1),
    groups: int = 1,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
    name: Optional[str] = None,
):
    conv = tf.keras.layers.Conv2D(
        filters,
        kernel_size,
        strides,
        padding,
        dilation_rate=dilation,
        groups=groups,
        use_bias=False if norm else True,
        kernel_initializer="he_normal",
        kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
        name=name,
    )
    x = build_conv(
        x, conv, weight_decay, norm, norm_weight_zero_init, activation, pre_activation
    )
    return x


def separable_conv2d(
    x,
    filters: int,
    kernel_size: Union[int, Tuple[int, int]],
    weight_decay: float,
    *,
    strides: Union[int, Tuple[int, int]] = (1, 1),
    padding: str = "same",
    dilation: Union[int, Tuple[int, int]] = (1, 1),
    depth_multiplier: int = 1,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
    name: Optional[str] = None,
):
    conv = tf.keras.layers.SeparableConv2D(
        filters,
        kernel_size,
        strides,
        padding,
        dilation_rate=dilation,
        depth_multiplier=depth_multiplier,
        use_bias=False if norm else True,
        depthwise_initializer="he_normal",
        pointwise_initializer="he_normal",
        depthwise_regularizer=tf.keras.regularizers.l2(weight_decay),
        pointwise_regularizer=tf.keras.regularizers.l2(weight_decay),
        name=name,
    )
    x = build_conv(
        x, conv, weight_decay, norm, norm_weight_zero_init, activation, pre_activation
    )
    return x

from typing import Optional, Tuple, Union

import tensorflow as tf

from tfcv.modeling.layers import normalization


__all__ = ["conv2d", "separable_conv2d"]


def norm_act(
    x,
    norm: Union[str, None],
    weight_decay: float,
    *,
    norm_weight_zero_init: bool = False,
    activation: Union[str, None],
):
    if norm is not None:
        x = normalization(
            x, norm, weight_decay, norm_weight_zero_init=norm_weight_zero_init
        )
    if activation is not None:
        x = tf.keras.layers.Activation(activation)(x)
    return x


def build_conv(
    x,
    conv: tf.keras.layers.Layer,
    weight_decay: float,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
):
    if pre_activation:
        assert activation is not None, "Activation must be enabled for pre-activation!"
        assert norm is not None, "Normalization must be enabled for pre-activation!"
        x = norm_act(
            x,
            norm,
            weight_decay,
            norm_weight_zero_init=norm_weight_zero_init,
            activation=activation,
        )
        x = conv(x)
    else:
        x = conv(x)
        x = norm_act(
            x,
            norm,
            weight_decay,
            norm_weight_zero_init=norm_weight_zero_init,
            activation=activation,
        )
    return x


def conv2d(
    x,
    filters: int,
    kernel_size: Union[int, Tuple[int, int]],
    weight_decay: float,
    *,
    strides: Union[int, Tuple[int, int]] = (1, 1),
    padding: str = "same",
    dilation: Union[int, Tuple[int, int]] = (1, 1),
    groups: int = 1,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
    name: Optional[str] = None,
):
    conv = tf.keras.layers.Conv2D(
        filters,
        kernel_size,
        strides,
        padding,
        dilation_rate=dilation,
        groups=groups,
        use_bias=False if norm else True,
        kernel_initializer="he_normal",
        kernel_regularizer=tf.keras.regularizers.l2(weight_decay),
        name=name,
    )
    x = build_conv(
        x, conv, weight_decay, norm, norm_weight_zero_init, activation, pre_activation
    )
    return x


def separable_conv2d(
    x,
    filters: int,
    kernel_size: Union[int, Tuple[int, int]],
    weight_decay: float,
    *,
    strides: Union[int, Tuple[int, int]] = (1, 1),
    padding: str = "same",
    dilation: Union[int, Tuple[int, int]] = (1, 1),
    depth_multiplier: int = 1,
    norm: Union[str, None] = "bn",
    norm_weight_zero_init: bool = False,
    activation: Union[str, None] = "relu",
    pre_activation: bool = False,
    name: Optional[str] = None,
):
    conv = tf.keras.layers.SeparableConv2D(
        filters,
        kernel_size,
        strides,
        padding,
        dilation_rate=dilation,
        depth_multiplier=depth_multiplier,
        use_bias=False if norm else True,
        depthwise_initializer="he_normal",
        pointwise_initializer="he_normal",
        depthwise_regularizer=tf.keras.regularizers.l2(weight_decay),
        pointwise_regularizer=tf.keras.regularizers.l2(weight_decay),
        name=name,
    )
    x = build_conv(
        x, conv, weight_decay, norm, norm_weight_zero_init, activation, pre_activation
    )
    return x
