#!/usr/bin/env python3
"""
Jaccard matcher v6 — file-based with keyword filtering.
Uses /shm/arxiv_texts.tsv for fast local matching.
Only uses keywords >= 6 chars to reduce noise.
"""

import json
import re
from collections import defaultdict
from pathlib import Path

CONCEPTS_PATH = "/tmp/cornfield_concepts.json"
ARXIV_TEXTS_PATH = "/shm/arxiv_texts.tsv"
TOP_K = 5
MIN_COVERAGE = 0.15
MIN_KEYWORD_LEN = 6  # Only use keywords >= 6 chars


def extract_keywords(concept: dict) -> set[str]:
    words = set()
    for tag in concept.get("tags", []):
        for w in re.split(r'[_\s\-/]+', tag):
            w = w.strip().lower()
            if len(w) >= MIN_KEYWORD_LEN:
                words.add(w)
    name = concept.get("name", "")
    for w in re.sub(r'([a-z])([A-Z])', r'\1 \2', name).split():
        w = w.strip().lower()
        if len(w) >= MIN_KEYWORD_LEN:
            words.add(w)
    desc = concept.get("description", "")
    for w in re.findall(r'[a-zA-Z]{6,}', desc):
        words.add(w.lower())
    novelty = concept.get("novelty_statement", "")
    for w in re.findall(r'[a-zA-Z]{6,}', novelty):
        words.add(w.lower())
    return words


def main():
    data = json.loads(Path(CONCEPTS_PATH).read_text())
    concepts = [c for c in data["concepts"] if c.get("id", "").startswith("cf_")]
    print(f"Processing {len(concepts)} cornfield concepts")

    # Extract keywords for all concepts
    concept_keywords = {}
    for c in concepts:
        kw = extract_keywords(c)
        if len(kw) >= 2:
            concept_keywords[c["id"]] = kw
    print(f"Keywords extracted for {len(concept_keywords)} concepts (min len {MIN_KEYWORD_LEN})")

    # Collect all unique keywords
    all_keywords = set()
    for kw_set in concept_keywords.values():
        all_keywords.update(kw_set)
    print(f"Total unique keywords: {len(all_keywords)}")

    # Phase 1: Build inverted index
    print("Phase 1: Building inverted index...")
    inverted = defaultdict(set)
    count = 0
    with open(ARXIV_TEXTS_PATH) as f:
        for line in f:
            parts = line.split('\t', 1)
            if len(parts) < 2:
                continue
            pid, text = parts[0], parts[1].lower()
            count += 1
            if count % 200000 == 0:
                print(f"  {count} papers scanned...")
            for kw in all_keywords:
                if kw in text:
                    inverted[kw].add(pid)
    print(f"  Done: {count} papers, {len(inverted)} keywords with hits")

    # Phase 2: Match concepts
    print("Phase 2: Matching concepts...")
    results = []
    for cid, kw_set in concept_keywords.items():
        candidates = defaultdict(set)
        for kw in kw_set:
            for pid in inverted.get(kw, set()):
                candidates[pid].add(kw)

        scored = []
        for pid, matched in candidates.items():
            cov = len(matched) / len(kw_set)
            if cov >= MIN_COVERAGE:
                scored.append((pid, cov, matched))
        scored.sort(key=lambda x: x[1], reverse=True)
        top = scored[:TOP_K]

        if top:
            print(f"  {cid}: {len(candidates)} candidates, top={top[0][1]:.3f}")
            results.append({
                "concept_id": cid,
                "keywords": sorted(kw_set),
                "candidates_found": len(candidates),
                "top_matches": [
                    {"paper_id": pid, "coverage": round(cov, 4), "matched_keywords": sorted(mk)}
                    for pid, cov, mk in top
                ]
            })
        else:
            print(f"  {cid}: no matches above {MIN_COVERAGE}")

    out = Path("/tmp/jaccard_matches.json")
    out.write_text(json.dumps(results, indent=2))
    total = sum(len(r["top_matches"]) for r in results)
    print(f"\nDone: {len(results)} concepts, {total} citations → {out}")


if __name__ == "__main__":
    main()