"""Scientific-paper classification, public v2 format: LIST FIRST for choices; paper first for yes/no.
Apache-2.0. No generated text, no remote model code. The Transformers backend does not reuse prefix states;
use vllm_fast.FastPaperClassifier for high-throughput shared-taxonomy batches.
"""
from __future__ import annotations
import json
import random
from pathlib import Path

SYSTEM = ("You classify scientific papers. You get a paper (its title, and its abstract when available), then a "
          "question. For a list of options, answer with the code of the option that best describes the research "
          "content of the paper, following the classification's own meaning of each option. For a single topic, answer "
          "yes if the topic is a central subject of the paper, else no. On topic: the paper studies the topic itself as "
          "a main subject: the phenomenon, its causes, consequences, mechanisms, measurement, methods, or how it is "
          "managed, in any system or organism. Not on topic: the topic is absent or only mentioned in passing; or the "
          "words mean something else in this paper; or the paper studies a neighbouring subject without studying the "
          "topic itself.")
NONE = "None of these (outside the options listed)"
QUESTION = "Which option best describes the research content of this paper?"
SCHEME = "Classification of research"


def paper_text(title, abstract=None, venue=None, year=None):
    """Benchmark paper representation: title, optional venue/year, abstract limited to 1,500 characters.
    Title-only inputs are valid. A complete paper is not required or read by this model."""
    s = f"Title: {title.strip()}"
    if venue:
        s += f"\nVenue: {venue.strip()}"
    if year:
        s += f"\nYear: {year}"
    abstract = " ".join((abstract or "").split())
    return s + (f"\nAbstract: {abstract[:1500]}" if abstract else "\n(no abstract)")


def messages(user):
    return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}]


def option_line(code, o, describe=True, parents=True):
    extra = [x for x in ((o.get("about") if describe else None),
                         (f"in {o['parent']}" if parents and o.get("parent") else None)) if x]
    return f"{code}. {o['name']}" + (" — " + "; ".join(extra) if extra else "")


def head_text(scheme, question, options, codes, describe=True, parents=True):
    lines = "\n".join(option_line(c, o, describe, parents) for c, o in zip(codes, options))
    return f"## Classification: {scheme}\n{question}\nOptions:\n{lines}\n"


def tail_text(paper):
    return f"\n## Paper\n{paper}\n\nAnswer with the option code only."


def yesno_prompt(paper, topic, definition=None):
    d = f" (what it covers: {definition})" if definition else ""
    return f"{paper}\n\n## Topic: {topic}{d}\nIs this topic a central subject of the paper? Answer yes or no."


def read_config(model, revision=None):
    p = Path(model) / "classifier.json"
    if not p.exists():
        from huggingface_hub import hf_hub_download
        p = Path(hf_hub_download(model, "classifier.json", revision=revision))
    cfg = json.loads(p.read_text())
    if cfg.get("format") != "list-first-filler-v1":
        raise ValueError("This helper is for the v2/list-first release, not the earlier paper-first models.")
    return cfg


class PaperInputs:
    """Exact token construction. No truncation of taxonomy/options or the answer prefix.
    Long paper text alone is capped at paper_tokens. Oversized prompts raise ValueError instead of dropping labels.
    The prefix is tokenized once and padded with individual newline token IDs to a cache-block boundary.
    Do not concatenate these newlines as a string and retokenize: token merging changes the prompt."""
    def __init__(self, tokenizer, config, block=None, paper_tokens=700, max_tokens=4000):
        self.tok, self.cfg = tokenizer, config
        self.block = int(block or config.get("block", 544))
        self.paper_tokens, self.max_tokens = paper_tokens, max_tokens
        if self.block < 1 or paper_tokens < 1 or max_tokens > 4096:
            raise ValueError("block/paper_tokens must be positive; this helper supports max_tokens <= 4096")
        probe = tokenizer.apply_chat_template(messages("XXUSERXX"), tokenize=False,
                    add_generation_prompt=True, enable_thinking=False)
        self.pre, self.post = probe.split("XXUSERXX")
        filler = self.encode("\n")
        if len(filler) != 1 or filler[0] != config["filler"]:
            raise ValueError("Tokenizer/newline token differs from the release.")
        self.filler = filler[0]
        for code, token in zip(config["codes"], config["code_ids"]):
            if self.encode(code) != [token]:
                raise ValueError(f"Wrong tokenizer for answer code {code}")
        for word in ("yes", "no"):
            if self.encode(word) != [config[word]]:
                raise ValueError(f"Wrong tokenizer for {word}")

    def encode(self, text):
        return self.tok.encode(text, add_special_tokens=False)

    def normalize_options(self, options, none_option=False):
        out = []
        for x in options:
            o = {"name": x} if isinstance(x, str) else dict(x)
            if not isinstance(o.get("name"), str) or not o["name"].strip():
                raise ValueError("Every option needs a nonempty name")
            o["id"] = o.get("id", o["name"])
            if not isinstance(o["id"], str):
                raise ValueError("Option IDs must be strings")
            out.append(o)
        if none_option:
            out.append({"name": NONE, "id": "__none__"})
        if not 2 <= len(out) <= len(self.cfg["codes"]):
            raise ValueError(f"Supply 2–{len(self.cfg['codes'])} options; larger taxonomies require multiple questions")
        if len({o["id"] for o in out}) != len(out):
            raise ValueError("Option names/IDs must be unique; give each duplicate name a distinct 'id'")
        return out

    def prefix(self, opts, scheme, question, describe=True, parents=True):
        codes = self.cfg["codes"][:len(opts)]
        h = self.encode(self.pre + head_text(scheme, question, opts, codes, describe, parents))
        return h + [self.filler] * ((-len(h)) % self.block)

    def pick(self, paper, prefix):
        p = self.encode(paper)
        if len(p) > self.paper_tokens:
            paper = self.tok.decode(p[:self.paper_tokens])
        ids = prefix + self.encode(tail_text(paper) + self.post)
        return self.check_length(ids)

    def yesno(self, paper, topic, definition=None):
        p = self.encode(paper)
        if len(p) > self.paper_tokens:
            paper = self.tok.decode(p[:self.paper_tokens])
        text = self.tok.apply_chat_template(messages(yesno_prompt(paper, topic, definition)), tokenize=False,
                                            add_generation_prompt=True, enable_thinking=False)
        return self.check_length(self.encode(text))

    def check_length(self, ids):
        if len(ids) > self.max_tokens:
            raise ValueError(f"Prompt has {len(ids)} tokens; limit {self.max_tokens}. Shorten descriptions/labels, "
                             "reduce paper_tokens, or split the taxonomy. No options were silently truncated.")
        return ids


class ClassifierTasks:
    """Shared task API for the Transformers and vLLM backends."""
    def classify_many(self, papers, options, scheme=SCHEME, question=QUESTION, *,
                      describe=True, parents=True, none_option=False, orders=1, seed=0):
        papers = list(papers)
        if not papers:
            return []
        if not isinstance(orders, int) or orders < 1:
            raise ValueError("orders must be a positive integer")
        opts = self.inputs.normalize_options(options, none_option)
        rng = random.Random(seed);k = len(opts)
        total = [[0.0] * k for _ in papers]
        for order in range(orders):
            perm = list(range(k)) if order == 0 else rng.sample(range(k), k)
            prefix = self.inputs.prefix([opts[j] for j in perm], scheme, question, describe, parents)
            rows = [self.inputs.pick(p, prefix) for p in papers]
            probs = self._scores(rows, [self.cfg["code_ids"][:k]] * len(rows), shared_prefix=True)
            for t, p in zip(total, probs):
                for pos, j in enumerate(perm):
                    t[j] += p[pos] / orders
        return [dict(sorted(((o["id"], p) for o, p in zip(opts, vals)), key=lambda x: -x[1])) for vals in total]

    def classify(self, paper, options, scheme=SCHEME, question=QUESTION, **kwargs):
        return self.classify_many([paper], options, scheme, question, **kwargs)[0]

    def is_on_topic_many(self, papers, topic, definition=None):
        rows = [self.inputs.yesno(p, topic, definition) for p in papers]
        if not rows:
            return []
        return [r[0] for r in self._scores(rows, [[self.cfg["yes"], self.cfg["no"]]] * len(rows))]

    def is_on_topic(self, paper, topic, definition=None):
        return self.is_on_topic_many([paper], topic, definition)[0]


class PaperClassifier(ClassifierTasks):
    def __init__(self, model="maximerivest/paper-classifier-2B-v2", *, revision=None, device=None,
                 dtype=None, batch_size=4, paper_tokens=700):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        dtype = dtype or (torch.bfloat16 if self.device.startswith("cuda") else torch.float32)
        if batch_size < 1:
            raise ValueError("batch_size must be positive")
        self.batch_size = batch_size
        self.cfg = read_config(model, revision)
        self.tok = AutoTokenizer.from_pretrained(model, revision=revision)
        self.tok.padding_side = "left"
        self.inputs = PaperInputs(self.tok, self.cfg, paper_tokens=paper_tokens)
        self.model = AutoModelForCausalLM.from_pretrained(model, revision=revision, dtype=dtype,
                                                          attn_implementation="sdpa").to(self.device).eval()

    def _scores(self, input_ids, answer_ids, shared_prefix=False):
        import torch
        out = []
        with torch.inference_mode():
            for i in range(0, len(input_ids), self.batch_size):
                enc = self.tok.pad({"input_ids": input_ids[i:i+self.batch_size]}, return_tensors="pt").to(self.device)
                h = self.model.model(**enc, use_cache=False).last_hidden_state[:, -1, :].float()
                W = self.model.get_output_embeddings().weight
                for row, ids in zip(h, answer_ids[i:i+self.batch_size]):
                    out.append(torch.softmax(row @ W[ids].float().T, -1).cpu().tolist())
        return out


if __name__ == "__main__":
    import sys
    model = sys.argv[1] if len(sys.argv)>1 else "maximerivest/paper-classifier-2B-v2"
    clf = PaperClassifier(model)
    paper = paper_text("Habitat fragmentation and tropical forest birds",
                       "We surveyed bird communities across 40 forest fragments. Species richness declined with "
                       "fragment area and increased with connectivity.")
    print(clf.classify(paper, ["Ecology", "Ornithology", "Particle physics", "Economics"], scheme="Fields of science"))
    print("P(topic is central):", clf.is_on_topic(paper, "habitat fragmentation"))