from __future__ import annotations
from typing import override
import torch
from torch import nn
from ..model.config import Config
from . import Module
from ..util.tensor import to2


class Qwen3VLPosEmbedding(Module):

    def __init__(
        self,
        config: Config,
        key: str,
        num_position_embeddings: int,
        hidden_size: int,
        spatial_merge_size: int,
        out_dtype: torch.dtype | None = None,
        qmap: str | None = None,
    ):
        super().__init__(config, key, None)
        self.module_name = "Qwen3VLPosEmbedding"
        self.qmap = qmap
        self.key = key

        self.num_position_embeddings = num_position_embeddings
        self.num_grid_per_side = num_position_embeddings ** 0.5
        self.spatial_merge_size = spatial_merge_size
        self.hidden_size = hidden_size
        self.out_dtype = out_dtype

        self.embedding = None


    def fast_pos_embed_interpolate(self, grid_thw):
        grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]

        idx_list = [[] for _ in range(4)]
        weight_list = [[] for _ in range(4)]

        for t, h, w in zip(grid_ts, grid_hs, grid_ws):
            h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
            w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
            h_idxs_floor = h_idxs.int()
            w_idxs_floor = w_idxs.int()
            h_idxs_ceil = (h_idxs.int() + 1).clip(max = self.num_grid_per_side - 1)
            w_idxs_ceil = (w_idxs.int() + 1).clip(max = self.num_grid_per_side - 1)
            dh = h_idxs - h_idxs_floor
            dw = w_idxs - w_idxs_floor

            base_h = h_idxs_floor * self.num_grid_per_side
            base_h_ceil = h_idxs_ceil * self.num_grid_per_side

            indices = [
                (base_h[None].T + w_idxs_floor[None]).flatten(),
                (base_h[None].T + w_idxs_ceil[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
            ]

            weights = [
                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
                ((1 - dh)[None].T * dw[None]).flatten(),
                (dh[None].T * (1 - dw)[None]).flatten(),
                (dh[None].T * dw[None]).flatten(),
            ]

            for i in range(4):
                idx_list[i].extend(indices[i].tolist())
                weight_list[i].extend(weights[i].tolist())

        idx_tensor = torch.tensor(idx_list, dtype = torch.long, device = self.embedding.weight.device)
        weight_tensor = torch.tensor(
            weight_list, dtype = self.embedding.weight.dtype, device = self.embedding.weight.device
        )
        pos_embeds = self.embedding(idx_tensor) * weight_tensor[:, :, None]
        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]

        patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])

        patch_pos_embeds_permute = []
        merge_size = self.spatial_merge_size
        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
            pos_embed = pos_embed.repeat(t, 1)
            pos_embed = (
                pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
                .permute(0, 1, 3, 2, 4, 5)
                .flatten(0, 4)
            )
            patch_pos_embeds_permute.append(pos_embed)
        patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
        return patch_pos_embeds


    @override
    def weights_numel(self):
        return self.num_position_embeddings + self.hidden_size


    def optimizer_targets(self):
        raise NotImplementedError()


    @override
    def load(self, device: torch.device, **kwargs):
        self.device = device
        weight = self.config.stc.get_tensor(self.key + ".weight", self.device, allow_bf16 = True)
        self.embedding = nn.Embedding(
            self.num_position_embeddings,
            self.hidden_size,
            device = "meta"
        )
        self.embedding.weight = nn.Parameter(weight)


    @override
    def unload(self):
        self.device = None
        self.embedding = None


    def forward(
        self,
        x: torch.Tensor,
        params: dict,
        out_dtype: torch.dtype | None = None
    ):
        pos_emb = self.fast_pos_embed_interpolate(params["grid_thw"])
        x += pos_emb.to(x.device)
        return to2(x, out_dtype, self.out_dtype)
