"""
modeling_qwen3_5vit_qwen3.py — Qwen3.5 Vision as SigLIP-compat vision_tower in LlavaOnevision.

设计：与 `modeling_qwen3vlvit_qwen3.py` 严格同构，仅换 vision backbone 源：
  - Qwen3VLVisionModel → Qwen3_5VisionModel（继承关系：Qwen3_5VisionModel(Qwen3VLVisionModel) 去 DeepStack）
  - Qwen3VLVisionConfig → Qwen3_5VisionConfig（父类用 AttributeError 哨兵屏蔽 deepstack_visual_indexes）

其余（Adapter 契约翻译、MLP projector + pre_norm、LlavaOnevision 继承 wire class）与
Qwen3-VL ViT pipeline 完全一致。两条 pipeline 并存意义：DeepStack ablation 天然实验组。

类层级：
  Qwen3_5ViTBackbone(Qwen3_5VisionModel)      — 去 merger，保持 NaViT 契约
  Qwen3_5ViTAsSiglipAdapter(nn.Module)        — 持有 Backbone，做 SigLIP ↔ NaViT 契约翻译

三方对比公平性：定 384×384 AnyRes tile + 同款 projector 骨架 + 同款 Qwen3-1.7B LLM。
"""

import math
import os
import sys
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    LlavaOnevisionConfig,
    LlavaOnevisionForConditionalGeneration,
    LlavaOnevisionModel,
    LlavaOnevisionPreTrainedModel,
    Qwen3Config,
)
from transformers.activations import ACT2FN
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5VisionConfig
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel

# Shared layout-permutation utility lives in declip_qwenvit (single source of
# truth — same code path runs in declip-training-side qk_cosine reorder).
# Add VisionEncoder repo root to sys.path so this modeling file is importable
# even when the package isn't pip-installed (ms-swift integration loads it
# via dynamic plugin path).
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _REPO_ROOT not in sys.path:
    sys.path.append(_REPO_ROOT)
from declip_qwenvit.model.qwen3vit_qk import block_merge_to_row_major_permutation  # noqa: E402


class LlavaQwen3_5ViTConfig(LlavaOnevisionConfig):
    """LlavaOnevisionConfig 子类，vision_config 类型换成 Qwen3_5VisionConfig。

    与 Qwen3-VL ViT 版的差异：
      - sub_configs.vision_config 用 Qwen3_5VisionConfig
      - 不再设置 deepstack_visual_indexes（Qwen3_5VisionConfig 用 AttributeError 哨兵屏蔽此字段）

    其余同 LlavaQwen3VLViTConfig（tile_size 默认 384；vision_feature_select_strategy='full'
    必须 override，Qwen3.5 ViT 无 CLS token）。
    """

    model_type = "llava_qwen3_5vit_qwen3"
    sub_configs = {"vision_config": Qwen3_5VisionConfig, "text_config": Qwen3Config}

    def __init__(
        self,
        vision_config=None,
        text_config=None,
        tile_size: int = 384,
        **kwargs,
    ):
        if isinstance(vision_config, dict):
            vision_config = Qwen3_5VisionConfig(**vision_config)
        elif vision_config is None:
            vision_config = Qwen3_5VisionConfig()
        # WHY 无 `vision_config.deepstack_visual_indexes = []`（对比 Qwen3-VL ViT 版）：
        # Qwen3_5VisionConfig 父类用 AttributeError() 哨兵显式屏蔽此字段，设置会报 AttributeError
        # LlavaOnevision.pack_image_features 用这个作为 tile 像素大小（不是 patch_size）
        vision_config.image_size = tile_size

        if isinstance(text_config, dict):
            text_config = Qwen3Config(**text_config)
        elif text_config is None:
            text_config = Qwen3Config()

        # 父类默认 select_strategy='default' 会跳首 token (CLS) — Qwen3.5 ViT 无 CLS 必须用 'full'
        kwargs.setdefault("vision_feature_select_strategy", "full")
        super().__init__(vision_config=vision_config, text_config=text_config, **kwargs)


class Qwen3_5ViTBackbone(Qwen3_5VisionModel):
    """Qwen3.5 Vision 去除原生 patch_merger 的 backbone 版本（V6 final_layernorm fix (2026-05-16): append final_layernorm）。

    构造时把 merger.norm 的预训练权重抠到 final_layernorm，然后 `del self.merger`
    释放 ~37M 参数（保留 norm 的 LN 焊到末端做 post_layernorm 角色，丢弃 spatial
    shuffle + linear_fc1/fc2，那对应 LlavaOV projector 的职责）。

    架构对称（V6 final_layernorm fix (2026-05-16) 修复）：
      SigLIP2:    encoder → post_layernorm → last_hidden_state → LlavaOV MLP → LLM
      V6 final_layernorm fix (2026-05-16):     encoder → final_layernorm → last_hidden_state → LlavaOV MLP → LLM

    forward 跑完 transformer blocks 后过 final_layernorm，再返回。下游 LlavaOnevision
    pack_image_features 的 AnyRes 2×2 pool 接管原 merger 的空间合并职责。

    输入输出契约与父类 Qwen3_5VisionModel 一致（NaViT flat）：
      forward(hidden_states=[L, patch_dim], grid_thw=[N, 3])
      → BaseModelOutput(last_hidden_state=[L, hidden_size])

    forward 主体 1:1 对照 `Qwen3_5VisionModel.forward`（已无 deepstack loop，比
    Qwen3VLVisionModel.forward 更短），仅跳过末尾 `self.merger(x)`，改为 final_layernorm。
    """

    def __init__(self, config):
        super().__init__(config)
        # V6 final_layernorm fix (2026-05-16): extract merger.norm pretrained weights into final_layernorm.
        # Default init: even if ckpt 阶段没 inject final_layernorm.* (e.g. stock
        # bootstrap path that filters merger.*), final_layernorm 仍持有 merger.norm
        # 的预训练值, 不是 random — 这是防止 silent corruption 的兜底.
        ln_w = self.merger.norm.weight.detach().clone()
        ln_b = self.merger.norm.bias.detach().clone()
        del self.merger
        self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.final_layernorm.weight.data.copy_(ln_w)
        self.final_layernorm.bias.data.copy_(ln_b)

    def forward(self, hidden_states, grid_thw, **kwargs):
        hidden_states = self.patch_embed(hidden_states)

        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states = hidden_states + pos_embeds

        rotary_pos_emb = self.rot_pos_emb(grid_thw)
        seq_len, _ = hidden_states.size()
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())

        cu_seqlens = torch.repeat_interleave(
            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
        ).cumsum(dim=0, dtype=torch.int32)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        for blk in self.blocks:
            hidden_states = blk(
                hidden_states,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        # V6 final_layernorm fix (2026-05-16): appended final LayerNorm — mirrors SigLIP2's post_layernorm.
        # Per-token affine; layout-invariant (reorder happens in adapter).
        hidden_states = self.final_layernorm(hidden_states)

        return BaseModelOutput(last_hidden_state=hidden_states)


class Qwen3_5ViTAsSiglipAdapter(nn.Module):
    """SigLIP 契约 → NaViT 契约的翻译层。持有 Qwen3_5ViTBackbone。

    对外暴露 SigLIP 式 forward(pixel_values=[N,3,H,W]) → BaseModelOutputWithPooling，
    供 LlavaOnevision 消费；对内按官方 _preprocess 的 reshape 链把 pixel_values
    转成 NaViT flat + grid_thw 喂给 Backbone。

    reshape 链 1:1 照抄 transformers 官方 Qwen2VLImageProcessorFast._preprocess
    (video_processing_qwen3_vl.py L227-252) —— Qwen3.5 无独立 image_processor，复用 Qwen3-VL 格式。
    """

    def __init__(self, vision_config: Qwen3_5VisionConfig):
        super().__init__()
        self.vision = Qwen3_5ViTBackbone(vision_config)
        self.config = vision_config

    @property
    def dtype(self):
        return next(self.parameters()).dtype

    @property
    def device(self):
        return next(self.parameters()).device

    def _flatten_navit(self, pixel_values: torch.Tensor):
        """[N, 3, H, W] → (flat=[N*L, patch_dim], grid_thw=[N, 3], shape=(N, L)).

        L = grid_t * grid_h * grid_w = 1 * (H/16) * (W/16)
        patch_dim = C * temporal_patch_size * patch_size^2 = 3 * 2 * 16 * 16 = 1536
        """
        pixel_values = pixel_values.to(dtype=self.dtype)
        tps = self.config.temporal_patch_size
        ps = self.config.patch_size
        ms = self.config.spatial_merge_size

        patches = pixel_values.unsqueeze(1)
        # 对单帧图像 T=1, pad=1 → expand 一帧使 T 整除 temporal_patch_size，
        # Conv3d 在复制帧上退化为等效 2D Conv（数学无损）
        T = patches.shape[1]
        pad = -T % tps
        if pad:
            repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
            patches = torch.cat((patches, repeats), dim=1)

        batch_size, t, channel, H, W = patches.shape
        grid_t = t // tps
        grid_h = H // ps
        grid_w = W // ps

        patches = patches.view(
            batch_size, grid_t, tps, channel,
            grid_h // ms, ms, ps,
            grid_w // ms, ms, ps,
        )
        patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
        flatten_patches = patches.reshape(
            batch_size,
            grid_t * grid_h * grid_w,
            channel * tps * ps * ps,
        )

        seq_len = grid_t * grid_h * grid_w
        flat = flatten_patches.reshape(batch_size * seq_len, -1)
        # on-device 构造小 tensor 再 expand，host→GPU 同步量 O(3) 而非 O(N*3)
        grid_unit = torch.tensor(
            [grid_t, grid_h, grid_w], dtype=torch.int32, device=pixel_values.device,
        )
        grid_thw = grid_unit.unsqueeze(0).expand(batch_size, -1).contiguous()
        return flat, grid_thw, (batch_size, seq_len)

    def forward(
        self,
        pixel_values: torch.Tensor,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> BaseModelOutputWithPooling:
        flat, grid_thw, (N, S) = self._flatten_navit(pixel_values)
        vision_out = self.vision(flat, grid_thw=grid_thw)
        hidden = vision_out.last_hidden_state.view(N, S, -1)

        # Block-merge → row-major reorder before handing to LlavaOnevision.
        # Internally the ViT runs in Qwen NaViT block-merge layout (pretrained
        # pos_embed + RoPE contract); downstream LlavaOV `pack_image_features`
        # (multi-tile AnyRes path, view(num_patch_h, num_patch_w, h, w, -1))
        # and `apply_pooling` (video path, view(B, h, w, -1) + bilinear) BOTH
        # assume row-major. Without this reorder, the multi-tile/video spatial
        # pool pulls together tokens that are NOT spatially adjacent — silent
        # corruption that doesn't fire on S1 single-tile path (line 348-351 of
        # modeling_llava_onevision.py just flattens [N,D] verbatim) but kills
        # S2 / eval quality.
        grid_h = int(grid_thw[0, 1].item())
        grid_w = int(grid_thw[0, 2].item())
        ms = getattr(self.config, "spatial_merge_size", 2)
        layout_perm = block_merge_to_row_major_permutation(
            grid_h, grid_w, ms=ms, device=hidden.device,
        )
        hidden = hidden[:, layout_perm, :]

        return BaseModelOutputWithPooling(
            last_hidden_state=hidden,
            # LlavaOnevision 索引 hidden_states[vision_feature_layer=-1]；tuple 长度 1 足够
            hidden_states=(hidden,),
            pooler_output=None,
        )


class LlavaQwen3_5ViTMultiModalProjector(nn.Module):
    """标准 LlavaOnevision projector（V6 final_layernorm fix (2026-05-16): pre_norm → Identity）。

    V6 final_layernorm fix (2026-05-16) 修复后, encoder 末端已自带 final_layernorm（与 SigLIP2 post_layernorm 对称），
    projector 不再需要补 LN — pre_norm 改为 nn.Identity，对齐 SigLIP2 plugin 的
    LlavaOnevision stock projector 结构（裸 linear_1 → GELU → linear_2），
    保证 SigLIP2 / Qwen3.5 / Qwen3-VL 三个 backbone 在 LlavaOV 设定下公平对比。

    （历史：V6.0.0~V6.0.4 时期 encoder 无 final LN，projector pre_norm 是补丁；
    现在补丁回到 encoder 内部，projector 回归 stock 形态。）
    """

    def __init__(self, config: LlavaQwen3_5ViTConfig):
        super().__init__()
        num_feature_layers = (
            1 if isinstance(config.vision_feature_layer, int) else len(config.vision_feature_layer)
        )
        vision_dim = config.vision_config.hidden_size * num_feature_layers
        text_dim = config.text_config.hidden_size
        bias = getattr(config, "multimodal_projector_bias", True)

        # V6 final_layernorm fix (2026-05-16): pre_norm = Identity (encoder 已自带 final_layernorm).
        self.pre_norm = nn.Identity()
        self.linear_1 = nn.Linear(vision_dim, text_dim, bias=bias)
        self.act = ACT2FN[config.projector_hidden_act]
        self.linear_2 = nn.Linear(text_dim, text_dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_2(self.act(self.linear_1(self.pre_norm(x))))


class LlavaQwen3_5ViTModel(LlavaOnevisionModel):
    """继承 LlavaOnevisionModel 但绕过其 __init__ 手动装配。

    父类 __init__ 调 `AutoModel.from_config(config.vision_config)` 会对 Qwen3_5VisionConfig
    抛 "Unrecognized configuration"（Qwen3.5 vision 没注册到 AutoModel）。手动装配避开
    这一步，同时省掉"先构造 Qwen3_5VisionModel 再被替换"的双重开销（~1.3GB init-peak）。

    装配顺序与父类一致：vision_tower / projector / image_newline / language_model / post_init。
    """

    config_class = LlavaQwen3_5ViTConfig

    def __init__(self, config: LlavaQwen3_5ViTConfig):
        # 跳过 LlavaOnevisionModel.__init__（AutoModel 不识别 Qwen3_5VisionConfig）
        LlavaOnevisionPreTrainedModel.__init__(self, config)
        self.vision_tower = Qwen3_5ViTAsSiglipAdapter(config.vision_config)
        self.multi_modal_projector = LlavaQwen3_5ViTMultiModalProjector(config)
        embed_std = 1 / math.sqrt(config.text_config.hidden_size)
        self.image_newline = nn.Parameter(
            torch.randn(config.text_config.hidden_size, dtype=self.dtype) * embed_std
        )
        self.vocab_size = config.text_config.vocab_size
        self.language_model = AutoModel.from_config(config.text_config)
        self.post_init()


class LlavaQwen3_5ViTForConditionalGeneration(LlavaOnevisionForConditionalGeneration):
    """继承 LlavaOnevisionForConditionalGeneration，只换 self.model。

    同样跳过父类 __init__（避免重复构造 LlavaOnevisionModel，根因见 LlavaQwen3_5ViTModel）。
    """

    config_class = LlavaQwen3_5ViTConfig

    def __init__(self, config: LlavaQwen3_5ViTConfig):
        LlavaOnevisionPreTrainedModel.__init__(self, config)
        self.model = LlavaQwen3_5ViTModel(config)
        self.lm_head = nn.Linear(
            config.text_config.hidden_size, config.text_config.vocab_size, bias=False
        )
        self.post_init()


AutoConfig.register(LlavaQwen3_5ViTConfig.model_type, LlavaQwen3_5ViTConfig)
AutoModelForCausalLM.register(LlavaQwen3_5ViTConfig, LlavaQwen3_5ViTForConditionalGeneration)


__all__ = [
    "LlavaQwen3_5ViTConfig",
    "Qwen3_5ViTBackbone",
    "Qwen3_5ViTAsSiglipAdapter",
    "LlavaQwen3_5ViTMultiModalProjector",
    "LlavaQwen3_5ViTModel",
    "LlavaQwen3_5ViTForConditionalGeneration",
]