import tensorflow as tf

from .pooling import AdaptiveAveragePooling2D

H_AXIS = -3
W_AXIS = -2


@tf.keras.utils.register_keras_serializable(package="gcvit")
class Mlp(tf.keras.layers.Layer):
    def __init__(
        self,
        hidden_features=None,
        out_features=None,
        act_layer="gelu",
        dropout=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.hidden_features = hidden_features
        self.out_features = out_features
        self.act_layer = act_layer
        self.dropout = dropout

    def build(self, input_shape):
        self.in_features = input_shape[-1]
        self.hidden_features = self.hidden_features or self.in_features
        self.out_features = self.out_features or self.in_features
        self.fc1 = tf.keras.layers.Dense(self.hidden_features, name="fc1")
        self.act = tf.keras.layers.Activation(self.act_layer, name="act")
        self.fc2 = tf.keras.layers.Dense(self.out_features, name="fc2")
        self.drop1 = tf.keras.layers.Dropout(self.dropout, name="drop1")
        self.drop2 = tf.keras.layers.Dropout(self.dropout, name="drop2")
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        x = self.fc1(inputs)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.drop2(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "hidden_features": self.hidden_features,
                "out_features": self.out_features,
                "act_layer": self.act_layer,
                "dropout": self.dropout,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class SE(tf.keras.layers.Layer):
    def __init__(self, oup=None, expansion=0.25, **kwargs):
        super().__init__(**kwargs)
        self.expansion = expansion
        self.oup = oup

    def build(self, input_shape):
        inp = input_shape[-1]
        self.oup = self.oup or inp
        self.avg_pool = AdaptiveAveragePooling2D(1, name="avg_pool")
        self.fc = [
            tf.keras.layers.Dense(
                int(inp * self.expansion), use_bias=False, name="fc_0"
            ),
            tf.keras.layers.Activation("gelu", name="fc_1"),
            tf.keras.layers.Dense(self.oup, use_bias=False, name="fc_2"),
            tf.keras.layers.Activation("sigmoid", name="fc_3"),
        ]
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        b, _, _, c = tf.unstack(tf.shape(inputs), num=4)
        x = tf.reshape(self.avg_pool(inputs), (b, c))
        for layer in self.fc:
            x = layer(x)
        x = tf.reshape(x, (b, 1, 1, c))
        return x * inputs

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "expansion": self.expansion,
                "oup": self.oup,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class ReduceSize(tf.keras.layers.Layer):
    def __init__(self, keep_dim=False, **kwargs):
        super().__init__(**kwargs)
        self.keep_dim = keep_dim

    def build(self, input_shape):
        dim = input_shape[-1]
        dim_out = dim if self.keep_dim else 2 * dim
        self.pad1 = tf.keras.layers.ZeroPadding2D(1, name="pad1")
        self.pad2 = tf.keras.layers.ZeroPadding2D(1, name="pad2")
        self.conv = [
            tf.keras.layers.DepthwiseConv2D(
                kernel_size=3,
                strides=1,
                padding="valid",
                use_bias=False,
                name="conv_0",
            ),
            tf.keras.layers.Activation("gelu", name="conv_1"),
            SE(name="conv_2"),
            tf.keras.layers.Conv2D(
                dim,
                kernel_size=1,
                strides=1,
                padding="valid",
                use_bias=False,
                name="conv_3",
            ),
        ]
        self.reduction = tf.keras.layers.Conv2D(
            dim_out,
            kernel_size=3,
            strides=2,
            padding="valid",
            use_bias=False,
            name="reduction",
        )
        self.norm1 = tf.keras.layers.LayerNormalization(
            axis=-1, epsilon=1e-05, name="norm1"
        )  # eps like PyTorch
        self.norm2 = tf.keras.layers.LayerNormalization(
            axis=-1, epsilon=1e-05, name="norm2"
        )
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        x = self.norm1(inputs)
        xr = self.pad1(
            x
        )  # if pad had weights it would've thrown error with .save_weights()
        for layer in self.conv:
            xr = layer(xr)
        x = x + xr
        x = self.pad2(x)
        x = self.reduction(x)
        x = self.norm2(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "keep_dim": self.keep_dim,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class FeatExtract(tf.keras.layers.Layer):
    def __init__(self, keep_dim=False, **kwargs):
        super().__init__(**kwargs)
        self.keep_dim = keep_dim

    def build(self, input_shape):
        dim = input_shape[-1]
        self.pad1 = tf.keras.layers.ZeroPadding2D(1, name="pad1")
        self.pad2 = tf.keras.layers.ZeroPadding2D(1, name="pad2")
        self.conv = [
            tf.keras.layers.DepthwiseConv2D(
                kernel_size=3,
                strides=1,
                padding="valid",
                use_bias=False,
                name="conv_0",
            ),
            tf.keras.layers.Activation("gelu", name="conv_1"),
            SE(name="conv_2"),
            tf.keras.layers.Conv2D(
                dim,
                kernel_size=1,
                strides=1,
                padding="valid",
                use_bias=False,
                name="conv_3",
            ),
        ]
        if not self.keep_dim:
            self.pool = tf.keras.layers.MaxPool2D(
                pool_size=3, strides=2, padding="valid", name="pool"
            )
        # else:
        #     # hack for PyTorch nn.Identity layer ;)
        #     self.pool = tf.keras.layers.Activation('linear', name='identity')
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        x = inputs
        xr = self.pad1(x)
        for layer in self.conv:
            xr = layer(xr)
        x = (
            x + xr
        )  # if pad had weights it would've thrown error with .save_weights()
        if not self.keep_dim:
            x = self.pad2(x)
            x = self.pool(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "keep_dim": self.keep_dim,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class GlobalQueryGen(tf.keras.layers.Layer):
    """
    Global query generator based on: "Hatamizadeh et al.,
    Global Context Vision Transformers <https://arxiv.org/abs/2206.09959>"
    """

    def __init__(self, keep_dims=False, **kwargs):
        super().__init__(**kwargs)
        self.keep_dims = keep_dims

    def build(self, input_shape):
        self.to_q_global = [
            FeatExtract(keep_dim, name=f"to_q_global_{i}")
            for i, keep_dim in enumerate(self.keep_dims)
        ]
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        x = inputs
        for layer in self.to_q_global:
            x = layer(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "keep_dims": self.keep_dims,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class Resizing(tf.keras.layers.Layer):
    def __init__(self, height, width, interpolation="bilinear", **kwargs):
        self.height = height
        self.width = width
        self.interpolation = interpolation
        super().__init__(**kwargs)

    def call(self, inputs):
        # tf.image.resize will always output float32 and operate more
        # efficiently on float32 unless interpolation is nearest, in
        # which case ouput type matches input type.
        if self.interpolation == "nearest":
            input_dtype = self.compute_dtype
        else:
            input_dtype = tf.float32
        inputs = tf.cast(inputs, dtype=input_dtype)
        size = [self.height, self.width]
        outputs = tf.image.resize(inputs, size=size, method=self.interpolation)
        return tf.cast(outputs, self.compute_dtype)

    def compute_output_shape(self, input_shape):
        input_shape = tf.TensorShape(input_shape).as_list()
        input_shape[H_AXIS] = self.height
        input_shape[W_AXIS] = self.width
        return tf.TensorShape(input_shape)

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "height": self.height,
                "width": self.width,
                "interpolation": self.interpolation,
            }
        )
        return config


@tf.keras.utils.register_keras_serializable(package="gcvit")
class FitWindow(tf.keras.layers.Layer):
    "Pad feature to fit window"

    def __init__(self, window_size, **kwargs):
        super().__init__(**kwargs)
        self.window_size = window_size

    def call(self, inputs):
        B, H, W, C = tf.unstack(tf.shape(inputs), num=4)
        # pad to multiple of window_size
        h_pad = (self.window_size - H % self.window_size) % self.window_size
        w_pad = (self.window_size - W % self.window_size) % self.window_size
        x = tf.pad(
            inputs,
            [
                [0, 0],
                [
                    h_pad // 2,
                    (h_pad // 2 + h_pad % 2),
                ],  # padding in both directions unlike tfgcvit
                [w_pad // 2, (w_pad // 2 + w_pad % 2)],
                [0, 0],
            ],
        )
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "window_size": self.window_size,
            }
        )
        return config
