#!/usr/bin/env python3
"""Yes/no forecast. One forward pass, then the fitted map. Returns P(YES) only."""
from __future__ import annotations

import math
import string

SYSTEM_PROMPT = "Apply the criterion to the evidence. Choose exactly one listed option. Answer with its letter only."
THINK_OFF = "<think>\n\n</think>\n\n"
TEMPERATURE = 1.421
# Fitted on the primary-election evaluation.
A = 1.099406044975045
C = 1.3783573144958756


def _logit(p: float) -> float:
    p = min(max(p, 1e-6), 1 - 1e-6)
    return math.log(p / (1 - p))


def _sigmoid(x: float) -> float:
    return 1.0 / (1.0 + math.exp(-x))


def _letter_ids(tokenizer) -> list[int]:
    probe = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Evidence:\n(none)\n\nQuestion: q\nOptions:\nA) a\nB) b"},
    ]
    text = tokenizer.apply_chat_template(probe, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    if not text.endswith(THINK_OFF):
        raise ValueError("chat template is missing the expected end marker")
    base = tokenizer.encode(text, add_special_tokens=False)
    ids = []
    for letter in string.ascii_uppercase:
        piece = tokenizer.encode(letter, add_special_tokens=False)
        if len(piece) != 1 or tokenizer.encode(text + letter, add_special_tokens=False) != base + piece:
            raise ValueError(f"{letter!r} is not a single token for this tokenizer")
        ids.append(piece[0])
    return ids


class Forecaster:
    """Load once. forecast(...) returns a float in [0, 1]."""

    def __init__(self, model_path: str = "YarrowLab/yarrow-forecast"):
        import torch
        from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration

        self._torch = torch
        self.tok = AutoTokenizer.from_pretrained(model_path)
        self.model = Qwen3_5ForConditionalGeneration.from_pretrained(
            model_path, dtype="auto", device_map="auto",
        ).eval()
        self._yes, self._no = _letter_ids(self.tok)[:2]

    def forecast(self, title: str, rules: str, evidence: str = "", as_of: str | None = None) -> float:
        return self.forecast_batch([{
            "title": title, "rules": rules, "evidence": evidence, "as_of": as_of,
        }])[0]

    def forecast_batch(self, questions: list[dict]) -> list[float]:
        return [self._one(q) for q in questions]

    def _one(self, q: dict) -> float:
        as_of = q.get("as_of") or "the question date"
        state = f"Forecast as of {as_of}."
        if q.get("evidence"):
            state += "\n\nEVIDENCE AVAILABLE AT FORECAST TIME:\n" + q["evidence"]
        question = (
            f"{q['title']}\n\nResolution rules (verbatim): {q['rules']}\n\n"
            "Will this resolve YES, based only on the evidence above?"
        )
        state = state if state.strip() else "(none)"
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "Evidence:\n" + state + "\n\nQuestion: " + question.strip() + "\nOptions:\nA) yes\nB) no"},
        ]
        text = self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
        if not text.endswith(THINK_OFF):
            raise ValueError("chat template is missing the expected end marker")
        ids = self.tok.encode(text, add_special_tokens=False)
        input_ids = self._torch.tensor([ids], device=self.model.device)
        with self._torch.no_grad():
            logits = self.model(input_ids=input_ids).logits[0, -1].float()
        pair = self._torch.stack([logits[self._yes], logits[self._no]])
        yes = self._torch.softmax(pair / TEMPERATURE, dim=0)[0].item()
        return _sigmoid(A * _logit(yes) + C)