GTI_GERENCIADOR_FISCAL/fiscal/central/app/ai/retrieval.py

"""Semantic candidate retrieval V2 — deterministic, zero LLM.

Three retrieval sources merged into one ranking:
  A. OFFICIAL: hierarchical search over the official NCM tree. Every 8-digit code is
     enriched with its ancestors' descriptions (chapter/position/subposition), so the
     very short legal text carries the vocabulary of its path (section 6).
  B. KB NEIGHBORS: the Knowledge Base acts as a commercial->official bridge — products
     whose commercial description resembles the query contribute their VALIDATED NCMs,
     weighted by similarity, tenant_count, observation_count and agreement (sections 2-3).
  C. HYBRID: single final ranking; the KB is only evidence for CANDIDATES, never a
     majority-vote fiscal decision.

Abbreviation lexicon (section 8) is derived statistically from the KB itself; commercial
normalization (section 7) strips units/sizes/colors/refs deterministically.
"""
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.orm import Session

from ..models import Dataset, KbKnowledge, KbObservation, KbProduct, NcmEntry, Tenant
from .fingerprint import normalize_description, token_stems

# ---------------- commercial normalization (section 7) ----------------

_UNITS = {"UN", "UNID", "UNIDADE", "CX", "CAIXA", "FD", "FDO", "FARDO", "PCT", "PACOTE", "PKG",
          "KG", "G", "GR", "GRAMA", "MG", "ML", "L", "LT", "LITRO", "LATA", "GARRAFA", "PET",
          "PC", "PECA", "DZ", "DUZIA", "PAR", "KIT", "JG", "PARA", "COM", "SEM", "C", "S"}
_DISCARD = _UNITS | {"MASC", "MASCULINO", "FEM", "FEMININO", "INF", "INFANTIL", "TAM", "TAMANHO",
                     "COR", "REF", "COD", "CODIGO", "MOD", "MODELO", "NOVA", "NOVO", "TIPO"}
_LEXEME_STRIP_PREFIXES = ("1", "2", "3", "4", "5", "6", "7", "8", "9", "0", ".")


def _is_size_or_code(token: str) -> bool:
    if token.isdigit():
        return True
    digits = "".join(char for char in token if char.isdigit())
    letters = "".join(char for char in token if char.isalpha())
    return bool(digits) and letters in _UNITS  # "1KG", "500ML", "2LT"


def commercial_tokens(description) -> set:
    """Stable commercial lexemes: units/sizes/colors/refs removed, stems kept."""
    tokens = []
    skip_next = False
    for token in normalize_description(description).split():
        if skip_next:
            skip_next = False
            continue
        if token in ("TAM", "COR", "REF", "COD", "MOD"):
            skip_next = True
            continue
        if token in _DISCARD or _is_size_or_code(token):
            continue
        tokens.append(token)
    return {_stem_safe(token) for token in tokens}


def _stem_safe(token: str) -> str:
    from .fingerprint import _stem
    return _stem(token)


def expand_stems(stems: set, lexicon: dict) -> set:
    expanded = set(stems)
    for stem in stems:
        target = lexicon.get(stem)
        if target:
            expanded.add(target)
    return expanded


def derive_abbreviation_lexicon(corpus: dict, min_support: int = 3) -> dict:
    """corpus: {commercial_description: ncm}. Statistical abbreviations only: a short token
    (3-5 chars) that is a strict prefix of a longer token, where BOTH forms occur across
    distinct products that share the same NCM. Support = (#products with the short form) +
    (#products with the long form); accepted when support >= min_support and both >= 1."""
    by_ncm: dict = {}
    for description, ncm in corpus.items():
        tokens = set(normalize_description(description).split())
        by_ncm.setdefault(ncm, []).append(tokens)
    best: dict = {}
    for descriptions in by_ncm.values():
        if len(descriptions) < 2:
            continue
        all_tokens = set().union(*descriptions)
        for short in all_tokens:
            if not (3 <= len(short) <= 5):
                continue
            short_count = sum(1 for tokens in descriptions if short in tokens)
            if short_count < 1:
                continue
            for long_token in all_tokens:
                if len(long_token) > len(short) and long_token.startswith(short):
                    long_count = sum(1 for tokens in descriptions if long_token in tokens)
                    support = short_count + long_count
                    if support >= min_support and long_count >= 1 and short_count >= 1:
                        current = best.get(short)
                        if current is None or support > current[1]:
                            best[short] = (long_token, support)
    return {short: target for short, (target, _) in best.items()}


# ---------------- enriched official corpus (sections 5-6) ----------------


@dataclass
class EnrichedNcmCorpus:
    documents: dict            # code -> enriched stems set
    descriptions: dict         # code -> official description
    positions: dict            # 4-digit position -> enriched stems (chapter + position)
    children: dict             # 4-digit position -> [8-digit codes]

    @staticmethod
    def build(session: Session) -> "EnrichedNcmCorpus":
        dataset = session.scalar(select(Dataset).where(Dataset.kind == "NCM", Dataset.status == "ACTIVE"))
        if dataset is None:
            return EnrichedNcmCorpus({}, {}, {}, {})
        entries = session.scalars(select(NcmEntry).where(NcmEntry.dataset_id == dataset.id)).all()
        by_code = {entry.code: entry.description for entry in entries}
        documents, positions, children = {}, {}, {}
        for code, description in by_code.items():
            if len(code) != 8:
                continue
            ancestors = [text for text in (by_code.get(code[:2]), by_code.get(code[:4]), by_code.get(code[:6])) if text]
            documents[code] = token_stems(description) | set().union(*[token_stems(text) for text in ancestors]) if ancestors else token_stems(description)
            positions.setdefault(code[:4], token_stems(by_code.get(code[:4], "")) | token_stems(by_code.get(code[:2], "")))
            children.setdefault(code[:4], []).append(code)
        return EnrichedNcmCorpus(documents, by_code, positions, children)

    def document(self, code: str) -> str:
        return " ".join(sorted(self.documents.get(code, set())))

    def items(self):
        return list(self.documents.items())


def _score(stems: set, document_stems: set, idf: dict | None = None) -> float:
    if not stems:
        return 0.0
    if idf:
        matched = sum(idf.get(stem, 1.0) for stem in stems if stem in document_stems)
        total = sum(idf.get(stem, 1.0) for stem in stems)
        return matched / total
    return len(stems & document_stems) / len(stems)


def hierarchical_candidates(corpus: EnrichedNcmCorpus, description, k: int = 20, max_positions: int = 15):
    """Two-stage: rank 4-digit positions first, then their 8-digit descendants (section 5)."""
    stems = commercial_tokens(description)
    if not stems:
        return []
    position_scores = sorted(((_score(stems, doc), position) for position, doc in corpus.positions.items()), reverse=True)
    ranked = []
    for position_score, position in position_scores[:max_positions]:
        if position_score <= 0:
            continue
        for code in corpus.children.get(position, []):
            code_score = _score(stems, corpus.documents[code])
            if code_score > 0:
                ranked.append((code, corpus.descriptions.get(code, ""), round(code_score, 4)))
    ranked.sort(key=lambda row: (-row[2], row[0]))
    return ranked[:k]


# ---------------- KB neighbor bridge (sections 2-3) ----------------


@dataclass
class KbNeighborIndex:
    entries: list  # (product_id, stems, ncm, tenant_ids:set, observation_count, has_conflict)

    @staticmethod
    def build(session: Session) -> "KbNeighborIndex":
        validated = {}
        for knowledge in session.scalars(select(KbKnowledge).where(KbKnowledge.knowledge_key == "validated_ncm")).all():
            validated.setdefault(knowledge.kb_product_id, set()).add(knowledge.value)
        conflicts = {product_id: len(values) > 1 for product_id, values in validated.items()}
        products = {product.id: product for product in session.scalars(select(KbProduct)).all()}
        observations = session.scalars(select(KbObservation)).all()
        per_product_tenants: dict = {}
        for observation in observations:
            per_product_tenants.setdefault(observation.kb_product_id, set()).add(observation.tenant_id)
        entries = []
        for product_id, values in validated.items():
            if conflicts.get(product_id):
                continue  # conflicting knowledge never feeds neighbors
            for value in values:
                product = products.get(product_id)
                if product is None or not product.normalized_description:
                    continue
                tenants = per_product_tenants.get(product_id, set())
                entries.append((product_id, commercial_tokens(product.normalized_description), value, tenants, len(tenants), False))
        return KbNeighborIndex(entries)

    def corpus(self) -> dict:
        return {" ".join(sorted(stems)): ncm for _, stems, ncm, _, _, _ in self.entries}


def kb_neighbor_candidates(index: KbNeighborIndex, description, k: int = 10, exclude_product_id: int | None = None):
    """Aggregate NCMs of lexically similar known products. `exclude_product_id` implements
    leave-one-out evaluation: the queried product never retrieves itself. Returns
    (ncm, neighbor_score, observation_count, tenant_count) — agreement is the union of
    distinct tenants behind the contributing neighbors, reported as evidence, never used
    as an automatic majority decision."""
    stems = commercial_tokens(description)
    if not stems:
        return []
    best: dict = {}
    for product_id, neighbor_stems, ncm, tenant_ids, observation_count, _ in index.entries:
        if exclude_product_id is not None and product_id == exclude_product_id:
            continue
        similarity = _score(stems, neighbor_stems)
        if similarity <= 0:
            continue
        current = best.get(ncm)
        if current is None:
            best[ncm] = [ncm, round(similarity, 4), observation_count, set(tenant_ids)]
        else:
            current[1] = max(current[1], round(similarity, 4))
            current[3] |= set(tenant_ids)
    ranked = sorted(((row[0], row[1], row[2], len(row[3])) for row in best.values()), key=lambda row: (-row[1], -row[3], row[0]))
    return ranked[:k]


# ---------------- hybrid ranking (section 4) ----------------

OFFICIAL_WEIGHT, KB_WEIGHT = 0.3, 0.7  # tuned on the V2 benchmark (KB bridge carries the signal;
# the official component breaks ties and covers products absent from the KB)


def hybrid_candidates(corpus: EnrichedNcmCorpus, index: KbNeighborIndex, description, k: int = 20, lexicon: dict | None = None, exclude_product_id: int | None = None):
    stems = expand_stems(commercial_tokens(description), lexicon or {})
    official = hierarchical_candidates(corpus, " ".join(sorted(stems)) or description, k=k)
    neighbors = kb_neighbor_candidates(index, description, k=k, exclude_product_id=exclude_product_id)
    combined: dict = {}
    for code, text, score in official:
        combined[code] = [code, text, OFFICIAL_WEIGHT * score, score, 0.0, 0]
    for code, similarity, observation_count, tenant_count in neighbors:
        row = combined.setdefault(code, [code, "", 0.0, 0.0, 0.0, 0])
        row[2] += KB_WEIGHT * similarity
        row[4] = similarity
        row[5] = tenant_count
    ranked = sorted(combined.values(), key=lambda row: (-round(row[2], 4), row[0]))
    return [(row[0], row[1], round(row[2], 4), row[3], row[4], row[5]) for row in ranked[:k]]