"""
Inference for Nawah-Router-v2 — an Arabic zero-shot router with a routing head.

The text and every category share one sequence. Each category's character span is mapped to token
indices through the tokenizer's offset mapping and mean-pooled into its own vector; a shared
scorer turns each into one logit, and the softmax runs over the categories actually supplied.
Because the scorer is shared across positions it reads category *content*, not slot index — which
is what makes the label set free text chosen at inference.

Layout is text-first, categories-second on purpose. The backbone is causal, so this ordering is
what lets every category token attend to the whole text; reversed, the categories would be
encoded blind to it.
"""
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

MAX_ROUTES = 9
MAX_LENGTH = 320


class RouterModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)
        h = self.backbone.config.hidden_size
        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))
        self.config = self.backbone.config

    def forward(self, input_ids, attention_mask, cat_pool, n_routes):
        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        logits = self.score(torch.bmm(cat_pool.to(hs.dtype), hs)).squeeze(-1)
        ar = torch.arange(logits.size(1), device=logits.device)[None, :]
        return logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)

    @classmethod
    def from_pretrained(cls, path, token=None):
        import os
        from huggingface_hub import hf_hub_download
        m = cls(path)
        w = (os.path.join(path, "router_model.pt") if os.path.isdir(path)
             else hf_hub_download(path, "router_model.pt", token=token))
        m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))
        return m.eval()


def build_text(text, routes):
    head = f"النص:\n{text}\n\nالفئات:\n"
    s, spans = head, []
    for c in routes:
        s += "- "
        spans.append((len(s), len(s) + len(c)))
        s += c + "\n"
    return s, spans


@torch.no_grad()
def route(model, tok, text, routes):
    """-> [{'route': str, 'score': float}] sorted high to low."""
    routes = [r for r in routes if r and r.strip()][:MAX_ROUTES]
    if not text.strip() or not routes:
        return []
    full, spans = build_text(text, routes)
    enc = tok(full, return_offsets_mapping=True, add_special_tokens=False,
              truncation=True, max_length=MAX_LENGTH)
    ids, offs = enc["input_ids"], enc["offset_mapping"]
    pool = torch.zeros(1, MAX_ROUTES, len(ids))
    for ci, (s, e) in enumerate(spans):
        idx = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]
        if idx:
            pool[0, ci, idx] = 1.0 / len(idx)
    logits = model(torch.tensor([ids]), torch.ones(1, len(ids), dtype=torch.long),
                   pool, torch.tensor([len(routes)]))
    probs = logits.softmax(-1)[0][: len(routes)].tolist()
    out = [{"route": r, "score": p} for r, p in zip(routes, probs)]
    return sorted(out, key=lambda x: -x["score"])


if __name__ == "__main__":
    M = "oddadmix/Nawah-Router-v2"
    tok = AutoTokenizer.from_pretrained(M)
    model = RouterModel.from_pretrained(M)
    for r in route(model, tok, "الطلب تأخر ساعة والسائق ما رد على الاتصال",
                   ["استفسار عن التوصيل", "شكوى تأخير", "مشكلة في الدفع"]):
        print(f"{r['score']:.3f}  {r['route']}")