import tensorflow as tf

from .block import GCViTBlock
from .feature import FitWindow
from .feature import GlobalQueryGen
from .feature import ReduceSize
from .feature import Resizing


@tf.keras.utils.register_keras_serializable(package="gcvit")
class GCViTLevel(tf.keras.layers.Layer):
    def __init__(
        self,
        depth,
        num_heads,
        window_size,
        keep_dims,
        downsample=True,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        path_drop=0.0,
        layer_scale=None,
        resize_query=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.depth = depth
        self.num_heads = num_heads
        self.window_size = window_size
        self.keep_dims = keep_dims
        self.downsample = downsample
        self.mlp_ratio = mlp_ratio
        self.qkv_bias = qkv_bias
        self.qk_scale = qk_scale
        self.drop = drop
        self.attn_drop = attn_drop
        self.path_drop = path_drop
        self.layer_scale = layer_scale
        self.resize_query = resize_query

    def build(self, input_shape):
        path_drop = (
            [self.path_drop] * self.depth
            if not isinstance(self.path_drop, list)
            else self.path_drop
        )
        self.blocks = [
            GCViTBlock(
                window_size=self.window_size,
                num_heads=self.num_heads,
                global_query=bool(i % 2),
                mlp_ratio=self.mlp_ratio,
                qkv_bias=self.qkv_bias,
                qk_scale=self.qk_scale,
                drop=self.drop,
                attn_drop=self.attn_drop,
                path_drop=path_drop[i],
                layer_scale=self.layer_scale,
                name=f"blocks_{i}",
            )
            for i in range(self.depth)
        ]
        self.down = ReduceSize(keep_dim=False, name="downsample")
        self.q_global_gen = GlobalQueryGen(self.keep_dims, name="q_global_gen")
        self.resize = Resizing(
            self.window_size, self.window_size, interpolation="bicubic"
        )
        self.fit_window = FitWindow(self.window_size)
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        H, W = tf.unstack(tf.shape(inputs)[1:3], num=2)
        # pad to fit window_size
        x = self.fit_window(inputs)
        # generate global query
        q_global = self.q_global_gen(x)  # (B, H, W, C)
        # - official impl issue:
        # https://github.com/NVlabs/GCVit/issues/13
        # TODO: resize query to fit key-value, but result in poor score
        # with official weights?
        if self.resize_query:
            # to avoid mismatch between feat_map and q_global:
            # https://github.com/NVlabs/GCVit/issues/9
            q_global = self.resize(q_global)
        # feature_map -> windows -> window_attention -> feature_map
        for i, blk in enumerate(self.blocks):
            if i % 2:
                x = blk([x, q_global])
            else:
                x = blk([x])
        x = x[:, :H, :W, :]  # https://github.com/NVlabs/GCVit/issues/9
        # set shape for [B, ?, ?, C]
        x.set_shape(
            inputs.shape
        )  # `tf.reshape` creates new tensor with new_shape
        # downsample
        if self.downsample:
            x = self.down(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "depth": self.depth,
                "num_heads": self.num_heads,
                "window_size": self.window_size,
                "keep_dims": self.keep_dims,
                "downsample": self.downsample,
                "mlp_ratio": self.mlp_ratio,
                "qkv_bias": self.qkv_bias,
                "qk_scale": self.qk_scale,
                "drop": self.drop,
                "attn_drop": self.attn_drop,
                "path_drop": self.path_drop,
                "layer_scale": self.layer_scale,
            }
        )
        return config
