"""
Vision Encoder — runs the ViT image encoder and MLP projector for VLMs.

Supports:
  - LLaVA-1.5:  CLIP ViT-L/14@336px  +  2-layer MLP projector
  - InternVL2:  InternViT-300M/6B     +  MLP projector
  - Qwen2-VL:   Built-in ViT          +  no separate projector

The vision encoder is small (~300MB–1GB) and stays permanently in GPU memory,
unlike the LLM backbone layers which are disk-offloaded.

Usage::

    encoder = VisionEncoder.from_tensors(model_path, config, device, compute_dtype)
    # pixel_values: [batch, 3, height, width]
    image_embeds = encoder.encode(pixel_values)
    # image_embeds: [batch, num_patches, lm_hidden_size]
"""

from __future__ import annotations

import logging
import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    h = x.float()
    h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps)
    return (h * weight.float()).to(x.dtype)


def _layer_norm(x: torch.Tensor, weight: torch.Tensor,
                bias: Optional[torch.Tensor], eps: float = 1e-6) -> torch.Tensor:
    return F.layer_norm(x, weight.shape, weight, bias, eps)


# ---------------------------------------------------------------------------
# CLIP-style ViT (used by LLaVA-1.5, InternVL2-small)
# ---------------------------------------------------------------------------


class CLIPAttention(nn.Module):
    def __init__(self, hidden: int, heads: int):
        super().__init__()
        self.heads = heads
        self.head_dim = hidden // heads
        self.scale = self.head_dim ** -0.5
        self.q_proj = nn.Linear(hidden, hidden, bias=True)
        self.k_proj = nn.Linear(hidden, hidden, bias=True)
        self.v_proj = nn.Linear(hidden, hidden, bias=True)
        self.out_proj = nn.Linear(hidden, hidden, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, s, _ = x.shape
        q = self.q_proj(x).view(b, s, self.heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(b, s, self.heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(b, s, self.heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.transpose(1, 2).contiguous().view(b, s, -1)
        return self.out_proj(out)


class CLIPMLP(nn.Module):
    def __init__(self, hidden: int, intermediate: int):
        super().__init__()
        self.fc1 = nn.Linear(hidden, intermediate, bias=True)
        self.fc2 = nn.Linear(intermediate, hidden, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(F.gelu(self.fc1(x), approximate="tanh"))


class CLIPEncoderLayer(nn.Module):
    def __init__(self, hidden: int, heads: int, intermediate: int, eps: float = 1e-5):
        super().__init__()
        self.self_attn = CLIPAttention(hidden, heads)
        self.layer_norm1 = nn.LayerNorm(hidden, eps=eps)
        self.mlp = CLIPMLP(hidden, intermediate)
        self.layer_norm2 = nn.LayerNorm(hidden, eps=eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.self_attn(self.layer_norm1(x))
        x = x + self.mlp(self.layer_norm2(x))
        return x


class CLIPVisionTransformer(nn.Module):
    """
    CLIP ViT-L/14 or ViT-B/16 encoder.
    Produces [batch, num_patches, hidden_size] patch features.
    CLS token is excluded (LLaVA drops it; configurable via select_layer).
    """

    def __init__(
        self,
        image_size: int = 336,
        patch_size: int = 14,
        hidden: int = 1024,
        heads: int = 16,
        num_layers: int = 24,
        intermediate: int = 4096,
        eps: float = 1e-5,
        select_layer: int = -2,      # which layer output to return (-2 = penultimate)
    ):
        super().__init__()
        self.patch_size = patch_size
        self.select_layer = select_layer
        num_patches = (image_size // patch_size) ** 2

        # class_embedding must be a plain Parameter, NOT inside ModuleDict
        # (nn.ModuleDict rejects nn.Parameter — it only accepts nn.Module subclasses)
        self.class_embedding = nn.Parameter(torch.zeros(hidden))
        self.embeddings = nn.ModuleDict({
            "patch_embedding": nn.Conv2d(3, hidden, patch_size, patch_size, bias=False),
            "position_embedding": nn.Embedding(num_patches + 1, hidden),
        })
        self.pre_layrnorm = nn.LayerNorm(hidden, eps=eps)
        self.encoder = nn.ModuleList([
            CLIPEncoderLayer(hidden, heads, intermediate, eps)
            for _ in range(num_layers)
        ])
        self.post_layernorm = nn.LayerNorm(hidden, eps=eps)

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        b = pixel_values.shape[0]
        # Patch embedding
        patches = self.embeddings["patch_embedding"](pixel_values)
        patches = patches.flatten(2).transpose(1, 2)  # [b, n_patches, hidden]

        # Class token
        cls = self.class_embedding.expand(b, 1, -1)
        x = torch.cat([cls, patches], dim=1)  # [b, n_patches+1, hidden]

        # Positional embedding
        pos_ids = torch.arange(x.shape[1], device=x.device)
        x = x + self.embeddings["position_embedding"](pos_ids)
        x = self.pre_layrnorm(x)

        # Transformer layers — stop at select_layer and return patch features.
        # LLaVA uses features BEFORE post_layernorm (post_layernorm applies to
        # the CLS token only in the full CLIP model, not to patch tokens).
        target = len(self.encoder) + self.select_layer if self.select_layer < 0 else self.select_layer
        selected = x
        for i, layer in enumerate(self.encoder):
            x = layer(x)
            if i == target:
                selected = x

        return selected[:, 1:, :]  # drop CLS, [b, n_patches, hidden]


# ---------------------------------------------------------------------------
# MLP Projector (LLaVA-style: vision_hidden → lm_hidden)
# ---------------------------------------------------------------------------


class MLPProjector(nn.Module):
    """2-layer GELU MLP to map vision features to LM embedding space."""

    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.linear_1 = nn.Linear(in_dim, out_dim, bias=True)
        self.linear_2 = nn.Linear(out_dim, out_dim, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_2(F.gelu(self.linear_1(x), approximate="tanh"))


# ---------------------------------------------------------------------------
# VisionEncoder — top-level class used by AutoModel / InferenceEngine
# ---------------------------------------------------------------------------


class VisionEncoder:
    """
    VLM vision encoder: ViT + optional MLP projector.
    Kept permanently in GPU (small, ~300MB–1GB).

    Supports:
      arch="clip"   — CLIP ViT (LLaVA-1.5, InternVL2-small)
      arch="intern" — InternViT-6B (InternVL2-large)
      arch="qwen2"  — Qwen2-VL built-in ViT

    Produces image_embeds shaped [batch, num_image_tokens, lm_hidden_size]
    ready to be merged into the LLM input sequence.
    """

    def __init__(
        self,
        vit: nn.Module,
        projector: Optional[nn.Module],
        device: torch.device,
        compute_dtype: torch.dtype,
        image_token_id: int = 32000,
        num_image_tokens: int = 576,    # patches per image (336px CLIP: 576)
    ):
        self.vit = vit.to(device=device, dtype=compute_dtype)
        self.projector = projector.to(device=device, dtype=compute_dtype) if projector else None
        self.device = device
        self.compute_dtype = compute_dtype
        self.image_token_id = image_token_id
        self.num_image_tokens = num_image_tokens

    def to(self, device: torch.device) -> "VisionEncoder":
        """Move ViT and projector to device (in-place, returns self)."""
        self.vit = self.vit.to(device=device)
        if self.projector is not None:
            self.projector = self.projector.to(device=device)
        self.device = device
        return self

    def cpu(self) -> "VisionEncoder":
        return self.to(torch.device("cpu"))

    def parameters(self):
        """Yield all ViT and projector parameters (for size estimation)."""
        yield from self.vit.parameters()
        if self.projector is not None:
            yield from self.projector.parameters()

    @torch.no_grad()
    def encode(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """
        Encode a batch of images into language-space embeddings.

        Args:
            pixel_values: [batch, 3, H, W], normalised float tensor.

        Returns:
            [batch, num_image_tokens, lm_hidden_size]
        """
        pv = pixel_values.to(device=self.device, dtype=self.compute_dtype)
        features = self.vit(pv)                     # [b, n_patches, vis_hidden]
        if self.projector is not None:
            features = self.projector(features)     # [b, n_patches, lm_hidden]
        return features

    # ------------------------------------------------------------------
    # Factory — build VisionEncoder from already-loaded weight dicts
    # ------------------------------------------------------------------

    @classmethod
    def from_tensors(
        cls,
        tensors: dict[str, torch.Tensor],
        config: dict,
        device: torch.device,
        compute_dtype: torch.dtype,
        arch: str = "clip",
    ) -> "VisionEncoder":
        """
        Build a VisionEncoder from a flat {param_name: tensor} dict.

        Args:
            tensors:  All vision-related tensors (vision_tower + mm_projector).
            config:   Full HuggingFace config dict.
            arch:     "clip" | "intern" | "qwen2"
        """
        if arch == "clip":
            return cls._build_clip(tensors, config, device, compute_dtype)
        raise NotImplementedError(f"VisionEncoder.from_tensors: arch={arch!r} not yet supported")

    @classmethod
    def _build_clip(
        cls,
        tensors: dict[str, torch.Tensor],
        config: dict,
        device: torch.device,
        compute_dtype: torch.dtype,
    ) -> "VisionEncoder":
        """Build a CLIP ViT + MLP projector from weight tensors."""
        vcfg = config.get("vision_config", config.get("vision_tower_cfg", {}))
        image_size   = vcfg.get("image_size", 336)
        patch_size   = vcfg.get("patch_size", 14)
        vis_hidden   = vcfg.get("hidden_size", 1024)
        vis_heads    = vcfg.get("num_attention_heads", 16)
        vis_layers   = vcfg.get("num_hidden_layers", 24)
        vis_inter    = vcfg.get("intermediate_size", 4096)
        eps          = vcfg.get("layer_norm_eps", 1e-5)
        select_layer = config.get("mm_vision_select_layer", -2)

        lm_hidden = (
            config.get("hidden_size")
            or config.get("text_config", {}).get("hidden_size", 4096)
        )
        image_token_id = config.get("image_token_index",
                                    config.get("image_patch_token", 32000))
        num_patches = (image_size // patch_size) ** 2

        # --- Build ViT skeleton ---
        vit = CLIPVisionTransformer(
            image_size=image_size,
            patch_size=patch_size,
            hidden=vis_hidden,
            heads=vis_heads,
            num_layers=vis_layers,
            intermediate=vis_inter,
            eps=eps,
            select_layer=select_layer,
        )

        # --- Load ViT weights ---
        vit_prefix = _find_prefix(tensors, [
            "vision_tower.vision_model.",
            "vision_model.",
            "visual.",
        ])
        _load_clip_weights(vit, tensors, vit_prefix)

        # --- Build MLP projector ---
        # LLaVA-hf uses projector_hidden_act instead of mm_projector_type
        mm_type = (
            config.get("mm_projector_type")
            or ("mlp2x_gelu" if config.get("projector_hidden_act") else None)
            or "mlp2x_gelu"
        )
        projector: Optional[nn.Module] = None
        if mm_type and "mlp" in mm_type.lower():
            projector = MLPProjector(vis_hidden, lm_hidden)
            proj_prefix = _find_prefix(tensors, [
                "mm_projector.", "projector.", "multi_modal_projector.",
            ])
            _load_mlp_projector_weights(projector, tensors, proj_prefix)
        elif mm_type and "linear" in mm_type.lower():
            projector = nn.Linear(vis_hidden, lm_hidden, bias=True)
            proj_prefix = _find_prefix(tensors, ["mm_projector.", "multi_modal_projector."])
            for name, t in tensors.items():
                if name.startswith(proj_prefix):
                    key = name[len(proj_prefix):]
                    if "weight" in key:
                        projector.weight = nn.Parameter(t)
                    elif "bias" in key:
                        projector.bias = nn.Parameter(t)

        logger.info(
            "VisionEncoder built: CLIP ViT-%s (%d layers, %d hidden) + %s projector. "
            "%d patches per image.",
            "L" if vis_hidden >= 1024 else "B",
            vis_layers, vis_hidden, mm_type or "none", num_patches,
        )

        return cls(
            vit=vit,
            projector=projector,
            device=device,
            compute_dtype=compute_dtype,
            image_token_id=image_token_id,
            num_image_tokens=num_patches,
        )


# ---------------------------------------------------------------------------
# Weight loading helpers
# ---------------------------------------------------------------------------


def _find_prefix(tensors: dict[str, torch.Tensor], candidates: list[str]) -> str:
    """Return the first candidate prefix found in the tensor dict."""
    for prefix in candidates:
        if any(k.startswith(prefix) for k in tensors):
            return prefix
    return ""


def _load_clip_weights(
    vit: CLIPVisionTransformer,
    tensors: dict[str, torch.Tensor],
    prefix: str,
) -> None:
    """Map flat tensor dict into CLIPVisionTransformer modules."""
    p = prefix  # shorthand

    def get(key: str) -> Optional[torch.Tensor]:
        return tensors.get(p + key)

    # Embeddings
    if (t := get("embeddings.patch_embedding.weight")) is not None:
        vit.embeddings["patch_embedding"].weight = nn.Parameter(t)
    if (t := get("embeddings.position_embedding.weight")) is not None:
        vit.embeddings["position_embedding"].weight = nn.Parameter(t)
    if (t := get("embeddings.class_embedding")) is not None:
        vit.class_embedding = nn.Parameter(t)

    # Pre LayerNorm
    if (t := get("pre_layrnorm.weight")) is not None:
        vit.pre_layrnorm.weight = nn.Parameter(t)
    if (t := get("pre_layrnorm.bias")) is not None:
        vit.pre_layrnorm.bias = nn.Parameter(t)

    # Post LayerNorm
    if (t := get("post_layernorm.weight")) is not None:
        vit.post_layernorm.weight = nn.Parameter(t)
    if (t := get("post_layernorm.bias")) is not None:
        vit.post_layernorm.bias = nn.Parameter(t)

    # Encoder layers
    for i, layer in enumerate(vit.encoder):
        ep = f"encoder.layers.{i}."

        def eg(k: str) -> Optional[torch.Tensor]:
            return tensors.get(p + ep + k)

        # Self-attention
        for proj in ("q_proj", "k_proj", "v_proj", "out_proj"):
            if (t := eg(f"self_attn.{proj}.weight")) is not None:
                getattr(layer.self_attn, proj).weight = nn.Parameter(t)
            if (t := eg(f"self_attn.{proj}.bias")) is not None:
                getattr(layer.self_attn, proj).bias = nn.Parameter(t)

        # MLP
        for fc in ("fc1", "fc2"):
            if (t := eg(f"mlp.{fc}.weight")) is not None:
                getattr(layer.mlp, fc).weight = nn.Parameter(t)
            if (t := eg(f"mlp.{fc}.bias")) is not None:
                getattr(layer.mlp, fc).bias = nn.Parameter(t)

        # LayerNorms
        for ln_name, ln_module in [("layer_norm1", layer.layer_norm1),
                                   ("layer_norm2", layer.layer_norm2)]:
            if (t := eg(f"{ln_name}.weight")) is not None:
                ln_module.weight = nn.Parameter(t)
            if (t := eg(f"{ln_name}.bias")) is not None:
                ln_module.bias = nn.Parameter(t)


def _load_mlp_projector_weights(
    proj: MLPProjector,
    tensors: dict[str, torch.Tensor],
    prefix: str,
) -> None:
    """Load 2-layer MLP projector weights."""
    p = prefix

    def get(k: str) -> Optional[torch.Tensor]:
        return tensors.get(p + k)

    # LLaVA-1.5 uses: 0.weight, 0.bias, 2.weight, 2.bias  (nn.Sequential with GELU at index 1)
    # We map them to linear_1 and linear_2 of MLPProjector
    for src, dst in [("0", "linear_1"), ("2", "linear_2")]:
        if (t := get(f"{src}.weight")) is not None:
            getattr(proj, dst).weight = nn.Parameter(t)
        if (t := get(f"{src}.bias")) is not None:
            getattr(proj, dst).bias = nn.Parameter(t)

    # Alternative: some models use linear_1.weight / linear_2.weight directly
    for dst in ("linear_1", "linear_2"):
        if (t := get(f"{dst}.weight")) is not None:
            getattr(proj, dst).weight = nn.Parameter(t)
        if (t := get(f"{dst}.bias")) is not None:
            getattr(proj, dst).bias = nn.Parameter(t)


# ---------------------------------------------------------------------------
# Image token merging
# ---------------------------------------------------------------------------


def merge_image_tokens(
    input_ids: torch.Tensor,
    text_embeds: torch.Tensor,
    image_embeds: torch.Tensor,
    image_token_id: int,
) -> torch.Tensor:
    """
    Replace image placeholder tokens in text_embeds with image_embeds.

    Each occurrence of image_token_id in input_ids is replaced by the
    corresponding image's patch embeddings (num_image_tokens patches).

    Args:
        input_ids:     [1, seq_len]  — original token IDs
        text_embeds:   [1, seq_len, hidden_size]  — embedded tokens
        image_embeds:  [num_images, num_patches, hidden_size]  — vision features
        image_token_id: Token ID used as placeholder (e.g. 32000 for LLaVA)

    Returns:
        merged: [1, new_seq_len, hidden_size]
           new_seq_len = seq_len - num_images + num_images * num_patches
    """
    ids = input_ids[0]   # [seq_len]
    emb = text_embeds[0] # [seq_len, hidden_size]

    image_positions = (ids == image_token_id).nonzero(as_tuple=True)[0]
    if len(image_positions) == 0:
        return text_embeds

    segments = []
    prev = 0
    for img_idx, pos in enumerate(image_positions):
        pos = pos.item()
        if pos > prev:
            segments.append(emb[prev:pos])
        if img_idx < image_embeds.shape[0]:
            segments.append(image_embeds[img_idx])   # [num_patches, hidden_size]
        prev = pos + 1
    if prev < emb.shape[0]:
        segments.append(emb[prev:])

    merged = torch.cat(segments, dim=0).unsqueeze(0)  # [1, new_seq_len, hidden_size]
    return merged
