import math
import torch
import torch.nn as nn
from .utils import IMAGE_INDICATOR_IDS


def build_vision_projector(config, delay_load=False, **kwargs):
    projector_type = getattr(config, 'mm_projector_type', 'linear')

    if projector_type == 'conv_adapter':
        return ConvAdapter(config.mm_hidden_size, config.hidden_size, getattr(config, "mlp_hidden_dim", None))
    elif projector_type == 'mlp_pixel_shuffle':
        return MlpPixelShuffle(config.mm_hidden_size, config.hidden_size,
                               config.pixelshuffle_downsample_ratio, getattr(config, "mlp_hidden_dim", None))
    elif projector_type == 'ovis_conv_adapter':
        return OvisConvAdapter(config.mm_hidden_size, config.hidden_size, getattr(config, "mlp_hidden_dim", 32000),
                               getattr(config, "tokenize_function", "softmax"))
    elif projector_type == 'ovis2_adapter':
        return Ovis2Adapter(config.mm_hidden_size, config.hidden_size, getattr(config, "mlp_hidden_dim", 66536), 
                            getattr(config, "hidden_stride", 2), getattr(config, "pooling_stride", 1), getattr(config, "tokenize_function", "softmax"))
    elif projector_type == 'ovis_conv_adapter_navit':
        return OvisConvAdapterNavit(1280, config.hidden_size, getattr(config, "mlp_hidden_dim", 32000), getattr(config, "tokenize_function", "softmax"))
    raise ValueError(f'Unknown projector type: {projector_type}')


class ConvAdapter(nn.Module):
    def __init__(self, dim_in, dim_out, mlp_hidden_dim=None):
        super().__init__()
        self.mm_projector_type = 'conv_adapter'
        if mlp_hidden_dim is None:
            self.mlp = nn.Sequential(
                nn.Linear(dim_in, dim_out),
                nn.GELU(),
                nn.Linear(dim_out, dim_out)
            )
        else:
            self.mlp = nn.Sequential(
                nn.Linear(dim_in, mlp_hidden_dim),
                nn.GELU(),
                nn.Linear(mlp_hidden_dim, dim_out)
            )
        self.conv = nn.Conv2d(dim_out, dim_out, kernel_size=(3, 3), stride=(2, 2), padding=1)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): image features
                shape (F, v, D)
        Returns:
            shape (F, n, D) where n is token_num that has been reduced
        """
        x = self.mlp(x)

        f, v, d = x.shape
        s = int(math.sqrt(v - 1))
        x = x[:, 1:, :]  # remove cls_token
        x = x.reshape(f, s, s, d).permute([0, 3, 1, 2])
        x = self.conv(x)
        x = x.permute([0, 2, 3, 1]).reshape(f, -1, d)
        return x


class MlpPixelShuffle(nn.Module):
    def __init__(self, dim_in, dim_out, pixelshuffle_downsample_ratio, mlp_hidden_dim=None):
        super().__init__()
        self.mm_projector_type = 'mlp_pixel_shuffle'
        if mlp_hidden_dim is None:
            self.mlp = nn.Sequential(
                nn.Linear(int(dim_in * (pixelshuffle_downsample_ratio ** 2)), dim_out),
                nn.GELU(),
                nn.Linear(dim_out, dim_out)
            )
        else:
            self.mlp = nn.Sequential(
                nn.Linear(int(dim_in * (pixelshuffle_downsample_ratio ** 2)), mlp_hidden_dim),
                nn.GELU(),
                nn.Linear(mlp_hidden_dim, dim_out)
            )
        self.scale_factor = pixelshuffle_downsample_ratio

    def pixel_shuffle(self, x, scale_factor=2):
        # change scale_factor from float to int

        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H / scale, C * scale
        x = x.view(n, w, int(h / scale_factor), int(c * scale_factor))
        # N, W, H / scale, C * scale --> N, H / scale, W, C * scale
        x = x.permute(0, 2, 1, 3).contiguous()
        # N, H / scale, W, C * scale --> N, H / scale, W / scale, C * (scale ** 2)
        x = x.view(n, int(h / scale_factor), int(w / scale_factor),
                   int(c * (scale_factor * scale_factor)))

        x = x.permute(0, 2, 1, 3).contiguous()

        return x

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): image features
                shape (F, v, D)
        Returns:
            shape (F, n, D) where n is token_num that has been reduced
        """
        x = x[:, 1:, :]  # remove cls_token
        h = w = int(x.shape[1] ** 0.5)
        x = x.view(x.shape[0], h, w, -1)
        x = self.pixel_shuffle(x, self.scale_factor)
        x = self.mlp(x)
        x = x.view(x.shape[0],-1,x.shape[-1])
        return x


class OvisConvAdapter(nn.Module):
    def __init__(self, dim_in, dim_out, vocab_size, tokenize_function="softmax"):
        super().__init__()
        self.mm_projector_type = 'ovis_conv_adapter'
        self.conv = nn.Conv2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2), padding=1)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(dim_in, vocab_size, bias=False),
            torch.nn.LayerNorm(vocab_size)
        )
        self.embedding = torch.nn.Embedding(vocab_size, dim_out)
        self.tokenize_function = tokenize_function

    def tokenize(self, logits):
        def st_argmax(y_soft, dim):  # straight-through softmax
            index = y_soft.max(dim, keepdim=True)[1]
            y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
            ret = y_hard - y_soft.detach() + y_soft
            return ret

        if self.tokenize_function == 'softmax':
            tokens = torch.nn.functional.softmax(logits, dim=-1)
        elif self.tokenize_function == 'gumbel_argmax':
            tokens = torch.nn.functional.gumbel_softmax(logits, tau=self.config.tau, hard=True)
        elif self.tokenize_function == 'st_argmax':
            tokens = st_argmax(logits, dim=-1)
        else:
            raise ValueError(
                'Invalid `max_type`, expected softmax or gumbel_argmax or st_argmax,'
                f' but got {self.config.tokenize_function}'
            )
        return tokens

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): image features
                shape (F, v, D)
        Returns:
            shape (F, n, D) where n is token_num that has been reduced
        """
        # conv
        f, v, d = x.shape
        s = int(math.sqrt(v - 1))
        x = x[:, 1:, :]  # remove cls_token
        x = x.reshape(f, s, s, d).permute([0, 3, 1, 2])
        x = self.conv(x)
        x = x.permute([0, 2, 3, 1]).reshape(f, -1, d)

        # tokenize
        logits = self.mlp(x)
        visual_tokens = self.tokenize(logits)

        # get embeddings
        out = torch.matmul(visual_tokens, self.embedding.weight)

        return out


class Ovis2Adapter(nn.Module):
    def __init__(self, dim_in, dim_out, vocab_size, hidden_stride=2, pooling_stride=1, tokenize_function="softmax"):
        super().__init__()
        head_dim = vocab_size - len(IMAGE_INDICATOR_IDS)
        self.mm_projector_type = 'ovis2_adapter'
        self.hidden_stride = hidden_stride
        self.tokenize_function = tokenize_function
        self.head = torch.nn.Sequential(
            torch.nn.Linear(
                dim_in * self.hidden_stride * self.hidden_stride, head_dim,
                bias=False
            ),
            torch.nn.LayerNorm(head_dim)
        )
        self.embedding = torch.nn.Embedding(vocab_size, dim_out)
        self.pool_s = pooling_stride
        print("pooling_stride: ", pooling_stride)
    
    def encode(self, features):
        # merge number of `hidden_stride * hidden_stride` hidden states together to reduce token sequence length
        # e.g., for hidden_stride=2, this leads to a token length reduction: 1024 -> 256 for aimv2
        features = features[:, 1:, :]
        if self.hidden_stride > 1:
            n, l, d = features.shape  # this `d` maybe different from the above `d
            sqrt_l = int(l ** 0.5)
            assert sqrt_l ** 2 == l, "The token sequence length should be a perfect square."
            features = features.reshape(n, sqrt_l, sqrt_l, d)
            pl = (self.hidden_stride - (sqrt_l % self.hidden_stride)) % self.hidden_stride
            features = torch.nn.functional.pad(features, (0, 0, 0, pl, 0, pl), "constant", 0)
            sqrt_l += pl
            features = features.reshape(n, sqrt_l // self.hidden_stride, self.hidden_stride,
                                        sqrt_l // self.hidden_stride, self.hidden_stride, d)
            features = features.permute(0, 1, 3, 2, 4, 5)  # [n, sqrt_l/hs, sqrt_l/hs, hs, hs, d]
            features = features.flatten(3)  # [n, sqrt_l/hs, sqrt_l/hs, hs*hs*d]
            features = features.reshape(
                n, -1, self.hidden_stride * self.hidden_stride * d)

        return features

    def tokenize(self, logits):
        def st_argmax(y_soft, dim):  # straight-through softmax
            index = y_soft.max(dim, keepdim=True)[1]
            y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
            ret = y_hard - y_soft.detach() + y_soft
            return ret

        if self.tokenize_function == 'softmax':
            tokens = torch.nn.functional.softmax(logits, dim=-1)
        elif self.tokenize_function == 'gumbel_argmax':
            tokens = torch.nn.functional.gumbel_softmax(logits, tau=self.config.tau, hard=True)  # here need to be check？？？
        elif self.tokenize_function == 'st_argmax':
            tokens = st_argmax(logits, dim=-1)
        else:
            raise ValueError(
                f'Invalid `max_type`, expected softmax or gumbel_argmax or st_argmax, but got {self.config.tokenize_function}')
        return tokens

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): image features
                shape (F, v, D)
        Returns:
            shape (F, n, D) where n is token_num that has been reduced
        """
        # pixelshuffle
        # f, v, d = x.shape
        # s = int(math.sqrt(v))
        x = self.encode(x)
        # tokenize
        logits = self.head(x)
        visual_tokens = self.tokenize(logits)
        batch_size, token_len, _ = visual_tokens.shape
        padding_tensor = torch.zeros(size=(batch_size, token_len, len(IMAGE_INDICATOR_IDS)),
                                dtype=visual_tokens.dtype,
                                device=visual_tokens.device,
                                layout=visual_tokens.layout,
                                requires_grad=False)
        visual_tokens = torch.cat([visual_tokens, padding_tensor], dim=2)
        # get embeddings here need to change argmax
        out = torch.matmul(visual_tokens, self.embedding.weight)
        
        if self.pool_s > 1:
            f, v, d = out.shape
            s = int(math.sqrt(v))
            out = out.reshape(f, s, s, d)
            out = out.reshape(f, s // self.pool_s, self.pool_s, s // self.pool_s, self.pool_s, d)
            out = out.permute([0, 1, 3, 5, 2, 4]).reshape(f, s // self.pool_s * s // self.pool_s, d, -1).mean(-1)
        return out

class OvisConvAdapterNavit(nn.Module):
    def __init__(self, dim_in, dim_out, vocab_size, tokenize_function="softmax"):
        super().__init__()
        self.mm_projector_type = 'ovis_conv_adapter_navit'
        self.conv = nn.Conv2d(dim_in, dim_in, kernel_size=(2, 2), stride=(2, 2))
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(dim_in, vocab_size, bias=False),
            torch.nn.LayerNorm(vocab_size)
        )
        self.embedding = torch.nn.Embedding(vocab_size, dim_out)
        self.tokenize_function = tokenize_function
    def tokenize(self, logits):
        def st_argmax(y_soft, dim):  # straight-through softmax
            index = y_soft.max(dim, keepdim=True)[1]
            y_hard = torch.zeros_like(y_soft, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
            ret = y_hard - y_soft.detach() + y_soft
            return ret
        if self.tokenize_function == 'softmax':
            tokens = torch.nn.functional.softmax(logits, dim=-1)
        elif self.tokenize_function == 'gumbel_argmax':
            tokens = torch.nn.functional.gumbel_softmax(logits, tau=self.config.tau, hard=True)
        elif self.tokenize_function == 'st_argmax':
            tokens = st_argmax(logits, dim=-1)
        else:
            raise ValueError(
                f'Invalid `max_type`, expected softmax or gumbel_argmax or st_argmax, but got {self.config.tokenize_function}')
        return tokens
    def forward(self, x):
        """
        Args:
            x (torch.Tensor): image features of navit
                shape (v, D)
        Returns:
            shape (n, D) where n is token_num that has been reduced
        """
        # conv
        _, d = x.shape
        x = x.reshape(-1, 2, 2, d).permute([0, 3, 1, 2])
        x = self.conv(x)
        x = x.permute([0, 2, 3, 1]).reshape(-1, d)
        # tokenize
        logits = self.mlp(x)
        visual_tokens = self.tokenize(logits)
        # get embeddings
        out = torch.matmul(visual_tokens, self.embedding.weight)
        
        return out