"""MetaEncoder-30B: match a task against candidates, over text, images, and video.

    from modeling_metaencoder import MetaEncoder

    model = MetaEncoder.from_pretrained("<this repo>")

    results = model.match(
        task={"text": "Which photo shows a bicycle?", "image": context_img},
        candidates=[
            {"image": img_a},
            {"image": img_b, "instruction": "Represent the photo"},
            {"text": "a red bicycle leaning on a wall"},
            {"video": frames},
        ],
    )
    results[0].index, results[0].score, results[0].candidate

A *task* is what you are asking: text, images, and video in one object. Put the question,
query or instruction into its `text` — that text is what grounds the task to the candidates.
A *candidate* is the same shape, plus an optional `instruction` describing how that candidate
should be encoded.

`match` returns candidates ranked by similarity. For a corpus you query repeatedly, encode it
once with `encode` and pass the resulting tensor as `candidates`.

Prompt construction, media routing, padding and pooling are handled internally and must match
training exactly; see `_encode_batch` for the contract.
"""

from __future__ import annotations

import unicodedata
from dataclasses import dataclass
from typing import Any, Mapping, Sequence, Union

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoProcessor
from transformers.models.muse_glimmer.modeling_muse_glimmer import (
    MuseGlimmerForConditionalGeneration,
)

# Held constant across every input. A system turn that varies per task is a task-ID channel the
# model keys on instead of reading the instruction, which makes the embeddings brittle to
# unseen phrasings. It is emitted rather than omitted because the chat template's no-system
# branch interpolates the current date, which would make a corpus encoded yesterday not match a
# query encoded today.
DEFAULT_SYSTEM_INSTRUCTION = "Represent the user's input."

IMAGE_TOKEN = "<|patch|>"
VIDEO_TOKEN = "<|video|>"

DEFAULT_MAX_LENGTH = 8192
DEFAULT_BATCH_SIZE = 8

# Frames arrive already sampled, so the video processor must be told not to sample again. It
# defaults to do_sample_frames=True against an inferred fps (24 when it cannot tell), which both
# drops frames and writes wrong wall-clock timestamps into the prompt.
VIDEO_FPS = 2.0

Item = Union[str, Image.Image, Mapping[str, Any]]

_ITEM_KEYS = {"text", "instruction", "image", "images", "video"}


@dataclass
class Match:
    """One ranked candidate."""

    index: int       # position in the candidates you passed in
    score: float     # cosine similarity to the task, in [-1, 1]
    candidate: Any   # the candidate object itself, or None if you passed embeddings

    def __repr__(self):
        return f"Match(index={self.index}, score={self.score:.4f})"


def _ensure_punct(s: str) -> str:
    s = (s or "").strip()
    if s and not unicodedata.category(s[-1]).startswith("P"):
        s += "."
    return s


def _user_text(instruction: str | None, text: str | None) -> str:
    # Media markers are placeholders, not content: drop them and collapse whitespace, matching
    # the training-time strip exactly (replace with a space, then re-join on whitespace) so an
    # embedded marker cannot fuse two words together.
    body = (text or "").replace(IMAGE_TOKEN, " ").replace(VIDEO_TOKEN, " ")
    body = " ".join(body.split())
    instr = _ensure_punct(instruction)
    return " ".join(p for p in (instr, body) if p)


def _as_items(inputs: Item | Sequence[Item], instruction: str | None) -> list[dict]:
    """Normalise any accepted input shape into one dict per embedding."""
    if inputs is None:
        raise ValueError("encode() got no inputs")
    if isinstance(inputs, (str, Image.Image, Mapping)):
        inputs = [inputs]
    else:
        try:
            inputs = list(inputs)
        except TypeError:
            # Not a sequence either; wrap it so the per-item check reports what was expected
            # instead of leaking "object is not iterable".
            inputs = [inputs]

    items = []
    for i, raw in enumerate(inputs):
        if isinstance(raw, str):
            raw = {"text": raw}
        elif isinstance(raw, Image.Image):
            raw = {"image": raw}
        elif not isinstance(raw, Mapping):
            raise TypeError(
                f"item {i} is {type(raw).__name__}; expected str, PIL.Image, or a dict with "
                f"keys {sorted(_ITEM_KEYS)}"
            )

        unknown = set(raw) - _ITEM_KEYS
        if unknown:
            hint = ""
            if "videos" in unknown:
                hint = " (use 'video' with a list of frames for a single clip)"
            raise ValueError(f"item {i} has unknown key(s) {sorted(unknown)}{hint}")

        images = []
        if raw.get("image") is not None:
            images.append(raw["image"])
        if raw.get("images") is not None:
            got = raw["images"]
            images.extend(got if isinstance(got, (list, tuple)) else [got])
        images = [im for im in images if im is not None]

        video = raw.get("video")
        if video is not None:
            video = [f for f in (video if isinstance(video, (list, tuple)) else [video])
                     if f is not None]
            if not video:
                video = None

        items.append({
            "text": raw.get("text"),
            "instruction": raw.get("instruction", instruction),
            "images": images or None,
            "video": video,
        })
    return items


class MetaEncoder:
    """Embedding model over text, images, and video."""

    def __init__(self, model, processor):
        self.model = model
        self.processor = processor
        # Last-token pooling reads the final column, which is only the last real token under
        # left padding. The processor defaults to right.
        self.processor.tokenizer.padding_side = "left"

    @classmethod
    def from_pretrained(cls, path, device_map="auto", dtype=torch.bfloat16,
                        attn_implementation=None):
        if attn_implementation is None:
            attn_implementation = "sdpa"
            # flash-attn can be installed and still be unusable: it is CUDA-only, so probing
            # the import alone picks a backend that fails at init on a CPU load.
            if torch.cuda.is_available():
                try:
                    import flash_attn  # noqa: F401

                    attn_implementation = "flash_attention_2"
                except ImportError:
                    pass
        model = MuseGlimmerForConditionalGeneration.from_pretrained(
            path,
            dtype=dtype,
            device_map=device_map,
            attn_implementation=attn_implementation,
            low_cpu_mem_usage=True,
        )
        model.config.use_cache = False
        model.eval()
        # transformers keeps the rotary inv_freq buffers in fp32 whatever `dtype` is, but every
        # reported number was measured after a `.to(bf16)`, which rounds them too. Rounding them
        # here reproduces those numbers exactly; with fp32 inv_freq, image and video embeddings
        # drift by up to ~1.5% cosine.
        for module in model.modules():
            for name, buf in list(module.named_buffers(recurse=False)):
                if "inv_freq" in name and buf.dtype != dtype:
                    setattr(module, name, buf.to(dtype))
        return cls(model, AutoProcessor.from_pretrained(path))

    @property
    def embedding_dim(self) -> int:
        return self.model.config.text_config.hidden_size

    def encode(
        self,
        inputs: Item | Sequence[Item],
        instruction: str | None = None,
        batch_size: int = DEFAULT_BATCH_SIZE,
        max_length: int = DEFAULT_MAX_LENGTH,
    ) -> torch.Tensor:
        """Embed items directly. Returns L2-normalised float32 (n_items, embedding_dim).

        Most uses want `match` instead. Reach for `encode` to embed a corpus once and pass the
        tensor to `match` as `candidates`, or when you need the vectors themselves.

        `inputs` is one item or a list of items. An item is a string, a PIL image, or a dict:

            {"text": str, "image": img, "images": [img, ...], "video": [frame, ...],
             "instruction": str}

        All keys optional. Use `video` for pre-sampled clips (temporal patching); use `images`
        only for genuinely independent stills.
        """
        items = _as_items(inputs, instruction)
        if not items:
            return torch.empty(0, self.embedding_dim, device=self.model.device)
        out = [self._encode_batch(items[i:i + batch_size], max_length)
               for i in range(0, len(items), batch_size)]
        return torch.cat(out, dim=0)

    def match(
        self,
        task: Item,
        candidates: Sequence[Item] | torch.Tensor,
        top_k: int | None = None,
        batch_size: int = DEFAULT_BATCH_SIZE,
    ) -> list[Match]:
        """Rank `candidates` against `task`, best first.

        `task` is a single item — text, images, video. Put the question, query or instruction
        in its `text`; that is what grounds the task to the candidates.

        `candidates` is a list of items, each optionally carrying its own `instruction` for
        how it should be encoded. It may instead be a tensor of embeddings from a previous
        `encode` call, which is how you avoid re-encoding a corpus on every query.

        Returns a list of `Match(index, score, candidate)` sorted by descending score.
        """
        scores = self.score(task, candidates, batch_size=batch_size)
        order = torch.argsort(scores, descending=True)
        if top_k is not None:
            order = order[:top_k]
        originals = None if isinstance(candidates, torch.Tensor) else list(candidates)
        return [
            Match(index=int(i),
                  score=float(scores[i]),
                  candidate=None if originals is None else originals[int(i)])
            for i in order
        ]

    def score(
        self,
        task: Item,
        candidates: Sequence[Item] | torch.Tensor,
        batch_size: int = DEFAULT_BATCH_SIZE,
    ) -> torch.Tensor:
        """Cosine similarity of `task` to each candidate, in input order. Shape (n_candidates,)."""
        if isinstance(task, (list, tuple)):
            raise TypeError(
                "task must be a single item, not a list. To score several tasks, call match() "
                "per task (encode the candidates once and pass the tensor back in)."
            )
        task_emb = self.encode(task, batch_size=1)
        if isinstance(candidates, torch.Tensor):
            cand_emb = candidates
            if cand_emb.ndim != 2 or cand_emb.shape[1] != task_emb.shape[1]:
                raise ValueError(
                    f"candidate embeddings must be (n, {task_emb.shape[1]}), "
                    f"got {tuple(cand_emb.shape)}"
                )
            cand_emb = cand_emb.to(task_emb.device)
        else:
            cand_emb = self.encode(candidates, batch_size=batch_size)
        return (task_emb @ cand_emb.T)[0]

    @staticmethod
    def similarity(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Cosine similarity matrix (n_a, n_b). Inputs are already unit-norm."""
        return a @ b.T

    # ---------------------------------------------------------------- internals

    def _conversations(self, items):
        conversations, flat_images, flat_videos = [], [], []
        for it in items:
            content = []
            for im in (it["images"] or []):
                content.append({"type": "image", "image": im})
                flat_images.append(im)
            if it["video"] is not None:
                content.append({"type": "video", "video": it["video"]})
                flat_videos.append(it["video"])

            ut = _user_text(it["instruction"], it["text"])
            if ut:
                content.append({"type": "text", "text": ut})
            elif not content:
                content.append({"type": "text", "text": " "})

            conversations.append([
                {"role": "system",
                 "content": [{"type": "text", "text": DEFAULT_SYSTEM_INSTRUCTION}]},
                {"role": "user", "content": content},
            ])
        return conversations, flat_images, flat_videos

    @torch.no_grad()
    def _encode_batch(self, items, max_length):
        conversations, flat_images, flat_videos = self._conversations(items)
        chat_texts = [
            self.processor.apply_chat_template(c, add_generation_prompt=False, tokenize=False)
            for c in conversations
        ]

        kwargs = dict(
            text=chat_texts,
            images=flat_images or None,
            videos=flat_videos or None,
            padding=True,
            return_tensors="pt",
            # The chat template already emits bos_token and the tokenizer has
            # add_bos_token=True; leaving this on yields a doubled <|begin_of_text|>.
            add_special_tokens=False,
        )
        if flat_videos:
            kwargs["do_sample_frames"] = False
            from transformers.video_utils import VideoMetadata

            kwargs["video_metadata"] = [
                VideoMetadata(
                    total_num_frames=len(v),
                    fps=VIDEO_FPS,
                    duration=len(v) / VIDEO_FPS,
                    frames_indices=list(range(len(v))),
                )
                for v in flat_videos
            ]
        if flat_images or flat_videos:
            # Truncation would desynchronise the expanded <|patch|>/<|video|> runs from the pixels.
            kwargs["truncation"] = False
        else:
            kwargs["truncation"] = True
            kwargs["max_length"] = max_length

        inputs = self.processor(**kwargs)
        inputs["input_ids"] = inputs["input_ids"].long()
        inputs["attention_mask"] = inputs["attention_mask"].long()
        inputs = {k: (v.to(self.model.device) if hasattr(v, "to") else v)
                  for k, v in inputs.items()}

        am = inputs["attention_mask"]
        assert bool((am[:, -1] == 1).all()), "expected left padding"

        out = self.model(
            **inputs,
            return_dict=True,
            output_hidden_states=True,
            # ForConditionalGeneration otherwise materialises full-sequence logits over a
            # 202k vocab, which dominates memory; only hidden states are needed here.
            logits_to_keep=1,
        )
        pooled = out.hidden_states[-1][:, -1, :]
        return F.normalize(pooled.float(), p=2, dim=-1)