chunk_analyzer.py

"""
Chunk Analyzer - Pre-analysis tool for RAG chunk quality.

Reads a saved pickle run and analyzes chunk content:
  - Keyword density (most common terms across all chunks)
  - TF-IDF top terms per chunk (what makes each chunk unique)
  - TF-IDF cosine similarity matrix (find near-duplicate chunks)
  - Boilerplate ratio per chunk (using known patterns)

Usage:
  python chunk_analyzer.py --run <run_name>                   # Analyze a saved run
  python chunk_analyzer.py --run <run_name> --top-terms 20    # Show top 20 keywords
  python chunk_analyzer.py --run <run_name> --duplicates 0.95 # Find chunks with sim > 0.95
  python chunk_analyzer.py --run <run_name> --export           # Export analysis to JSON
"""

import argparse
import os
import sys
import json
import pickle
import math
import re
from collections import Counter
from datetime import datetime

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
RUNS_DIR = os.path.join(SCRIPT_DIR, "runs")

# Terminal colors
C = {
    "RESET": "\033[0m",
    "DIM": "\033[2m",
    "BOLD": "\033[1m",
    "CYAN": "\033[36m",
    "GREEN": "\033[32m",
    "YELLOW": "\033[33m",
    "MAGENTA": "\033[35m",
    "RED": "\033[31m",
}


def log(tag, msg):
    ts = datetime.now().strftime("%H:%M:%S")
    colors = {"LOAD": C["CYAN"], "TF-IDF": C["GREEN"], "KEYWORD": C["MAGENTA"],
              "DUPES": C["YELLOW"], "INFO": C["DIM"], "ERROR": C["RED"]}
    color = colors.get(tag, C["RESET"])
    print(f"{C['DIM']}{ts}{C['RESET']} {color}[{tag}]{C['RESET']}  {msg}")


def separator(title):
    print(f"\n{C['BOLD']}{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}{C['RESET']}\n")


# ============================================================================
# LOAD RUN
# ============================================================================
def load_run(name):
    """Load texts and metadatas from a saved pickle run."""
    path = os.path.join(RUNS_DIR, f"{name}.pkl")
    if not os.path.exists(path):
        log("ERROR", f"Run not found: runs/{name}.pkl")
        available = [f[:-4] for f in os.listdir(RUNS_DIR) if f.endswith(".pkl")]
        if available:
            log("INFO", f"Available: {', '.join(sorted(available))}")
        sys.exit(1)

    with open(path, "rb") as f:
        data = pickle.load(f)

    log("LOAD", f"Run: {name} ({data['num_vectors']} chunks, {data['dims']} dims)")
    log("LOAD", f"  params: {json.dumps(data.get('params', {}))}")
    log("LOAD", f"  cost: ${data.get('est_cost_usd', 0):.6f} | time: {data.get('embed_seconds', 0):.1f}s")
    return data["texts"], data["metadatas"], data.get("params", {})


# ============================================================================
# TOKENIZER (simple whitespace + punctuation split)
# ============================================================================
STOP_WORDS = {
    "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for",
    "of", "with", "by", "from", "as", "is", "was", "are", "were", "be",
    "been", "being", "have", "has", "had", "do", "does", "did", "will",
    "would", "could", "should", "may", "might", "shall", "can", "need",
    "it", "its", "this", "that", "these", "those", "i", "you", "he", "she",
    "we", "they", "me", "him", "her", "us", "them", "my", "your", "his",
    "our", "their", "what", "which", "who", "whom", "how", "when", "where",
    "why", "all", "each", "every", "both", "few", "more", "most", "other",
    "some", "such", "no", "nor", "not", "only", "own", "same", "so", "than",
    "too", "very", "just", "about", "above", "after", "again", "also",
    "any", "because", "before", "between", "during", "if", "into", "new",
    "over", "then", "through", "under", "up", "out", "s", "t", "re", "ve",
    "ll", "d", "m", "o", "don", "didn", "doesn", "won", "isn", "aren",
    "wasn", "weren", "hasn", "haven", "hadn", "couldn", "shouldn", "wouldn",
}


def tokenize(text):
    """Split text into lowercase tokens, removing punctuation and stop words."""
    tokens = re.findall(r'[a-zA-Z][a-zA-Z0-9+#._-]{1,}', text.lower())
    return [t for t in tokens if t not in STOP_WORDS and len(t) > 1]


# ============================================================================
# KEYWORD DENSITY
# ============================================================================
def analyze_keywords(texts, top_n=30):
    """Count term frequency across all chunks."""
    separator("KEYWORD DENSITY")

    global_counts = Counter()
    for text in texts:
        tokens = tokenize(text)
        global_counts.update(tokens)

    total_tokens = sum(global_counts.values())
    unique_tokens = len(global_counts)
    log("KEYWORD", f"{total_tokens:,} total tokens | {unique_tokens:,} unique terms")

    top = global_counts.most_common(top_n)
    print(f"\n  {'Term':<30} {'Count':>8}  {'% of all':>8}")
    print(f"  {'-'*30} {'-'*8}  {'-'*8}")
    for term, count in top:
        pct = count / total_tokens * 100
        bar = "#" * min(int(pct * 2), 40)
        print(f"  {term:<30} {count:>8,}  {pct:>7.2f}%  {C['DIM']}{bar}{C['RESET']}")
    print()

    return global_counts


# ============================================================================
# TF-IDF
# ============================================================================
def compute_tfidf(texts):
    """Compute TF-IDF vectors for all chunks."""
    separator("TF-IDF ANALYSIS")

    log("TF-IDF", f"Computing TF-IDF for {len(texts)} chunks...")

    # Tokenize all documents
    docs = [tokenize(text) for text in texts]

    # Build vocabulary
    vocab = sorted(set(token for doc in docs for token in doc))
    vocab_idx = {term: i for i, term in enumerate(vocab)}
    log("TF-IDF", f"Vocabulary: {len(vocab)} terms")

    # Document frequency (how many docs contain each term)
    df = Counter()
    for doc in docs:
        df.update(set(doc))

    n_docs = len(docs)

    # Compute TF-IDF vectors
    tfidf_vectors = []
    for doc in docs:
        tf = Counter(doc)
        doc_len = len(doc) if doc else 1
        vec = [0.0] * len(vocab)
        for term, count in tf.items():
            idx = vocab_idx[term]
            tf_val = count / doc_len
            idf_val = math.log((n_docs + 1) / (df[term] + 1)) + 1  # smoothed IDF
            vec[idx] = tf_val * idf_val
        # L2 normalize
        norm = math.sqrt(sum(v * v for v in vec))
        if norm > 0:
            vec = [v / norm for v in vec]
        tfidf_vectors.append(vec)

    # Show top TF-IDF terms for first few chunks
    log("TF-IDF", "Top TF-IDF terms per chunk (first 5):")
    for i in range(min(5, len(docs))):
        scored = [(vocab[j], tfidf_vectors[i][j]) for j in range(len(vocab)) if tfidf_vectors[i][j] > 0]
        scored.sort(key=lambda x: x[1], reverse=True)
        top5 = ", ".join(f"{t}({s:.3f})" for t, s in scored[:5])
        print(f"  {C['DIM']}chunk[{i}]{C['RESET']} {top5}")
    print()

    return tfidf_vectors, vocab


# ============================================================================
# NEAR-DUPLICATE DETECTION
# ============================================================================
def find_duplicates(tfidf_vectors, texts, metadatas, threshold=0.90):
    """Find chunk pairs with TF-IDF cosine similarity above threshold."""
    separator(f"NEAR-DUPLICATE DETECTION (threshold={threshold})")

    n = len(tfidf_vectors)
    log("DUPES", f"Comparing {n} chunks (cosine sim > {threshold})...")

    duplicates = []
    # For large datasets, only check a sample to avoid O(n^2) blowup
    check_limit = min(n, 2000)
    if n > check_limit:
        log("DUPES", f"Sampling first {check_limit} chunks (full dataset too large for O(n^2))")

    for i in range(check_limit):
        for j in range(i + 1, check_limit):
            # Dot product (vectors are already L2-normalized)
            sim = sum(a * b for a, b in zip(tfidf_vectors[i], tfidf_vectors[j]))
            if sim >= threshold:
                duplicates.append({
                    "chunk_a": i,
                    "chunk_b": j,
                    "similarity": round(sim, 6),
                    "source_a": metadatas[i].get("source", "?") if i < len(metadatas) else "?",
                    "source_b": metadatas[j].get("source", "?") if j < len(metadatas) else "?",
                    "preview_a": texts[i][:80].replace("\n", " "),
                    "preview_b": texts[j][:80].replace("\n", " "),
                })

    duplicates.sort(key=lambda x: x["similarity"], reverse=True)

    log("DUPES", f"Found {len(duplicates)} pairs above {threshold}")

    for d in duplicates[:20]:
        same = "SAME" if d["source_a"] == d["source_b"] else "DIFF"
        print(f"  {C['BOLD']}{d['similarity']:.4f}{C['RESET']} [{same}] {d['source_a']} <-> {d['source_b']}")
        print(f"    A: \"{d['preview_a']}...\"")
        print(f"    B: \"{d['preview_b']}...\"")
        print()

    if len(duplicates) > 20:
        log("DUPES", f"  ... and {len(duplicates) - 20} more (use --export to see all)")

    return duplicates


# ============================================================================
# EXPORT
# ============================================================================
def export_analysis(run_name, keywords, duplicates, params):
    """Export analysis results to JSON."""
    path = os.path.join(SCRIPT_DIR, f"chunk_analysis_{run_name}.json")

    report = {
        "run_name": run_name,
        "analyzed_at": datetime.now().isoformat(),
        "pipeline_params": params,
        "keyword_density": {
            "total_unique_terms": len(keywords),
            "top_50": [{"term": t, "count": c} for t, c in keywords.most_common(50)],
        },
        "near_duplicates": {
            "total_pairs": len(duplicates),
            "pairs": duplicates[:200],
        },
    }

    with open(path, "w") as f:
        json.dump(report, f, indent=2)

    size_kb = os.path.getsize(path) / 1024
    log("INFO", f"Exported: {os.path.basename(path)} ({size_kb:.1f} KB)")
    return path


# ============================================================================
# MAIN
# ============================================================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chunk Analyzer - pre-analysis for RAG tuning")
    parser.add_argument("--run", required=True, help="Name of the saved run to analyze")
    parser.add_argument("--top-terms", type=int, default=30, help="Number of top keywords to show (default: 30)")
    parser.add_argument("--duplicates", type=float, default=0.90, help="Cosine similarity threshold for duplicates (default: 0.90)")
    parser.add_argument("--export", action="store_true", help="Export analysis to chunk_analysis_<run>.json")
    args = parser.parse_args()

    texts, metadatas, params = load_run(args.run)

    keywords = analyze_keywords(texts, top_n=args.top_terms)
    tfidf_vectors, vocab = compute_tfidf(texts)
    duplicates = find_duplicates(tfidf_vectors, texts, metadatas, threshold=args.duplicates)

    if args.export:
        export_analysis(args.run, keywords, duplicates, params)

    separator("SUMMARY")
    print(f"  Chunks analyzed:     {len(texts):,}")
    print(f"  Unique terms:        {len(keywords):,}")
    print(f"  Near-duplicate pairs: {len(duplicates)}")
    print(f"  Duplicate threshold:  {args.duplicates}")
    print()