"""Run the quantized ONNX models from this repository with ONNX Runtime.

Needs only `onnxruntime`, `tokenizers` and `numpy` (plus `huggingface_hub` to download by repo id).

    from inference import load, maxsim

    emb = load("nynefold/pplx-embed-v1-0.6b-ONNX")             # recommended variant
    q = emb.embed_queries(["how do I reset the router?"])
    d = emb.embed_documents(["Hold the reset button for ten seconds.", "The router has four ports."])
    print(q @ d.T)                                               # cosine similarities

    rr = load("nynefold/zerank-1-small-ONNX", variant="mix30-gptq-sq50-a4")
    print(rr.rerank("how do I reset the router?", ["Hold the reset button ...", "The router has ..."]))

    late = load("nynefold/pplx-embed-v2-late-0.6b-ONNX")         # late interaction: one vector per token
    docs = late.embed_documents(["Hold the reset button for ten seconds.", "The router has four ports."])
    print(maxsim(late.embed_queries(["how do I reset the router?"])[0], docs))

`load` takes a repo id or a local folder. A repo id downloads the variant into
~/.cache/nynefold/<repo id> (or $NYNEFOLD_CACHE/<repo id>) as plain files. All model-specific details (prompts, truncation, output
transform) come from `config.json`, which is part of the contract: the graphs expect
exactly these inputs.
"""
import json
import os
from pathlib import Path

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

CONFIG = "config.json"  # also the file the Hub counts downloads by


def _local(repo_or_dir, variant):
    p = Path(repo_or_dir)
    if p.is_dir():
        return p
    # Plain files in a local folder, not the Hub cache: its symlinks put model.onnx and model.onnx.data
    # in different blob folders, and ONNX Runtime refuses external data outside the model's folder.
    from huggingface_hub import hf_hub_download, snapshot_download
    dst = Path(os.environ.get("NYNEFOLD_CACHE", Path.home() / ".cache" / "nynefold")) / repo_or_dir
    cfg = json.loads(Path(hf_hub_download(repo_or_dir, CONFIG, local_dir=dst)).read_text())
    v = variant or cfg["recommended_variant"]
    return Path(snapshot_download(repo_or_dir, local_dir=dst,
                                  allow_patterns=["*.json", "*.txt", "*.jinja", f"{v}/*"]))


def load(repo_or_dir, variant=None, threads=None):
    """A ready model: Embedder or Reranker. variant=None picks the repo's recommended variant."""
    root = _local(repo_or_dir, variant)
    cfg = json.loads((root / CONFIG).read_text())
    variant = variant or cfg["recommended_variant"]
    vdir = root / variant
    tok_file = vdir / "tokenizer.json" if (vdir / "tokenizer.json").exists() else root / "tokenizer.json"
    so = ort.SessionOptions()
    if threads:  # otherwise ORT's default: one thread per physical core
        so.intra_op_num_threads = threads
    so.inter_op_num_threads = 1
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess = ort.InferenceSession(str(vdir / "model.onnx"), so, providers=["CPUExecutionProvider"])
    cls = {"embed": Embedder, "rerank": Reranker, "late": LateEmbedder}[cfg["kind"]]
    return cls(cfg, Tokenizer.from_file(str(tok_file)), sess)


class _Model:
    def __init__(self, cfg, tok, sess):
        self.cfg, self.tok, self.sess = cfg, tok, sess
        self.max_len = cfg["max_length"]
        self.pad_id = tok.token_to_id(cfg["pad_token"])

    def _ids(self, text, special=True):
        return self.tok.encode(text, add_special_tokens=special).ids

    def _run(self, ids, max_tokens=8192, max_batch=64, per_token=False):
        """Right-padded, length-sorted batches through the graph; one output row per input
        (per_token: a list with one [tokens, dim] array per input, padding removed)."""
        out = [None] * len(ids)
        order = sorted(range(len(ids)), key=lambda i: -len(ids[i]))
        i = 0
        while i < len(order):
            T = len(ids[order[i]])
            idx = order[i : i + max(1, min(max_batch, max_tokens // max(T, 1)))]
            inp = np.full((len(idx), T), self.pad_id, dtype=np.int64)
            msk = np.zeros((len(idx), T), dtype=np.int64)
            for r, j in enumerate(idx):
                inp[r, : len(ids[j])] = ids[j]
                msk[r, : len(ids[j])] = 1
            y = self.sess.run(None, {"input_ids": inp, "attention_mask": msk})[0]
            for r, j in enumerate(idx):
                out[j] = y[r, : len(ids[j])] if per_token else y[r]
            i += len(idx)
        return out if per_token else np.stack(out)


class Embedder(_Model):
    """Pooling is inside the graph. Returns L2-normalized float32 vectors; compare by dot product."""

    def encode_ids(self, texts, role):
        prefix = self.cfg["query_prefix"] if role == "query" else self.cfg["document_prefix"]
        out = []
        for t in texts:
            x = self._ids(prefix + t, special=self.cfg["add_special_tokens"])
            if len(x) > self.max_len:
                # last-token pooling reads the final (end) token, so keep it
                x = x[: self.max_len - 1] + x[-1:] if self.cfg["pooling"] == "last" else x[: self.max_len]
            out.append(x)
        return out

    def _embed(self, texts, role):
        e = self._run(self.encode_ids(list(texts), role)).astype(np.float32)
        if self.cfg["output"] == "int8_tanh":  # the base model's native int8 output: round(127 * tanh(x))
            e = np.clip(np.round(127 * np.tanh(e)), -128, 127)
        return e / np.linalg.norm(e, axis=1, keepdims=True)

    def embed_queries(self, texts):
        return self._embed(texts, "query")

    def embed_documents(self, texts):
        return self._embed(texts, "document")


class Reranker(_Model):
    """Scores (query, document) pairs; higher is more relevant."""

    def encode_ids(self, pairs):
        c = self.cfg
        out = []
        if c["format"] == "chat":  # prompt from the template; if too long, only the document's tokens are cut
            head, tail = c["template"].split("{document}")
            suf = self._ids(tail, False)
            for q, d in pairs:
                q, d = q[: c["query_max_chars"]].strip(), d[: c["document_max_chars"]].strip()
                x = self._ids(head.format(query=q) + d + tail, special=False)
                if len(x) > self.max_len:  # keep the end of the template: the score is read at its last token
                    pre = self._ids(head.format(query=q), False)
                    x = pre + self._ids(d, False)[: max(0, self.max_len - len(pre) - len(suf))] + suf
                out.append(x)
        else:  # fixed prefix and suffix around a body; only the body is truncated
            pre, suf = self._ids(c["prefix"], False), self._ids(c["suffix"], False)
            room = self.max_len - len(pre) - len(suf)
            for q, d in pairs:
                body = self._ids(c["body"].format(instruction=c["instruction"], query=q, document=d), False)
                out.append(pre + body[:room] + suf)
        return out

    def logits(self, query, documents):
        return self._run(self.encode_ids([(query, d) for d in documents])).reshape(-1).astype(np.float32)

    def rerank(self, query, documents):
        """Relevance scores in [0, 1], one per document, in input order."""
        z = self.logits(query, documents) * self.cfg["score_scale"]
        return 1.0 / (1.0 + np.exp(-z))


class LateEmbedder(_Model):
    """Late interaction (ColBERT-style, pplx-embed-v2-late): one L2-normalized 128-d vector per token.
    Score a query against documents with `maxsim`. Models that share a token space (the 0.6B and 9B)
    can be mixed: index documents with one, embed queries with the other."""

    def __init__(self, cfg, tok, sess):
        super().__init__(cfg, tok, sess)
        self.skip = np.array(sorted({i for i in (tok.token_to_id(t) for t in cfg["document_skiplist"]) if i is not None}))

    def encode_ids(self, texts, role):
        c = self.cfg
        prefix, n = ((c["query_prefix"], c["query_max_length"]) if role == "query"
                     else (c["document_prefix"], c["document_max_length"]))
        return [self._ids(prefix + t, special=c["add_special_tokens"])[:n] for t in texts]

    def _embed(self, texts, role):
        ids = self.encode_ids(list(texts), role)
        out = [v.astype(np.float32) for v in self._run(ids, per_token=True)]
        if role == "document":  # punctuation tokens don't take part in scoring
            out = [v[~np.isin(x, self.skip)] for v, x in zip(out, ids)]
        return out

    def embed_queries(self, texts):
        """A list of [query tokens, 128] float32 arrays."""
        return self._embed(texts, "query")

    def embed_documents(self, texts):
        """A list of [document tokens, 128] float32 arrays (punctuation tokens removed): the index."""
        return self._embed(texts, "document")


def maxsim(query, documents):
    """Late-interaction scores of one query ([tokens, dim]) against documents (a list of [tokens, dim]):
    for each query token the best-matching document token's dot product, summed. Higher is more relevant."""
    return np.array([float((query @ d.T).max(1).sum()) for d in documents], dtype=np.float32)