"""Calibrare politica k-NN pentru sugestia embeddings (top-1 vs vot ponderat top-k).

Metodologie: leave-one-out pe corpusul SILVER (`mapping_suggestions`) folosind
vectorii persistati in `embedding_cache` (zero re-embedding). Pentru fiecare
exemplu: se exclude din corpus, se prezice din vecinii ramasi, se compara cu
eticheta lui SILVER (cod sau NUL).

ATENTIE la interpretare: SILVER e etichetat de LLM (~95% calitate la scara),
NU ground-truth uman — valorile absolute sunt optimiste; comparatia RELATIVA
intre politici si calibrarea pragului raman valide (aceeasi tinta pentru toate).

Politici evaluate:
  top1       — eticheta primului vecin (comportamentul curent din enrich_suggestions)
  vote<k>    — vot ponderat cu similaritatea pe top-k vecini peste prag;
               NUL e eticheta proprie (castiga -> supresie)

Metrice (per politica x prag):
  coverage    — % din exemple pentru care se emite o predictie (cod sau NUL)
  precizie    — % predictii corecte din cele emise
  cod-gresit  — % din TOATE exemplele unde s-a emis un COD diferit de gold
                (critic: cod gresit preselectat = risc FINALIZATA eronata)

Rulare:
    python3 tools/mapare-llm/knn_calibrate.py --db data/autopass.db
    python3 tools/mapare-llm/knn_calibrate.py --db data/autopass.db --sample 5000 --seed 42
"""
from __future__ import annotations

import argparse
import os
import random
import sqlite3
import sys

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.abspath(os.path.join(_HERE, "..", ".."))
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)

from app import embedding_cache as cache  # noqa: E402
from app.embeddings import FASTEMBED_MODEL  # noqa: E402

NUL = "NUL"


def load_corpus(db_path: str):
    """Incarca (etichete, matrice vectori) aliniate pozitional din SQLite."""
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT denumire_normalizata, cod_prestatie, is_nul FROM mapping_suggestions "
        "WHERE denumire_normalizata IS NOT NULL"
    ).fetchall()
    vecmap = {
        h: np.frombuffer(b, dtype=np.float32)
        for h, b in conn.execute(
            "SELECT text_hash, vector FROM embedding_cache WHERE model=?", (FASTEMBED_MODEL,)
        )
    }
    labels: list[str] = []
    vecs: list[np.ndarray] = []
    lipsa = 0
    for r in rows:
        v = vecmap.get(cache.text_hash(r["denumire_normalizata"]))
        if v is None:
            lipsa += 1
            continue
        labels.append(NUL if r["is_nul"] else str(r["cod_prestatie"]))
        vecs.append(v)
    if lipsa:
        print(f"AVERTISMENT: {lipsa} exemple fara vector in cache (excluse)")
    matrix = np.vstack(vecs)
    norms = np.linalg.norm(matrix, axis=1)
    norms[norms == 0.0] = 1.0
    return labels, matrix / norms[:, None]


def predict_top1(nbr_labels: list[str], nbr_sims: np.ndarray, prag: float) -> str | None:
    """Politica curenta: eticheta celui mai apropiat vecin, daca trece pragul."""
    if nbr_sims[0] < prag:
        return None
    return nbr_labels[0]


def predict_vote(nbr_labels: list[str], nbr_sims: np.ndarray, prag: float) -> str | None:
    """Vot ponderat cu similaritatea pe vecinii peste prag."""
    scoruri: dict[str, float] = {}
    for lab, sim in zip(nbr_labels, nbr_sims):
        if sim < prag:
            continue
        scoruri[lab] = scoruri.get(lab, 0.0) + float(sim)
    if not scoruri:
        return None
    return max(scoruri, key=lambda k: scoruri[k])


def evalueaza(labels, unit, indices, k_max, politici, praguri):
    """LOO: pentru fiecare index, top-k_max vecini (fara el insusi), apoi
    aplica fiecare (politica, prag) pe aceiasi vecini. Intoarce metrice agregate."""
    stats = {(p, prag): {"emis": 0, "corect": 0, "cod_gresit": 0} for p in politici for prag in praguri}
    n = len(indices)
    for count, i in enumerate(indices, 1):
        sims = unit @ unit[i]
        sims[i] = -1.0  # exclude exemplul insusi (leave-one-out)
        top = np.argpartition(-sims, k_max)[:k_max]
        top = top[np.argsort(-sims[top])]
        nbr_labels = [labels[j] for j in top]
        nbr_sims = sims[top]
        gold = labels[i]
        for (pname, prag), st in stats.items():
            pred = politici[pname](nbr_labels, nbr_sims, prag)
            if pred is None:
                continue
            st["emis"] += 1
            if pred == gold:
                st["corect"] += 1
            elif pred != NUL:
                st["cod_gresit"] += 1  # a emis un COD gresit (sau cod in loc de NUL)
        if count % 2000 == 0:
            print(f"  ...{count}/{n}", file=sys.stderr)
    return stats, n


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--db", default=os.path.join(_ROOT, "data", "autopass.db"))
    ap.add_argument("--sample", type=int, default=0, help="LOO doar pe un esantion (0 = tot corpusul)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--k", type=int, default=5, help="k pentru politica de vot")
    args = ap.parse_args()

    labels, unit = load_corpus(args.db)
    print(f"corpus: {len(labels)} exemple ({sum(1 for l in labels if l == NUL)} NUL)")

    indices = list(range(len(labels)))
    if args.sample and args.sample < len(indices):
        random.Random(args.seed).shuffle(indices)
        indices = indices[: args.sample]
    print(f"evaluare LOO pe {len(indices)} exemple, k={args.k}")

    politici = {
        "top1": predict_top1,
        f"vote{args.k}": predict_vote,
    }
    praguri = [0.5, 0.7, 0.8, 0.85, 0.88, 0.9, 0.92, 0.95]
    stats, n = evalueaza(labels, unit, indices, args.k, politici, praguri)

    print(f"\n{'politica':<8} {'prag':>5} {'coverage':>9} {'precizie':>9} {'cod-gresit':>11}")
    for pname in politici:
        for prag in praguri:
            st = stats[(pname, prag)]
            cov = st["emis"] / n
            prec = (st["corect"] / st["emis"]) if st["emis"] else 0.0
            wc = st["cod_gresit"] / n
            print(f"{pname:<8} {prag:>5.2f} {cov:>8.1%} {prec:>8.1%} {wc:>10.2%}")


if __name__ == "__main__":
    main()