"""Stufe 3: Entscheidungen direkt auf Embeddings (Nearest Centroid mit Konfidenz)."""

from __future__ import annotations

import json
import math
from collections import defaultdict
from pathlib import Path

from .embeddings import Vector, cosine


class EmbeddingClassifier:
    def __init__(self, centroids: dict[str, Vector], counts: dict[str, int] | None = None):
        if not centroids:
            raise ValueError("Classifier braucht mindestens eine Klasse")
        self.centroids = centroids
        self.counts = counts or {}

    @classmethod
    def train(cls, examples: list[tuple[Vector, str]]) -> EmbeddingClassifier:
        sums: dict[str, Vector] = {}
        counts: dict[str, int] = defaultdict(int)
        for vec, label in examples:
            if label not in sums:
                sums[label] = [0.0] * len(vec)
            sums[label] = [a + b for a, b in zip(sums[label], vec, strict=True)]
            counts[label] += 1
        centroids = {}
        for label, total in sums.items():
            norm = math.sqrt(sum(v * v for v in total)) or 1.0
            centroids[label] = [v / norm for v in total]
        return cls(centroids, dict(counts))

    def predict(self, vec: Vector) -> tuple[str, float]:
        """Liefert (Label, Konfidenz). Konfidenz = Softmax-Abstand der besten zwei Klassen (0..1)."""
        scores = sorted(((cosine(vec, c), label) for label, c in self.centroids.items()), reverse=True)
        best_score, best_label = scores[0]
        if len(scores) == 1:
            return best_label, max(0.0, best_score)
        exps = [math.exp(s * 20) for s, _ in scores]  # Temperatur 0.05
        return best_label, exps[0] / sum(exps)

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps({"centroids": self.centroids, "counts": self.counts}))

    @classmethod
    def load(cls, path: str | Path) -> EmbeddingClassifier:
        data = json.loads(Path(path).read_text())
        return cls(data["centroids"], data.get("counts"))