"""MiniCPM5-V (unofficial): MiniCPM5-1B LLM + SigLIP2 vision encoder + trained projector.

Not affiliated with OpenBMB. Vision encoder wiring adapted from the nanoVLM project (MIT).
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

from vision_encoder import ViT, ModalityProjector, VIT_IMG_SIZE

LLM_REPO = "openbmb/MiniCPM5-1B"
IMAGE_TOKEN = "<|image|>"
IMAGE_TOKEN_LEN = 64  # (512/16)^2 / 4^2 = 1024/16 = 64 tokens per image


class MiniCPM5V(nn.Module):
    def __init__(self, freeze_llm=True, freeze_vision=True):
        super().__init__()
        print("Loading MiniCPM5-1B...")
        self.llm = AutoModelForCausalLM.from_pretrained(LLM_REPO, dtype=torch.bfloat16)
        self.tokenizer = AutoTokenizer.from_pretrained(LLM_REPO)

        num_added = self.tokenizer.add_special_tokens({"additional_special_tokens": [IMAGE_TOKEN]})
        if num_added:
            self.llm.resize_token_embeddings(len(self.tokenizer))
        self.image_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)

        print("Loading SigLIP2 vision encoder...")
        self.vision_encoder = ViT.from_pretrained().to(torch.bfloat16)
        self.projector = ModalityProjector(self.llm.config.hidden_size)

        if freeze_vision:
            for p in self.vision_encoder.parameters():
                p.requires_grad_(False)
        if freeze_llm:
            for p in self.llm.parameters():
                p.requires_grad_(False)

    def trainable_parameters(self):
        return [p for p in self.parameters() if p.requires_grad]

    def _embed(self, input_ids, pixel_values):
        token_embd = self.llm.get_input_embeddings()(input_ids)
        if pixel_values is not None:
            with torch.set_grad_enabled(any(p.requires_grad for p in self.vision_encoder.parameters())):
                vis_feats = self.vision_encoder(pixel_values)
            image_embd = self.projector(vis_feats).to(token_embd.dtype)
            mask = (input_ids == self.image_token_id).unsqueeze(-1)
            flat = image_embd.reshape(-1, image_embd.size(-1))
            token_embd = token_embd.masked_scatter(mask, flat)
        return token_embd

    def forward(self, input_ids, attention_mask=None, pixel_values=None, labels=None):
        inputs_embeds = self._embed(input_ids, pixel_values)
        out = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels)
        return out

    @torch.no_grad()
    def generate(self, input_ids, attention_mask=None, pixel_values=None, max_new_tokens=128, min_new_tokens=5, **kw):
        inputs_embeds = self._embed(input_ids, pixel_values)
        return self.llm.generate(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            min_new_tokens=min_new_tokens,
            **kw,
        )

    def make_prompt_ids(self, prompt_text, device):
        """Builds `<image>*64 + prompt` ids for a single image + text prompt."""
        image_prefix = IMAGE_TOKEN * IMAGE_TOKEN_LEN
        text = image_prefix + "\n" + prompt_text
        ids = self.tokenizer(text, return_tensors="pt").input_ids.to(device)
        return ids