"""Decision Index engine for prism-system-one: one predict() call per request, refusing (never truncating)."""
import math
import os

import torch
from decision_index.engines import Engine, Unsupported
from modeling_prism import PrismDecisionModel


def _finite(response):
    for a in response["answers"].values():
        values = [a["noul"]] if a["type"] == "noul" else list(a["probabilities"].values())
        if not all(math.isfinite(v) for v in values):
            return False
    return True


class PrismEngine(Engine):
    name = "prism-system-one"
    latency = ("One forward pass per request on one GPU (fp16 on T4, bf16 on Ampere and newer), "
               "including tokenization.")

    def __init__(self, path, device=None, **options):
        super().__init__(**options)
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        cap = torch.cuda.get_device_capability()[0] if self.device == "cuda" else 0
        self.amp = self.device == "cuda"
        self.dtype = torch.bfloat16 if cap >= 8 else torch.float16
        if self.amp and cap < 8:
            torch.backends.cudnn.enabled = False          # avoids the depthwise-conv crash seen on T4
        self.model = PrismDecisionModel.from_pretrained(path, device=self.device).eval()
        self.provenance = {"repo": "paulprt/prism-system-one",
                           "revision": os.path.basename(os.path.normpath(path)),
                           "device": self.device,
                           "dtype": str(self.dtype) if self.amp else "float32",
                           "max_length": self.model.decision_model.max_length}

    def _predict(self, request, amp):
        with torch.autocast("cuda", dtype=self.dtype, enabled=amp):
            return self.model.predict([request], strict=True)[0]

    def __call__(self, state, questions):
        request = {"state": state if state is not None else "", "questions": questions}
        try:
            response = self._predict(request, self.amp)
            if self.amp and not _finite(response):        # fp16 overflow: redo this request in float32
                response = self._predict(request, False)
        except ValueError as e:
            raise Unsupported(str(e))                     # too long or outside the declared limits
        return response, None