"""HevForDecisions: frozen DiffusionGemma backbone + Hev head, exposing Jev-shaped decide()."""
from __future__ import annotations

import math
import random
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import PreTrainedModel

from .cache import collate
from .configuration_hev import HevConfig
from .features import NormCapture, _slot_token, encode_example, extra_layer_names
from .head import HevEnsemble, HevHead, HevHeadConfig
from .prompting import PROMPT_FORMAT_VERSION, alias_token_ids
from .schema import Example, Question


# Hub download options that from_pretrained forwards to snapshot_download and HevConfig.from_pretrained.
HUB_KWARGS = ("revision", "token", "cache_dir", "local_files_only", "force_download")


def _is_local_dir(name: str) -> bool:
    try:
        return Path(name).is_dir()
    except (OSError, ValueError):
        return False


def _vocab_size(backbone) -> int:
    """Backbone vocabulary size, read as hev.features.main reads it for random slot draws (default 262144)."""
    return getattr(getattr(backbone.config, "text_config", None), "vocab_size", 262144)


def _param_device(module) -> torch.device:
    """Device of the module's first parameter, or CPU when it has none."""
    params = module.parameters() if hasattr(module, "parameters") else iter(())
    first = next(params, None)
    return first.device if first is not None else torch.device("cpu")


class HevForDecisions(PreTrainedModel):
    config_class = HevConfig

    def __init__(self, config: HevConfig, backbone=None, tokenizer=None):
        super().__init__(config)
        members = []
        temps = list(getattr(config, "member_temperatures", None) or [config.temperatures])
        for i in range(int(getattr(config, "n_heads", 1))):
            head = HevHead(HevHeadConfig.from_dict(config.head_config))
            with torch.no_grad():
                head.log_temp.copy_(torch.log(torch.tensor(temps[i] if i < len(temps) else config.temperatures,
                                                           dtype=torch.float32)))
            members.append(head)
        self.head = HevEnsemble(members)
        self.backbone = backbone
        self.tokenizer = tokenizer
        self._cap = None
        self._alias_ids = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, backbone_cls=None, tokenizer_cls=None,
                        dtype=torch.bfloat16, device_map="auto", load_in_4bit: bool = False, config=None, **kwargs):
        """Load the Hev config, head, frozen backbone and tokenizer from a local folder or a Hub repo.

        `revision`, `token`, `cache_dir`, `local_files_only` and `force_download` are forwarded to the Hub
        download and to HevConfig.from_pretrained. A HevConfig passed as `config` (AutoModel does this) is used as is.
        """
        hub = {k: kwargs[k] for k in HUB_KWARGS if kwargs.get(k) is not None}
        name = str(pretrained_model_name_or_path)
        path = Path(name)
        from_hub = not path.exists()
        if from_hub:
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(name, **hub))
        if not isinstance(config, HevConfig):
            config = HevConfig.from_pretrained(str(path), **hub)
        if from_hub:
            config._name_or_path = name  # the repo id, not the local snapshot folder
        if backbone_cls is None or tokenizer_cls is None:
            from transformers import AutoConfig, AutoTokenizer, DiffusionGemmaForBlockDiffusion
            backbone_cls = backbone_cls or DiffusionGemmaForBlockDiffusion
            tokenizer_cls = tokenizer_cls or AutoTokenizer
            bb_cfg = AutoConfig.for_model(**config.backbone_config)
        else:
            bb_cfg = None
        bb_kwargs = {"config": bb_cfg, "dtype": dtype, "device_map": device_map}
        if load_in_4bit:
            from transformers import BitsAndBytesConfig
            bb_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=dtype)
        backbone = backbone_cls.from_pretrained(str(path), **bb_kwargs).eval()
        tokenizer = tokenizer_cls.from_pretrained(str(path))
        model = cls(config, backbone=backbone, tokenizer=tokenizer)
        if len(model.head.members) == 1 and (path / "hev_head.safetensors").exists():
            model.head.members[0].load_state_dict(load_file(str(path / "hev_head.safetensors")))
        else:
            for i, member in enumerate(model.head.members):
                member.load_state_dict(load_file(str(path / f"hev_head_{i}.safetensors")))
        model.head.to(_param_device(backbone))
        return model.eval()

    @staticmethod
    def confidence(probs) -> float:
        k = len(probs)
        if k < 2:
            return 1.0
        h = -sum(p * math.log(p) for p in probs if p > 0)
        return max(0.0, min(1.0, 1.0 - h / math.log(k)))

    def _extra_layers(self) -> dict:
        extra = getattr(self.config, "extra_layers", None) or {}
        return {"encoder": list(extra.get("encoder", [])), "decoder": list(extra.get("decoder", []))}

    def _ensure_capture(self):
        if self._cap is None:
            self._cap = NormCapture(self.backbone, self.config.encoder_norm, self.config.decoder_norm,
                                    extra=extra_layer_names({"extra_layers": self._extra_layers()}))
            self._alias_ids = alias_token_ids(self.tokenizer, len(self.config.aliases))

    @torch.no_grad()
    def decide(self, state: str, questions: dict) -> dict:
        """Answer Jev-shaped questions about `state` with the backbone features and the trained head.

        The answer slots are seeded exactly as hev.features.extract seeded them in training, by
        `config.seed_mode`: the pad token (0) for "pad", the first alias token for "alias", and an average
        over RANDOM_DRAWS forward passes with random slot tokens (fixed rng seed 1234) for "random".
        `config.slot_token_id` is not read here; seed_mode governs.
        """
        version = getattr(self.config, "prompt_format_version", 1)
        if version != PROMPT_FORMAT_VERSION:
            raise RuntimeError(
                f"this head was trained with prompt format {version}, but the installed hev code builds prompt format "
                f"{PROMPT_FORMAT_VERSION}; load the code revision the head was published with, or retrain the head")
        self._ensure_capture()
        qs = [Question.from_jev(qid, spec) for qid, spec in questions.items()]
        ex = Example(id="request", source="api", state=state, questions=qs, gold={})
        info = {"canvas_length": self.config.canvas_length, "encoder_norm": self.config.encoder_norm,
                "decoder_norm": self.config.decoder_norm, "extra_layers": self._extra_layers()}
        slot_tokens = _slot_token(self.config.seed_mode, self._alias_ids, random.Random(1234),
                                  _vocab_size(self.backbone))
        feats, rec = encode_example(self.backbone, self.tokenizer, ex, self._cap, info,
                                    slot_tokens=slot_tokens, alias_ids=self._alias_ids,
                                    max_tokens=self.config.max_prompt_tokens)
        batch = collate([(rec, feats)]).to(next(self.head.parameters()).device)
        logits = self.head(batch, use_temperature=True)
        answers = {}
        for q, lg in zip(qs, logits):
            probs = torch.softmax(lg.float(), 0).tolist()
            if q.type == "noul":
                answers[q.qid] = {"type": "noul", "noul": probs[0]}
            elif q.type == "choice":
                best = max(range(q.k), key=lambda i: probs[i])
                answers[q.qid] = {"type": "choice", "choice": q.options[best],
                                  "probabilities": dict(zip(q.options, probs)), "confidence": self.confidence(probs)}
            else:
                answers[q.qid] = {"type": "score", "score": sum(i * p for i, p in enumerate(probs)),
                                  "probabilities": {str(i): p for i, p in enumerate(probs)},
                                  "confidence": self.confidence(probs), "legend": list(q.descriptions)}
        name = getattr(self.config, "_name_or_path", "") or ""
        # A local folder name would leak a filesystem path into every response, so only a Hub repo id is reported.
        return {"model": name if name and not _is_local_dir(name) else "hev",
                "answers": answers, "usage": {"input_tokens": rec.n_prompt, "output_tokens": 0}}