"""High-level API: Decider.from_pretrained(...).decide(request or [requests])."""
import os

import torch

from . import _core as C


class Decider:
    """A loaded Decider 2B model.

    decide() takes one Jev/SystemOne request (returns one response) or a list (returns a list in the
    same order). Questions that do not fit next to the full state in one pass are split across
    passes, each with the whole state; nothing is truncated. Passes from all requests are batched
    on the device under a padded-token budget.
    """

    def __init__(self, model, name="decider-2b", batch_tokens=32768):
        self.model = model
        self.name = name
        self.batch_tokens = int(batch_tokens)

    @classmethod
    def from_pretrained(cls, model="Jeetu456/decider-2b", device=None, dtype=None, quantization=None,
                        abstain_threshold=None, revision=None, token=None, batch_tokens=32768, name="decider-2b"):
        """Load from a Hugging Face repo id or a local folder.

        device: "cuda", "mps" or "cpu" (default: cuda if available, else cpu).
        quantization: None, "8bit" or "4bit" (bitsandbytes, CUDA only).
        abstain_threshold: override the tuned "I don't know" threshold (0 turns it off).
        """
        path = model
        if not os.path.isdir(path):
            from huggingface_hub import snapshot_download
            path = snapshot_download(model, revision=revision, token=token)
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        loaded = C.load_release_model(path, device=device, dtype=dtype, quantization=quantization)
        if abstain_threshold is not None:
            C.ABSTAIN_THRESHOLD = float(abstain_threshold)
        return cls(loaded, name=name, batch_tokens=batch_tokens)

    @property
    def temperature(self):
        return self.model.temperature

    @property
    def device(self):
        return self.model.head_device

    def _plan(self, request):
        C.validate_request(request)
        questions = request["questions"]
        state = C.request_state_text(request["state"], questions)
        rendered = C.items(questions)
        passes, state_tokens, sizes = C.plan(self.model.tokenizer, state, rendered, self.model.max_length)
        return questions, state, rendered, passes, state_tokens, sizes

    @torch.inference_mode()
    def _logits(self, requests):
        """Raw logits per request: a list (one per question, in request order) of option-logit lists."""
        plans = [self._plan(r) for r in requests]
        units = []  # (request index, item indices, estimated length)
        for r, (_, _, _, passes, state_tokens, sizes) in enumerate(plans):
            for indices in passes:
                units.append((r, indices, state_tokens + sum(sizes[i] for i in indices) + 32))
        out = [[None] * len(p[2]) for p in plans]
        order = sorted(range(len(units)), key=lambda u: units[u][2])
        start = 0
        while start < len(order):
            end, width = start, 0
            while end < len(order):
                width = max(width, units[order[end]][2])
                if end > start and width * (end - start + 1) > self.batch_tokens:
                    break
                end += 1
            chosen = [units[u] for u in order[start:end]]
            batch = self.model.encode([{"state": plans[r][1], "items": [plans[r][2][i] for i in idx]}
                                       for r, idx, _ in chosen])
            logits = self.model(**{k: v.to(self.device) for k, v in batch.items()}).float().cpu()
            row = 0
            for r, idx, _ in chosen:
                for i in idx:
                    out[r][i] = logits[row, : len(plans[r][2][i]["labels"])]
                    row += 1
            start = end
        return plans, out

    def logits(self, requests):
        """Raw option logits before temperature: {question id: [logit per option, in request order]}."""
        single = isinstance(requests, dict)
        plans, out = self._logits([requests] if single else requests)
        result = [{item["id"]: values.tolist() for item, values in zip(plan[2], rows)} for plan, rows in zip(plans, out)]
        return result[0] if single else result

    def decide(self, requests, temperature=None):
        """Jev/SystemOne response(s): {"model", "answers", "usage"} per request."""
        single = isinstance(requests, dict)
        batch = [requests] if single else list(requests)
        plans, out = self._logits(batch)
        t = float(temperature or self.model.temperature)
        responses = []
        for request, (questions, _, rendered, passes, state_tokens, sizes), rows in zip(batch, plans, out):
            probabilities = [(values / t).softmax(-1).tolist() for values in rows]
            responses.append(C.response(request.get("model") or self.name, questions, probabilities,
                                        len(passes) * state_tokens + sum(sizes)))
        return responses[0] if single else responses