"""Two-checkpoint ensemble for the Typed Decisions specialist."""

from __future__ import annotations

import json
import os

import numpy as np
import torch
import torch.nn.functional as F

from .schema import KIND_TO_ID, parse_questions
from .task_specialist import TaskSpecialistEngine


class EnsembleEngine:
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, device: str | None = None, token: str | None = None, **kwargs) -> "EnsembleEngine":
        return cls(pretrained_model_name_or_path, device=device, token=token)

    def __init__(self, config_path: str, device: str | None = None, token: str | None = None):
        path = config_path
        if not os.path.exists(path):
            from huggingface_hub import snapshot_download
            path = snapshot_download(path, token=token)
        if os.path.isdir(path):
            candidate = os.path.join(path, "ensemble_config.json")
            path = candidate if os.path.exists(candidate) else path
        base = os.path.dirname(os.path.abspath(path))
        cfg = json.load(open(path, encoding="utf-8"))
        self.config_path = path
        self.device = device

        def resolve(p: str) -> str:
            if os.path.isabs(p) or os.path.exists(p):
                return p
            return os.path.join(base, p)

        self.engines = [TaskSpecialistEngine(resolve(p), device=device) for p in cfg["models"]]
        self.weights = cfg.get("weights", [1.0 / len(self.engines)] * len(self.engines))
        self.temperatures = cfg.get("temperature", {})
        self.model_name = cfg.get("name", "pluto-typed-ensemble")
        self.max_len = cfg.get("max_len", 512)

    @property
    def info(self) -> dict:
        return {
            "name": self.model_name,
            "models": [e.info["name"] for e in self.engines],
            "weights": self.weights,
            "temperatures": self.temperatures,
            "device": str(self.engines[0].device) if self.engines else str(self.device),
            "params": sum(e.info["params"] for e in self.engines),
        }

    def _blend(self, per_model: list[list[dict]], items) -> list[dict]:
        out = []
        for i, (_, q) in enumerate(items):
            ps = []
            for model_results in per_model:
                ps.append(model_results[i]["probs"][: q.n_options])
            p = np.zeros(q.n_options, dtype=np.float64)
            for w, pi in zip(self.weights, ps):
                p += w * pi
            p = p / max(p.sum(), 1e-9)
            t = self.temperatures.get(q.kind, 1.0)
            if t != 1.0:
                p = np.power(np.clip(p, 1e-12, 1.0), 1.0 / t)
                p = p / p.sum()
            top = int(np.argmax(p))
            if q.kind == "choice":
                value = q.keys[top]
            elif q.kind == "noul":
                value = float(p[1])
            else:
                value = float(np.dot(np.arange(q.n_options), p))
            out.append(
                {
                    "index": i,
                    "top": top,
                    "value": value,
                    "probs": p.astype(np.float32),
                    "confidence": float(p[top]),
                    "accept_probability": float(p[top]),
                }
            )
        return out

    def predict_many(self, items, batch_size: int = 64, max_len: int | None = None) -> list[dict]:
        items = list(items)
        per_model = [e.predict_many(items, batch_size=batch_size, max_len=max_len) for e in self.engines]
        return self._blend(per_model, items)

    def predict(self, state, questions, max_len: int | None = None, alpha: float | None = None) -> dict:
        qs = parse_questions(questions) if not isinstance(questions, list) else questions
        items = [(state, q) for q in qs]
        results = self.predict_many(items, max_len=max_len)
        answers = {}
        escalate = False
        for i, q in enumerate(qs):
            p = results[i]["probs"][: q.n_options]
            top = int(np.argmax(p))
            conf = float(p[top])
            item = {
                "type": q.kind,
                "confidence": conf,
                "accept_probability": conf,
                "escalate": conf < 0.5,
                "probabilities": {k: float(v) for k, v in zip(q.keys, p)},
            }
            if q.kind == "choice":
                item["choice"] = q.keys[top]
            elif q.kind == "noul":
                item["noul"] = float(p[1])
            else:
                item["score"] = float(np.dot(np.arange(q.n_options), p))
                item["score_normalized"] = item["score"] / max(1, q.n_options - 1)
            answers[q.name] = item
            escalate = escalate or item["escalate"]
        return {"answers": answers, "escalate": escalate, "routing": {"model": self.model_name}}