"""Load and run the model in this folder: per question, a probability distribution over its options.

    from typed_decisions import TypedDecisions
    m = TypedDecisions.load(".")                      # this folder, or a downloaded snapshot of it
    m.predict("invoice_processing", state_dict_or_json)
    # {"discrepancy_severity": {"0": 0.71, "1": 0.22, ...}, "disposition": {...}, ...}
"""
from __future__ import annotations

import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from torch import nn
from transformers import AutoConfig, AutoModel, AutoTokenizer


class TypedDecisions(nn.Module):
    def __init__(self, folder: Path, cfg: dict):
        super().__init__()
        enc_cfg = AutoConfig.from_pretrained(folder)
        self.enc = AutoModel.from_config(enc_cfg)
        h = enc_cfg.hidden_size
        self.schema: dict[str, dict[str, list[str]]] = cfg["schema"]
        self.heads = nn.ModuleDict({f"{wf}/{q}": nn.Linear(h, len(o)) for wf, qs in self.schema.items() for q, o in qs.items()})
        self.max_len = cfg["max_len"]
        self.tok = AutoTokenizer.from_pretrained(folder)

    @classmethod
    def load(cls, folder: str | Path, device: str = "cuda" if torch.cuda.is_available() else "cpu") -> "TypedDecisions":
        folder = Path(folder)
        m = cls(folder, json.loads((folder / "typed_decisions.json").read_text()))
        m.load_state_dict(load_file(str(folder / "model.safetensors")))
        return m.to(device).eval()

    @torch.no_grad()
    def predict(self, workflow: str, state: dict | str) -> dict[str, dict[str, float]]:
        """The input is the workflow name, a newline, then the state as compact JSON (as in training)."""
        s = json.dumps(json.loads(state) if isinstance(state, str) else state, ensure_ascii=False)
        dev = next(self.parameters()).device
        enc = self.tok([f"{workflow}\n{s}"], truncation=True, max_length=self.max_len, return_tensors="pt").to(dev)
        with torch.autocast(dev.type, dtype=torch.bfloat16, enabled=dev.type == "cuda"):
            hs = self.enc(**enc).last_hidden_state
        m = enc.attention_mask.unsqueeze(-1).to(hs.dtype)
        z = ((hs * m).sum(1) / m.sum(1)).float()
        return {q: dict(zip(opts, torch.softmax(self.heads[f"{workflow}/{q}"](z)[0], -1).tolist()))
                for q, opts in self.schema[workflow].items()}