"""分级头（T-G4）：LinearHead 纯 Python 推理 + LoraRemote 预留（T-G8）。

D-G5：serving 路径零新依赖——线性头推理是纯 Python 点积 + softmax（768 维
≈0.1ms）。工件为 JSON（weights 3×dim / bias 3 / dim / version），由
scripts/train_tier_head.py（numpy 离线训练）产出。
缺失/损坏 -> ArtifactMissing（D-G4 降级信号）。
"""
from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Dict, List, Optional

from gateway.sense.errors import ArtifactMissing

TIERS = ("t1", "t2", "t3")


class LinearHead:
    """线性有序三分类头（softmax；类别序 t1<t2<t3 与升级阶梯一致）。"""

    def __init__(self, weights: List[List[float]], bias: List[float],
                 version: str = "dev", labels: Optional[List[str]] = None):
        self.weights = weights
        self.bias = bias
        self.version = version
        self.labels = labels or ["t1", "t2", "t3"]
        if len(self.weights) != len(self.bias):
            raise ArtifactMissing("线性头权重与偏置维度不一致")

    @classmethod
    def load(cls, path: str | Path, version: str = "") -> "LinearHead":
        p = Path(path)
        if not p.exists():
            raise ArtifactMissing(f"线性头工件不存在: {p}")
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
            weights = data["weights"]
            bias = data["bias"]
            if not weights or not bias:
                raise ValueError("空权重")
            return cls(weights=weights, bias=bias,
                       version=str(data.get("version") or version or p.parent.name),
                       labels=data.get("labels"))
        except ArtifactMissing:
            raise
        except Exception as e:  # noqa: BLE001
            raise ArtifactMissing(f"线性头工件损坏: {type(e).__name__}: {e}") from e

    def predict(self, vec: List[float]) -> Dict[str, float]:
        """点积 + softmax -> {t1, t2, t3} 概率（和为 1）。"""
        logits = []
        for w, b in zip(self.weights, self.bias):
            n = min(len(w), len(vec))
            logits.append(sum(wi * vi for wi, vi in zip(w[:n], vec[:n])) + b)
        m = max(logits)
        exps = [math.exp(z - m) for z in logits]
        total = sum(exps)
        return {label: e / total for label, e in zip(self.labels, exps)}


class LoraRemote:
    """LoRA 远程分类（vLLM /v1/classify）——T-G8 可选实验，本版预留。"""

    def __init__(self, base_url: str, model: str, api_key: Optional[str] = None):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key

    def predict(self, vec: List[float]) -> Dict[str, float]:
        raise NotImplementedError("T-G8 可选实验（LoRA/vLLM 分类服务）")