"""ChoiceEncoder: choice label -> semantic embedding in the core space.

Choices are *inputs*: "yes", "no", "database", "move up" are encoded from their
text, so the same weights serve 2, 4 or 10 options and unseen labels. With a
pretrained text encoder the choice embedding comes from the same backbone as the
question (shared weights), so "sì" and "yes" start close together.

`EmbeddingCache` memoises backbone outputs per string. With a frozen backbone a
question or a choice label is encoded once, ever: Snake's 4 moves and a fixed
question cost nothing at inference, only the core runs.
"""
from __future__ import annotations

from collections import OrderedDict

import torch
from torch import nn

from minijev.encoders.text import Projection, TextPooler


class ChoiceEncoder(nn.Module):
    def __init__(self, d: int, embed: nn.Embedding | None = None, max_len: int = 16, d_text: int | None = None):
        super().__init__()
        self.byte = embed is not None
        if self.byte:
            self.pooler = TextPooler(embed, d, max_len)
        else:
            self.proj = Projection(d_text, d, hidden=2 * d)

    def forward(self, ids=None, pooled=None):
        """byte: ids (..., L) -> (..., d); pretrained: backbone pooled (..., d_text) -> (..., d)."""
        return self.pooler(ids) if self.byte else self.proj(pooled)


class EmbeddingCache:
    """LRU cache: text -> backbone output (tokens, mask, pooled) as CPU tensors."""

    def __init__(self, max_items: int = 4096):
        self.max_items = max_items
        self.data: OrderedDict[str, tuple] = OrderedDict()
        self.hits = self.misses = 0

    def get_many(self, texts, compute):
        """compute(list[str]) -> (tok (n, L, dt), mask (n, L), pooled (n, dt)) for the missing texts."""
        missing = [t for t in dict.fromkeys(texts) if t not in self.data]
        self.hits += len(texts) - len(missing)
        self.misses += len(missing)
        if missing:
            tok, mask, pooled = compute(missing)
            for i, t in enumerate(missing):
                n = int(mask[i].sum())
                self.data[t] = (tok[i, :n].detach().cpu(), pooled[i].detach().cpu())
                if len(self.data) > self.max_items:
                    self.data.popitem(last=False)
        for t in texts:
            self.data.move_to_end(t)
        rows = [self.data[t] for t in texts]
        L = max(1, max(r[0].shape[0] for r in rows))
        dt = rows[0][1].shape[-1]
        tok = torch.zeros(len(rows), L, dt)
        mask = torch.zeros(len(rows), L, dtype=torch.bool)
        for i, (t, _) in enumerate(rows):
            tok[i, : t.shape[0]] = t
            mask[i, : t.shape[0]] = True
        return tok, mask, torch.stack([r[1] for r in rows])