"""ALLSmartIntent: what a chat message wants and what a good reply needs, in one pass.

    pip install torch transformers safetensors huggingface_hub
    python beispiel.py "und morgen?"

Input: one sentence about the purpose of the bot, the last messages (optional) and the new message.
Output: dialogue act, reference to the conversation, needs, mood, safety signals and language.
"""
import json, sys

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from torch import nn
from transformers import AutoModel, AutoTokenizer

MODELL = "allsmart/Intent"
MAX_VERLAUF_ZEICHEN = 600   # per earlier message; longer ones keep start and end (the question is usually at the end)
GERAET = "cuda" if torch.cuda.is_available() else "cpu"

meta = json.load(open(hf_hub_download(MODELL, "intent.json")))
AKT, BEZUG, JA_NEIN, SPRACHEN, SCHWELLEN = (meta[k] for k in ("akt", "bezug", "ja_nein", "sprachen", "schwellen"))


def eingabe(zweck: str, verlauf: list[dict], nachricht: str) -> str:
    """Same format as in training: new message first, then purpose, then history newest first.
    Truncation therefore drops the oldest history, never the message."""
    teile = [f"Nachricht: {nachricht.strip()}", f"Zweck: {zweck.strip()}"]
    if verlauf:
        teile.append("Verlauf (neueste zuerst):")
        for v in reversed(verlauf):
            t = v["text"].strip()
            if len(t) > MAX_VERLAUF_ZEICHEN:
                t = t[:150] + " … " + t[-(MAX_VERLAUF_ZEICHEN - 150):]
            teile.append(f"{v['rolle']}: {t}")
    else:
        teile.append("Verlauf: keiner")
    return "\n".join(teile)


class Intent(nn.Module):
    """mmBERT-base, mean over the tokens, four linear heads."""

    def __init__(self):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(MODELL)
        h = self.encoder.config.hidden_size
        self.koepfe = nn.ModuleDict({"akt": nn.Linear(h, len(AKT)), "bezug": nn.Linear(h, len(BEZUG)),
                                     "ja_nein": nn.Linear(h, len(JA_NEIN)), "sprache": nn.Linear(h, len(SPRACHEN))})
        self.koepfe.load_state_dict(load_file(hf_hub_download(MODELL, "koepfe.safetensors")))

    def forward(self, input_ids, attention_mask):
        x = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        m = attention_mask.unsqueeze(-1).to(x.dtype)
        v = (x * m).sum(1) / m.sum(1).clamp(min=1)
        return {k: kopf(v) for k, kopf in self.koepfe.items()}


tok = AutoTokenizer.from_pretrained(MODELL)
modell = Intent().to(GERAET).eval()


@torch.no_grad()
def entscheide(faelle: list[dict], batch: int = 32) -> list[dict]:
    """faelle: [{"zweck": str, "verlauf": [{"rolle": "nutzer"|"bot", "text": str}], "nachricht": str}]"""
    aus = []
    for i in range(0, len(faelle), batch):
        teil = faelle[i:i + batch]
        k = tok([eingabe(f["zweck"], f.get("verlauf", []), f["nachricht"]) for f in teil], truncation=True,
                max_length=512, padding=True, return_tensors="pt").to(GERAET)
        with torch.autocast(GERAET, dtype=torch.bfloat16, enabled=GERAET == "cuda"):
            o = modell(k["input_ids"], k["attention_mask"])
        akt, bez, spr = (o[n].float().softmax(-1).cpu() for n in ("akt", "bezug", "sprache"))
        jn = o["ja_nein"].float().sigmoid().cpu().tolist()
        for j in range(len(teil)):
            p = dict(zip(JA_NEIN, jn[j]))
            aus.append({
                "akt": AKT[akt[j].argmax()], "bezug": BEZUG[bez[j].argmax()], "sprache": SPRACHEN[spr[j].argmax()],
                "ja": [n for n in JA_NEIN if p[n] >= SCHWELLEN[n]],   # heads above their threshold
                "werte": {n: round(x, 3) for n, x in p.items()},
            })
    return aus


if __name__ == "__main__":
    zweck = "Flight assistant on an airline website: helps passengers with questions about their flight."
    verlauf = [{"rolle": "nutzer", "text": "Wie wird das Wetter heute in Wien bei der Landung?"},
               {"rolle": "bot", "text": "In Wien sind es bei der Landung um 14:20 etwa 18 Grad, leicht bewölkt."}]
    nachricht = sys.argv[1] if len(sys.argv) > 1 else "und morgen?"
    print(json.dumps(entscheide([{"zweck": zweck, "verlauf": verlauf, "nachricht": nachricht}])[0],
                     ensure_ascii=False, indent=1))