src/fidelity.py

"""
fidelity.py — Min-Aggregated Fidelity Scoring

Implements equation 23 from the paper:
    F(S, S') = min(F_jaccard, F_cosine, F_nli)

The min-aggregation is the key design choice: a signal must pass ALL
three checks, not just one. This prevents gaming (e.g., high cosine
with destroyed modal operators).

All three metrics work without transformer models:
  - Jaccard: set overlap on commitment canonical forms
  - Cosine: TF-IDF vectors on commitment text
  - NLI proxy: structural entailment check on modal operators + key terms

When transformer-based NLI is available (e.g., on HuggingFace),
it replaces the proxy. The interface is the same.
"""

import re
import math
from typing import Set, Dict, List, Optional
from collections import Counter


# ---------------------------------------------------------------------------
# Jaccard fidelity — exact canonical match
# ---------------------------------------------------------------------------

def fidelity_jaccard(original: Set[str], transformed: Set[str]) -> float:
    """
    Jaccard index on canonical commitment strings.
    
    This is the strictest metric: requires exact canonical match.
    Returns 1.0 if both empty (vacuous truth — no commitments to lose).
    Returns 0.0 if one is empty and the other isn't.
    """
    if not original and not transformed:
        return 1.0
    if not original or not transformed:
        return 0.0
    intersection = len(original & transformed)
    union = len(original | transformed)
    return intersection / union


# ---------------------------------------------------------------------------
# Cosine fidelity — TF-IDF word vectors
# ---------------------------------------------------------------------------

def _tokenize(text: str) -> List[str]:
    """Simple word tokenizer. Lowercase, split on non-alphanumeric."""
    return re.findall(r'[a-z0-9]+', text.lower())


def _tf(tokens: List[str]) -> Dict[str, float]:
    """Term frequency."""
    counts = Counter(tokens)
    total = len(tokens)
    if total == 0:
        return {}
    return {t: c / total for t, c in counts.items()}


def _idf(doc_tokens_list: List[List[str]]) -> Dict[str, float]:
    """Inverse document frequency."""
    n_docs = len(doc_tokens_list)
    if n_docs == 0:
        return {}
    
    df = Counter()
    for tokens in doc_tokens_list:
        unique = set(tokens)
        for t in unique:
            df[t] += 1
    
    return {t: math.log(n_docs / count) + 1.0 for t, count in df.items()}


def _tfidf_vector(tf: Dict[str, float], idf: Dict[str, float], vocab: Set[str]) -> Dict[str, float]:
    """TF-IDF vector over shared vocabulary."""
    return {t: tf.get(t, 0.0) * idf.get(t, 0.0) for t in vocab}


def _cosine_sim(v1: Dict[str, float], v2: Dict[str, float]) -> float:
    """Cosine similarity between two sparse vectors."""
    keys = set(v1.keys()) | set(v2.keys())
    dot = sum(v1.get(k, 0.0) * v2.get(k, 0.0) for k in keys)
    norm1 = math.sqrt(sum(v ** 2 for v in v1.values())) or 1e-10
    norm2 = math.sqrt(sum(v ** 2 for v in v2.values())) or 1e-10
    return dot / (norm1 * norm2)


def fidelity_cosine(original: Set[str], transformed: Set[str]) -> float:
    """
    Cosine similarity on TF-IDF vectors of commitment text.
    
    Concatenates all commitments into a single document per set,
    computes TF-IDF, returns cosine similarity.
    
    More forgiving than Jaccard — catches paraphrased commitments
    that share vocabulary but differ in exact wording.
    """
    if not original and not transformed:
        return 1.0
    if not original or not transformed:
        return 0.0
    
    orig_text = ' '.join(original)
    trans_text = ' '.join(transformed)
    
    orig_tokens = _tokenize(orig_text)
    trans_tokens = _tokenize(trans_text)
    
    if not orig_tokens or not trans_tokens:
        return 0.0
    
    # Build IDF from both documents
    idf = _idf([orig_tokens, trans_tokens])
    vocab = set(idf.keys())
    
    tf_orig = _tf(orig_tokens)
    tf_trans = _tf(trans_tokens)
    
    v_orig = _tfidf_vector(tf_orig, idf, vocab)
    v_trans = _tfidf_vector(tf_trans, idf, vocab)
    
    return _cosine_sim(v_orig, v_trans)


# ---------------------------------------------------------------------------
# NLI proxy — structural entailment without transformer
# ---------------------------------------------------------------------------

# Key terms that must survive: modal operators, numbers, named entities
MODAL_TERMS = {
    'must', 'shall', 'cannot', 'required', 'prohibited', 'forbidden',
    'always', 'never', 'not', 'no',
}

NUMBER_RE = re.compile(r'\$?\d[\d,.]*')
TIME_RE = re.compile(r'\b(?:monday|tuesday|wednesday|thursday|friday|saturday|sunday|'
                     r'january|february|march|april|may|june|july|august|september|'
                     r'october|november|december|\d{1,2}(?:st|nd|rd|th)?)\b', re.I)


def _extract_key_terms(text: str) -> Set[str]:
    """Extract terms that are structurally significant for commitment identity."""
    tokens = set(_tokenize(text))
    
    key_terms = set()
    
    # Modal operators present
    key_terms.update(tokens & MODAL_TERMS)
    
    # Numbers (amounts, thresholds, counts)
    for match in NUMBER_RE.finditer(text):
        key_terms.add(match.group().lower())
    
    # Time references
    for match in TIME_RE.finditer(text):
        key_terms.add(match.group().lower())
    
    return key_terms


def fidelity_nli_proxy(original: Set[str], transformed: Set[str]) -> float:
    """
    Structural entailment proxy for NLI.
    
    Checks whether the KEY TERMS (modals, numbers, time references)
    from original commitments survive in transformed commitments.
    
    This is not full NLI — it's a conservative proxy that catches
    the most common failure mode: losing the modal operator or
    the specific quantity/deadline while retaining general topic words.
    
    When a real NLI model is available, replace this function.
    """
    if not original and not transformed:
        return 1.0
    if not original or not transformed:
        return 0.0
    
    orig_text = ' '.join(original)
    trans_text = ' '.join(transformed)
    
    orig_keys = _extract_key_terms(orig_text)
    trans_keys = _extract_key_terms(trans_text)
    
    if not orig_keys:
        # No structural terms to check — can't assess, return neutral
        return 0.5
    
    # What fraction of original key terms survived?
    preserved = len(orig_keys & trans_keys)
    total = len(orig_keys)
    
    return preserved / total


# ---------------------------------------------------------------------------
# Min-aggregated fidelity — equation 23
# ---------------------------------------------------------------------------

def fidelity_score(original: Set[str], transformed: Set[str]) -> float:
    """
    Min-aggregated fidelity score per equation 23:
        F(S, S') = min(F_jaccard, F_cosine, F_nli)
    
    A signal must pass ALL three checks. This prevents:
    - High Jaccard with semantically different content (false exact match)
    - High cosine with destroyed modal operators (topic match, no commitment)
    - High NLI with completely reworded unrelated commitments
    
    Returns a float in [0.0, 1.0].
    """
    f_j = fidelity_jaccard(original, transformed)
    f_c = fidelity_cosine(original, transformed)
    f_n = fidelity_nli_proxy(original, transformed)
    
    return min(f_j, f_c, f_n)


def fidelity_breakdown(original: Set[str], transformed: Set[str]) -> dict:
    """
    Return all three component scores plus the min-aggregated score.
    Useful for diagnostics.
    """
    f_j = fidelity_jaccard(original, transformed)
    f_c = fidelity_cosine(original, transformed)
    f_n = fidelity_nli_proxy(original, transformed)
    
    return {
        'jaccard': f_j,
        'cosine': f_c,
        'nli_proxy': f_n,
        'min_aggregated': min(f_j, f_c, f_n),
    }


# ---------------------------------------------------------------------------
# Legacy interface
# ---------------------------------------------------------------------------

def jaccard(a: Set[str], b: Set[str]) -> float:
    """Backward compatible."""
    return fidelity_jaccard(a, b)

def jaccard_index(a, b) -> float:
    """Backward compatible."""
    return fidelity_jaccard(set(a), set(b))