"""Riconoscitore per similarita' semantica con embedding (sentence-transformers).

Tecnica diversa: non classifica ne' genera. Calcola gli embedding di
*frammenti candidati* del testo (n-grammi di parole) e li confronta, via
similarita' del coseno, con gli embedding di *concetti seme* del dominio
appalti. Un frammento sufficientemente vicino a un concetto viene emesso come
span del tipo associato.

Utile per cogliere varianti lessicali non previste da regole/etichette fisse
(es. "soggetto aggiudicatore" ~ stazione appaltante) senza un classificatore
addestrato.

- Modello iniettato come `encoder(list[str]) -> matrice (n, d)` normalizzabile
  (Dependency Inversion).
- I concetti seme sono iniettati (config), non cablati nella classe.
"""

from __future__ import annotations

import re
from collections.abc import Sequence
from typing import Any

from config.detection.category_scores import category_score
from core.contracts import LayerPriority, Span, build_span, cap_score
from layers.embedding.config import EMBEDDING_LOADER_POLICY, LAYER_SCORING, EmbeddingLoaderPolicy

_EMBEDDING_SCORING = LAYER_SCORING
_LAYER_KEY = "embedding"  # chiave per la tabella di cap per-categoria

# Genera i candidati: finestre scorrevoli di 1..4 parole (riancorate al testo).
_WORD = re.compile(r"\S+")


def _candidate_spans(text: str, max_words: int = 4) -> list[tuple[int, int, str]]:
    tokens = [(m.start(), m.end()) for m in _WORD.finditer(text)]
    out: list[tuple[int, int, str]] = []
    n = len(tokens)
    for i in range(n):
        for w in range(1, max_words + 1):
            j = i + w - 1
            if j >= n:
                break
            start, end = tokens[i][0], tokens[j][1]
            frag = text[start:end]
            if 3 <= len(frag) <= 80:
                out.append((start, end, frag))
    return out


def _l2_normalize(mat):
    """Normalizza L2 le righe di una matrice (n, d); righe nulle restano nulle.

    Dopo la normalizzazione il coseno tra due righe e' un semplice prodotto
    scalare: cosi' l'intera matrice di similarita' candidati×seed diventa un
    unico matmul, indipendentemente dal fatto che l'encoder iniettato normalizzi
    gia' o meno (stessa semantica del vecchio confronto coppia-a-coppia).
    """
    import numpy as np

    norm = np.linalg.norm(mat, axis=1, keepdims=True)
    norm[norm == 0] = 1.0
    return mat / norm


class EmbeddingSimilarityRecognizer:
    """Riconoscitore per vicinanza semantica a concetti seme (embedding)."""

    layer = LayerPriority.EMBEDDING
    name = "embedding_similarity"

    def __init__(
        self,
        encoder,
        seeds_by_type: dict[str, list[str]],
        threshold: float = _EMBEDDING_SCORING.threshold,
        max_words: int = 4,
        score_cap: float = _EMBEDDING_SCORING.score_cap,
    ) -> None:
        # `encoder`: callable list[str] -> matrice embedding (n, d).
        self._encoder = encoder
        self._threshold = threshold
        self._max_words = max_words
        self._score_cap = score_cap
        self._seed_types: list[str] = []
        self._seed_texts: list[str] = []
        for etype, seeds in seeds_by_type.items():
            for s in seeds:
                self._seed_types.append(etype)
                self._seed_texts.append(s)
        self._seed_vecs: Any = None  # calcolati lazy al primo uso

    def _ensure_seeds(self) -> None:
        if self._seed_vecs is None:
            self._seed_vecs = self._encoder(self._seed_texts)

    def analyze(self, text: str) -> Sequence[Span]:
        import numpy as np

        candidates = _candidate_spans(text, self._max_words)
        if not candidates:
            return []
        # Dedup dei frammenti prima dell'encoding: su testo ripetitivo (es. molti
        # token "appalto"/"comune") la stessa stringa verrebbe encodata piu' volte.
        # Encoda gli unici, poi rimappa ai candidati per posizione.
        unique_frags = list(dict.fromkeys(c[2] for c in candidates))
        try:
            self._ensure_seeds()
            uniq_vecs = self._encoder(unique_frags)
        except Exception:  # encoder non disponibile -> nessuno span
            return []
        cand = np.asarray(uniq_vecs, dtype=float)
        seeds = np.asarray(self._seed_vecs, dtype=float)
        if cand.ndim != 2 or seeds.ndim != 2 or seeds.shape[0] == 0:
            return []
        # Similarita' coseno = matmul su righe normalizzate: (n_unici × n_seed).
        sims = _l2_normalize(cand) @ _l2_normalize(seeds).T
        best_idx = sims.argmax(axis=1)
        best_sim = sims[np.arange(sims.shape[0]), best_idx]
        frag_pos = {frag: i for i, frag in enumerate(unique_frags)}

        spans: list[Span] = []
        for start, end, frag in candidates:
            pos = frag_pos[frag]
            score = float(best_sim[pos])
            if score >= self._threshold:
                best_type = self._seed_types[int(best_idx[pos])]
                # Tetto PER CATEGORIA (config/detection/category_scores): la
                # similarità è solo un voto di supporto, calibrato per concetto.
                cap = category_score(_LAYER_KEY, best_type, fallback=self._score_cap)
                spans.append(
                    build_span(
                        start=start,
                        end=end,
                        text=frag,
                        entity_type=best_type,
                        score=cap_score(score, cap),
                        layer=self.layer,
                        source=f"{self.name}:{best_type}",
                        slice_text=False,
                    )
                )
        return spans


def load_sentence_encoder(
    model_name: str,
    device: str,
    policy: EmbeddingLoaderPolicy = EMBEDDING_LOADER_POLICY,
):
    """Carica un modello sentence-transformers (import pesante lazy).

    Ritorna un callable `encode(list[str]) -> ndarray (n, d)`.
    `normalize_embeddings` arriva dalla policy centralizzata (config/models): la
    normalizzazione resta coerente sostituendo encoder della stessa famiglia.
    """
    from sentence_transformers import SentenceTransformer  # type: ignore

    model = SentenceTransformer(model_name, device=device)

    def encode(texts: list[str]):
        return model.encode(
            texts,
            normalize_embeddings=policy.normalize_embeddings,
            show_progress_bar=False,
        )

    return encode