"""BaguettotronVLM model — self-contained for HuggingFace Hub."""
from __future__ import annotations

import torch
import torch.nn as nn
from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
)
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_baguettotron_vlm import BaguettotronVLMConfig


class PixelUnshuffleProjector(nn.Module):
    """Reduces ViT tokens 4× via PixelUnshuffle then projects to LLM dim."""

    def __init__(self, in_dim: int, out_dim: int, factor: int):
        super().__init__()
        self.factor = factor
        self.unshuffle = nn.PixelUnshuffle(factor)
        self.mlp = nn.Sequential(
            nn.Linear(in_dim * factor * factor, out_dim),
            nn.GELU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, D = x.shape
        spatial = int(N ** 0.5)
        x = x.reshape(B, spatial, spatial, D).permute(0, 3, 1, 2)
        x = self.unshuffle(x)
        x = x.flatten(2).transpose(1, 2)
        return self.mlp(x)


class BaguettotronVLMForConditionalGeneration(PreTrainedModel):
    """
    BaguettotronVLM: InternViT-300M + PixelUnshuffle projector + Baguettotron-321M.

    ~628M total parameters. The same architecture is shipped for all three
    training stages; only the checkpoint weights and `config.chat_style`
    differ between them.

    Load with:
        from transformers import AutoModelForImageTextToText
        model = AutoModelForImageTextToText.from_pretrained(
            "andreagemelli/Baguettotron-VLM",
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
        )
    """

    config_class = BaguettotronVLMConfig
    _no_split_modules = ["InternVisionEncoderLayer", "LlamaDecoderLayer"]
    # Tell HF Trainer not to pass num_items_in_batch (loss handled internally)
    model_accepts_loss_kwargs: bool = False

    def __init__(self, config: BaguettotronVLMConfig):
        super().__init__(config)

        self.vit = AutoModel.from_pretrained(
            config.vit_model_id,
            dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
        self.projector = PixelUnshuffleProjector(
            in_dim=config.vit_hidden,
            out_dim=config.llm_hidden,
            factor=config.unshuffle_factor,
        )

        tokenizer = AutoTokenizer.from_pretrained(config.llm_model_id)
        tokenizer.add_special_tokens(
            {"additional_special_tokens": [config.image_token, "</image>"]}
        )
        raw_id = tokenizer.convert_tokens_to_ids(config.image_token)
        self.image_token_id: int = raw_id if isinstance(raw_id, int) else int(raw_id[0])

        self.llm = AutoModelForCausalLM.from_pretrained(
            config.llm_model_id, dtype=torch.bfloat16
        )
        self.llm.resize_token_embeddings(len(tokenizer))
        # Break weight tying — safetensors rejects shared-storage tensors
        self.llm.lm_head.weight = nn.Parameter(self.llm.lm_head.weight.data.clone())

        self._tokenizer = tokenizer

    def _init_weights(self, module: nn.Module) -> None:
        # Pretrained components are initialised from their respective hubs;
        # the projector weights come from the saved checkpoint — skip random init.
        pass

    # ------------------------------------------------------------------
    # Training interface
    # ------------------------------------------------------------------

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        labels: torch.Tensor | None = None,
        pixel_values: torch.Tensor | None = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        inputs_embeds = self.llm.get_input_embeddings()(input_ids)

        if pixel_values is not None:
            with torch.no_grad():
                vit_out = self.vit(pixel_values)
            image_features = vit_out.last_hidden_state
            if image_features.shape[1] == self.config.vit_tokens + 1:
                image_features = image_features[:, 1:, :]
            visual_tokens = self.projector(image_features.float())
            image_mask = input_ids == self.image_token_id
            inputs_embeds[image_mask] = visual_tokens.reshape(
                -1, self.config.llm_hidden
            ).to(inputs_embeds.dtype)

        return self.llm(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels,
            return_dict=True,
            use_cache=False,
        )

    # ------------------------------------------------------------------
    # Inference interface
    # ------------------------------------------------------------------

    @torch.no_grad()
    def chat(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        pixel_values: torch.Tensor | None = None,
        max_new_tokens: int = 256,
        repetition_penalty: float = 1.3,
        **generate_kwargs,
    ) -> str:
        """Inject visual tokens, generate autoregressively, return decoded string."""
        inputs_embeds = self.llm.get_input_embeddings()(input_ids)

        if pixel_values is not None:
            vit_out = self.vit(pixel_values)
            image_features = vit_out.last_hidden_state
            if image_features.shape[1] == self.config.vit_tokens + 1:
                image_features = image_features[:, 1:, :]
            visual_tokens = self.projector(image_features.float())
            image_mask = input_ids == self.image_token_id
            inputs_embeds[image_mask] = visual_tokens.reshape(
                -1, self.config.llm_hidden
            ).to(inputs_embeds.dtype)

        im_end_id = int(self._tokenizer.convert_tokens_to_ids("<|im_end|>"))
        output_ids = self.llm.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=repetition_penalty,
            eos_token_id=[self._tokenizer.eos_token_id, im_end_id],
            **generate_kwargs,
        )
        decoded = self._tokenizer.decode(output_ids[0], skip_special_tokens=False)
        if "<|im_end|>" in decoded:
            decoded = decoded[: decoded.index("<|im_end|>")]
        return decoded.strip()