"""Decision Index engine for FINAL-Bench/Darwin-27B-ZTC.

Use with the Decision Index kit (github.com/apolinario/decision-index):

    PYTHONPATH=<snapshot of this repo> python -m decision_index pipeline \
        --engine ztc_engine:DarwinZTCEngine --option model=FINAL-Bench/Darwin-27B-ZTC --out runs/darwin-27b-ztc

One request is one row: every question of the row is answered with one forward pass of the
backbone (no generated tokens), and the questions of a row are batched together.
A question whose prompt is longer than max_length tokens is declared unsupported (never truncated).
"""

import math
import sys
import time
from pathlib import Path

import torch

from decision_index.engines.base import Engine, Unsupported


class DarwinZTCEngine(Engine):
    name = "darwin-27b-ztc"
    latency = "In-process request wall time (prompt construction, tokenisation and forward passes; excludes model loading)."

    def __init__(self, model="FINAL-Bench/Darwin-27B-ZTC", revision=None, device="cuda", max_length=32768,
                 batch=8, token_budget=16000, **options):
        super().__init__(**options)
        path = Path(model)
        if not path.exists():
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(model, revision=revision))
        if str(path) not in sys.path:
            sys.path.insert(0, str(path))
        from autojev.model import DecisionModel, answer
        self._answer = answer
        self.model = DecisionModel(checkpoint=path, device=device).eval()
        self.max_length, self.batch, self.token_budget = int(max_length), int(batch), int(token_budget)
        self.provenance = {"model": model, "revision": revision, "path": str(path), "temperature": self.model.temperature,
                           "max_length": self.max_length, "batch": self.batch, "token_budget": self.token_budget}

    def runtime(self):
        return {"torch": torch.__version__, "cuda": torch.version.cuda,
                "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"}

    def synchronize(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    @staticmethod
    def _est(row):
        import json
        return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192

    @torch.inference_mode()
    def _probs(self, rows):
        b = self.model.prepare(rows, max_length=self.max_length)
        d = (self.model(b) / self.model.temperature).softmax(-1).cpu().tolist()
        return [v[:c] for v, c in zip(d, b.counts)]

    def _one(self, row):
        try:
            return self._probs([row])[0]
        except ValueError as e:
            raise Unsupported(str(e))

    def __call__(self, state, questions):
        keys = list(questions)
        rows = [{"state": state, "question": questions[k]} for k in keys]
        order = sorted(range(len(rows)), key=lambda i: self._est(rows[i]))
        dists = {}
        i = 0
        while i < len(order):
            group = [order[i]]; tok = self._est(rows[order[i]])
            while i + len(group) < len(order) and len(group) < self.batch and tok + self._est(rows[order[i + len(group)]]) <= self.token_budget:
                tok += self._est(rows[order[i + len(group)]]); group.append(order[i + len(group)])
            i += len(group)
            try:
                for j, p in zip(group, self._probs([rows[j] for j in group])):
                    dists[j] = p
            except (ValueError, torch.OutOfMemoryError):
                torch.cuda.empty_cache()
                for j in group:
                    dists[j] = self._one(rows[j])
        answers = {}
        for j, k in enumerate(keys):
            p = dists[j]
            if not all(math.isfinite(x) and x >= 0 for x in p) or sum(p) <= 0:
                p = [1.0] * len(p)
            answers[k] = self._answer(questions[k], p)
        return {"model": self.name, "answers": answers}, None