phase1/scripts/build_citation_graph.py

#!/usr/bin/env python3
"""Rebuild the citation graph from BODY TEXT — edges_v2.jsonl + case_aliases.json.

The legacy graph (edges.jsonl) parses only the headnote's "Case Law Cited" list:
measured ~3.6 edges/doc vs ~96.6 in-text citation mentions/doc (~4% capture). This
script regexes every chunk of every judgment for SCR / INSC / AIR-SC / SCC citations,
resolves them against the same normalized resolver the runtime uses
(neutral_citation + equivalent_citations), keeps the surrounding sentence as `para`
(the citing court's own description of the precedent — later embedded as the
citation-context retrieval arm), and merges the legacy edges so real treatment labels
survive.

Also emits case_aliases.json: for heavily-cited targets, the short name phrase that
most often precedes their citation in other judgments ("Kesavananda Bharati",
"Maneka Gandhi") -> doc_id, for known-item lookup of famous-name queries.

Run:  python phase1/scripts/build_citation_graph.py [data_dir]   (~5-10 min, CPU)
Out:  <data_dir>/edges_v2.jsonl, <data_dir>/case_aliases.json
Needs: corpus_ledger.jsonl (for sibling clusters; run build_ledger.py first).
"""
import json, os, re, sys, time
from collections import Counter, defaultdict

data_dir = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("THEMIS_DATA", "phase1/data/thor_artifacts")
t0 = time.time()

# ---------- resolvers (identical normalization to tools.py:66-68) ----------
meta = {}
for line in open(os.path.join(data_dir, "escr_meta.jsonl"), encoding="utf-8"):
    m = json.loads(line)
    meta.setdefault(m["doc_id"], m)
resolver = {}
for d, m in meta.items():
    for k in [m.get("neutral_citation")] + (m.get("equivalent_citations") or []):
        if not k: continue
        nk = re.sub(r"\s+", " ", k.replace(".", "")).strip().upper()
        resolver.setdefault(nk, d)
        # SUPP volumes are often printed WITHOUT the volume digit ("[1973] Supp. SCR 1"
        # for "[1973] SUPP. 1 S.C.R. 1") — register the volume-less variant too
        if " SUPP " in nk:
            resolver.setdefault(re.sub(r"(SUPP) \d+ (SCR)", r"\1 \2", nk), d)

cluster_of = {}
lp = os.path.join(data_dir, "corpus_ledger.jsonl")
if os.path.exists(lp):
    for line in open(lp, encoding="utf-8"):
        r = json.loads(line)
        if r.get("cluster_id"): cluster_of[r["doc_id"]] = r["cluster_id"]

# citation patterns; SUPP volumes included. INSC never appears in body text (modern
# administrative id), so the resolvable anchors are SCR (in meta) and, via the learned
# crosswalk below, SCC/AIR (the dominant forms in body text, absent from meta).
SCR_P = re.compile(r"\[\s*\d{4}\s*\]\s*(?:SUPP\.?\s*)?\d*\s*S\.?\s*C\.?\s*R\.?\s*\d+", re.I)
SCC_P = re.compile(r"\(\s*\d{4}\s*\)\s*\d+\s*S\.?C\.?C\.?\s*\d+", re.I)
AIR_P = re.compile(r"A\.?I\.?R\.?\s+\d{4}\s+S\.?C\.?\s+\d+", re.I)
CITE = re.compile(f"({SCR_P.pattern})|({SCC_P.pattern})|({AIR_P.pattern})", re.I)

def norm_key(s):
    return re.sub(r"\s+", " ", s.replace(".", "")).strip().upper()

# crosswalk learned from parallel citations printed side by side in judgments/headnotes
# ("... [1983] 1 SCR 145 : (1980) 2 SCC 684 : AIR 1980 SC 898 ...")
crosswalk = {}

def resolve(s):
    k = norm_key(s)
    return resolver.get(k) or crosswalk.get(k)

# alias: the full "Petitioner v. Respondent" name immediately before the citation;
# alias = the petitioner side (the memorable half: "Kesavananda Bharati", "Maneka Gandhi")
NAME_BEFORE = re.compile(
    r"([A-Z][A-Za-z.'&()-]*(?:\s+[A-Za-z.'&()-]+){0,6}\s+v\.?s?\.?\s+"
    r"[A-Z][A-Za-z.'&()-]*(?:\s+[A-Za-z.'&()-]+){0,5})\s*[,(\[]?\s*$")
_TITLE = re.compile(r"\b(shri|smt|sri|mst|dr|mr|mrs|ms|justice|his|holiness|m/s|the|state|union|of|india)\b\.?", re.I)
_AL_GENERIC = {"state", "union", "india", "government", "collector", "commissioner", "corporation",
               "municipal", "board", "authority", "bank", "company", "ltd", "limited", "co"}

# ---------- stream chunks grouped by doc ----------
edges_best = {}                 # (from,target) -> longest context
mentions = unresolved = 0
alias_votes = defaultdict(Counter)   # target -> Counter(alias phrase)
cur_doc, buf = None, []
pat = re.compile(r'"doc_id":\s*"([^"]+)"')

from collections import Counter as _Counter
crosswalk_votes = defaultdict(_Counter)
_SEP = re.compile(r"^[\s:;,=]*$")

def learn_crosswalk(text):
    """Parallel citations printed adjacently teach SCC/AIR -> doc mappings. Guarded:
    only pure separator chars between the two cites (dense citation LISTS put a case
    name between different cases' cites), majority vote + year sanity applied after."""
    anchors = [(m.start(), m.end(), resolver.get(norm_key(m.group(0)))) for m in SCR_P.finditer(text)]
    anchors = [(s, e, d) for s, e, d in anchors if d]
    for p in (SCC_P, AIR_P):
        for m in p.finditer(text):
            for s, e, d in anchors:
                if 0 <= m.start() - e <= 6 and _SEP.match(text[e:m.start()]):
                    crosswalk_votes[norm_key(m.group(0))][d] += 1; break
                if 0 <= s - m.end() <= 6 and _SEP.match(text[m.end():s]):
                    crosswalk_votes[norm_key(m.group(0))][d] += 1; break

def settle_crosswalk(decision_year):
    for k, votes in crosswalk_votes.items():
        (top, n), total = votes.most_common(1)[0], sum(votes.values())
        if n < 2 or n / total < 0.67: continue
        ym = re.search(r"\d{4}", k)
        dy = decision_year.get(top)
        if ym and dy and abs(int(ym.group(0)) - dy) > 2: continue     # AIR/SCC year must match the case
        crosswalk[k] = top

def petitioner_alias(full_name):
    pet = re.split(r"\s+v\.?s?\.?\s+", full_name, 1, flags=re.I)[0]
    pet = _TITLE.sub(" ", pet)
    words = [w.strip(".,'()&-") for w in pet.split()]
    words = [w for w in words if len(w) >= 3 and w[0].isupper() and w.lower() not in _AL_GENERIC
             and not (len(w) <= 3 and w.isupper())]              # drop initials like "K.S."
    if not 1 <= len(words) <= 3: return None
    a = " ".join(words)
    return a if len(a) >= 6 else None

def flush(doc, text):
    global mentions, unresolved
    if not doc or not text: return
    for m in CITE.finditer(text):
        mentions += 1
        tgt = resolve(m.group(0))
        if not tgt: unresolved += 1; continue
        if tgt == doc: continue
        if cluster_of.get(doc) and cluster_of.get(doc) == cluster_of.get(tgt): continue
        s, e = m.start(), m.end()
        ctx = re.sub(r"\s+", " ", text[max(0, s - 160):min(len(text), e + 60)]).strip()
        key = (doc, tgt)
        if key not in edges_best or len(ctx) > len(edges_best[key]):
            edges_best[key] = ctx
        nm = NAME_BEFORE.search(text[max(0, s - 90):s].strip())
        if nm:
            a = petitioner_alias(re.sub(r"\s+", " ", nm.group(1)))
            if a: alias_votes[tgt][a] += 1

def stream(handler):
    global cur_doc, buf
    cur_doc, buf = None, []
    with open(os.path.join(data_dir, "escr_chunks.jsonl"), encoding="utf-8") as f:
        for line in f:
            d = pat.search(line[:120]).group(1)
            if d != cur_doc:
                handler(cur_doc, " ".join(buf)); cur_doc, buf = d, []
            buf.append(json.loads(line)["text"])
    handler(cur_doc, " ".join(buf))

# pass 1: learn the SCC/AIR crosswalk from parallel citations
stream(lambda d, t: learn_crosswalk(t) if t else None)
_dy = {}
for d, m in meta.items():
    dt = m.get("date") or ""
    if dt[:4].isdigit(): _dy[d] = int(dt[:4])
settle_crosswalk(_dy)
print(f"[graph] crosswalk learned: {len(crosswalk)} SCC/AIR keys "
      f"(from {len(crosswalk_votes)} candidates), {time.time()-t0:.0f}s", flush=True)

# pass 2: extract edges + contexts + aliases with the full resolver
stream(flush)
print(f"[graph] body pass: {mentions} mentions, {len(edges_best)} unique edges, "
      f"{unresolved} unresolved ({unresolved/max(mentions,1):.0%}), {time.time()-t0:.0f}s", flush=True)

# ---------- merge legacy edges (keep their treatment labels) ----------
legacy_kept = 0
legacy = {}
for line in open(os.path.join(data_dir, "edges.jsonl"), encoding="utf-8"):
    e = json.loads(line)
    key = (e["from"], e["target"])
    legacy[key] = e
out = os.path.join(data_dir, "edges_v2.jsonl")
with open(out, "w", encoding="utf-8") as f:
    for key, ctx in edges_best.items():
        le = legacy.pop(key, None)
        treatment = (le or {}).get("treatment") or "cited"
        f.write(json.dumps({"from": key[0], "target": key[1], "treatment": treatment,
                            "method": "body", "para": ctx}, ensure_ascii=False) + "\n")
    for key, e in legacy.items():                       # headnote-only edges body regex missed
        legacy_kept += 1
        f.write(json.dumps({"from": key[0], "target": key[1],
                            "treatment": e.get("treatment") or "cited",
                            "method": "headnote", "para": e.get("para") or ""},
                           ensure_ascii=False) + "\n")
n_edges = len(edges_best) + legacy_kept
print(f"[graph] wrote {n_edges} edges ({len(edges_best)} body + {legacy_kept} headnote-only) "
      f"-> {out} | avg {n_edges/len(meta):.1f}/doc (was 3.6)", flush=True)

# ---------- aliases ----------
indeg = Counter()
for (f_, t_) in edges_best: indeg[t_] += 1
aliases, alias_votes_n = {}, {}
for tgt, votes in alias_votes.items():
    if indeg.get(tgt, 0) < 8: continue
    phrase, n = votes.most_common(1)[0]
    if n >= 4 and len(phrase) >= 6:
        key = phrase.lower()
        if key not in aliases or n > alias_votes_n.get(key, 0):
            aliases[key] = tgt; alias_votes_n[key] = n
with open(os.path.join(data_dir, "case_aliases.json"), "w", encoding="utf-8") as f:
    json.dump(aliases, f, ensure_ascii=False, indent=1)
print(f"[graph] aliases: {len(aliases)} -> case_aliases.json | {time.time()-t0:.0f}s total", flush=True)
for probe in ("kesavananda bharati", "maneka gandhi", "bachan singh"):
    hits = {k: v for k, v in aliases.items() if probe.split()[0] in k}
    if hits: print("   probe:", dict(list(hits.items())[:3]), flush=True)