#!/usr/bin/env python3
"""Merge all search/snowball results into one de-duplicated screening sheet.

Usage: python3 merge.py            (reads ../search/results/*.csv)

Writes ../search/screening.csv. Columns:
  found_by        which searches / snowball steps returned the paper
  auto_stage      what the automatic filters decided (see protocol.md section 5):
                    EXCLUDE-pre2020 | EXCLUDE-offtopic | SCREEN
  ai_proposal     filled later: first-pass include/exclude suggestion
  decision        YOURS to fill: include / exclude / maybe
  reason          short reason for the decision

Existing decision/ai_proposal/reason values are preserved on re-runs, so you
can re-merge after new searches without losing screening work.
"""
import csv
import glob
import os
import re

HERE = os.path.dirname(os.path.abspath(__file__))
SEARCH = os.path.normpath(os.path.join(HERE, "..", "search"))
OUT = os.path.join(SEARCH, "screening.csv")

# Topic filter (protocol.md section 5, criterion I2): a paper goes to manual
# screening only if its title/abstract touches attention-logit control or
# training stability of transformers.
TOPIC = re.compile(
    r"qk[- ]?(norm|layer ?norm)|query[- ]key norm|attention logit|logit (growth|explosion|cap|soft)"
    r"|entropy collapse|attention entropy|training (instabilit|stabilit)|loss spike|divergen"
    r"|learning[- ]rate (sensitiv|stabil|transfer)|maximal update|\bmu-?p\b|μp|z-loss"
    r"|massive activation|attention sink|layer ?norm|rms ?norm|normali[sz]ation",
    re.I,
)
TRANSFORMER = re.compile(r"transformer|attention|language model|\bllm|\bvit\b|gpt", re.I)
LLM_SEARCH = re.compile(r"systematic|literature|screening|review", re.I)
# Seed set (protocol.md section 3): always screened, whatever the filters say.
SEEDS = {"2010.04245", "2302.05442", "2303.06296", "2309.14322", "2602.02522", "2410.16682"}


def key(r):
    if r.get("arxiv", "").lower().startswith("10.48550/arxiv."):  # old rows: DOI in arxiv column
        r["arxiv"] = r["arxiv"].lower().split("arxiv.")[-1]
    if not r.get("arxiv") and "10.48550/arxiv." in r.get("doi", "").lower():
        r["arxiv"] = r["doi"].lower().split("arxiv.")[-1]
    if r.get("arxiv"):
        return "arxiv:" + r["arxiv"].lower()
    if r.get("doi"):
        return "doi:" + r["doi"].lower()
    return "title:" + re.sub(r"[^a-z0-9]", "", r["title"].lower())


def main():
    merged = {}
    for path in sorted(glob.glob(os.path.join(SEARCH, "results", "*.csv"))):
        name = os.path.basename(path)[:-4]
        for r in csv.DictReader(open(path)):
            if not r.get("title"):
                continue
            k = key(r)
            m = merged.setdefault(k, {**r, "found_by": set()})
            m["found_by"].add(name)
            for f in ("arxiv", "doi", "abstract", "venue"):
                if not m.get(f) and r.get(f):
                    m[f] = r[f]
            try:
                m["citations"] = max(int(m.get("citations") or 0), int(r.get("citations") or 0))
            except ValueError:
                pass
    # also merge title-duplicates that slipped through on different id types
    by_title = {}
    for k, m in list(merged.items()):
        t = re.sub(r"[^a-z0-9]", "", m["title"].lower())
        if t in by_title:
            by_title[t]["found_by"] |= m["found_by"]
            for f in ("arxiv", "doi", "abstract"):
                by_title[t][f] = by_title[t].get(f) or m.get(f)
            del merged[k]
        else:
            by_title[t] = m

    old = {}
    if os.path.exists(OUT):
        for r in csv.DictReader(open(OUT)):
            old[r["key"]] = r

    rows = []
    for k, m in merged.items():
        text = f"{m['title']} {m.get('abstract', '')}"
        lit_only = all(f.startswith("L") for f in m["found_by"])
        try:
            year = int(m.get("year") or 0)
        except ValueError:
            year = 0
        if any(sd in k for sd in SEEDS) or "findings-emnlp.379" in k:
            stage = "SCREEN"
        elif lit_only:
            stage = "SCREEN" if LLM_SEARCH.search(m["title"]) else "EXCLUDE-offtopic"
        elif year and year < 2020:
            stage = "EXCLUDE-pre2020"
        elif TOPIC.search(text) and TRANSFORMER.search(text):
            stage = "SCREEN"
        else:
            stage = "EXCLUDE-offtopic"
        prev = old.get(k, {})
        rows.append({
            "key": k, "year": m.get("year", ""), "title": m["title"],
            "first_author": m.get("first_author", ""), "venue": m.get("venue", ""),
            "citations": m.get("citations", ""), "found_by": ";".join(sorted(m["found_by"])),
            "n_sources": len(m["found_by"]), "auto_stage": stage,
            "ai_proposal": prev.get("ai_proposal", ""), "decision": prev.get("decision", ""),
            "reason": prev.get("reason", ""), "abstract": m.get("abstract", ""),
        })
    rows.sort(key=lambda r: (r["auto_stage"] != "SCREEN", -r["n_sources"], str(r["year"])))
    with open(OUT, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    from collections import Counter
    c = Counter(r["auto_stage"] for r in rows)
    print(f"{len(rows)} unique records -> {os.path.relpath(OUT, SEARCH)}: {dict(c)}")


if __name__ == "__main__":
    main()