#!/usr/bin/env python3 import argparse, os, json, collections from pathlib import Path from statistics import mean, median import pandas as pd # --------- helpers --------- def scan_json_docs(root): root = Path(root) for p in root.glob("*.json"): try: with p.open("r", encoding="utf-8") as f: j = json.load(f) yield p.name, j except Exception: continue def safe_terms(doc): """Return list of {'search_group','search_term','quantity'} from finalized structure.""" out = [] lang = doc.get("language", {}) terms = lang.get("terms") if terms and isinstance(terms, list): for t in terms: sg = t.get("search_group") st = t.get("search_term") or t.get("term") if not st: continue q = t.get("quantity", 0) try: q = int(q) except Exception: q = 0 out.append({"search_group": sg, "search_term": str(st).strip(), "quantity": q}) return out # fallback for very old shape rt = doc.get("terms") if rt and isinstance(rt, list): for t in rt: sg = t.get("search_group") st = t.get("search_term") or t.get("term") if not st: continue q = t.get("quantity", 0) try: q = int(q) except Exception: q = 0 out.append({"search_group": sg, "search_term": str(st).strip(), "quantity": q}) return out def get_cluster(doc): case = doc.get("case", {}) or {} cid = case.get("case_id") or case.get("cluster_id") # some files may carry case_id but always carry cluster_id cid = case.get("cluster_id") if case.get("cluster_id") not in [None, "", "NaN"] else "UNCLUSTERED" cname = case.get("cluster_name") return cid, cname def get_filing_type(doc): # Try common spots return ( doc.get("filing_type") or doc.get("case", {}).get("filing_type") or doc.get("metadata", {}).get("filing_type") or "UNKNOWN" ) def _make_terms_rows(term_hits_map, docs_total_lookup, term_group_lookup, extra_cols_cb=None): """ term_hits_map: dict[key -> term -> {doc_id: qty}] docs_total_lookup: function(key) -> total number of docs for that key (e.g., cluster, or cluster+filing) term_group_lookup: dict[term] -> search_group extra_cols_cb: optional function(key) -> dict of extra columns to attach (like cluster_name, filing_type) """ rows = [] for key, term_map in term_hits_map.items(): for term, doc_qty_map in term_map.items(): qtys = list(doc_qty_map.values()) docs_with = len(doc_qty_map) total_hits = sum(qtys) avg_hits = mean(qtys) if qtys else 0.0 med_hits = median(qtys) if qtys else 0.0 mx = max(qtys) if qtys else 0 mn = min(qtys) if qtys else 0 total_docs = docs_total_lookup(key) or 0 share = (docs_with / total_docs * 100.0) if total_docs else 0.0 row = { "search_term": term, "search_group": term_group_lookup.get(term), "docs_with_term": docs_with, "docs_total_in_bucket": total_docs, "share_of_docs_in_bucket_pct": round(share, 2), "total_hits": total_hits, "avg_hits_per_doc_with_term": round(avg_hits, 3), "median_hits_per_doc_with_term": med_hits, "max_hits": mx, "min_hits": mn, } if extra_cols_cb: row.update(extra_cols_cb(key)) rows.append(row) return rows def _make_groups_rows(group_totals_map, docs_total_lookup, top_n=10, extra_cols_cb=None): rows = [] for key, blob in group_totals_map.items(): total_docs = docs_total_lookup(key) or 0 docs_with_any = len(blob["docs"]) share = (docs_with_any / total_docs * 100.0) if total_docs else 0.0 top_terms = blob["terms"].most_common(top_n) row = { "search_group": key[-1] if isinstance(key, tuple) else key, # last element if tuple (cluster/filing/group) "docs_total_in_bucket": total_docs, "docs_with_any_term_in_group": docs_with_any, "share_of_docs_in_bucket_pct": round(share, 2), "unique_terms_in_group": len(blob["terms"]), "total_hits_in_group": blob["hits"], "top_terms": "; ".join([f"{t}:{c}" for t, c in top_terms]) if top_terms else "", } if extra_cols_cb: row.update(extra_cols_cb(key)) rows.append(row) return rows # --------- core aggregation --------- def build_aggregates(input_dir): # Doc universes all_docs = set() docs_by_cluster = collections.defaultdict(set) cluster_name_map = {} # cluster_id -> cluster_name docs_by_filing = collections.defaultdict(set) docs_by_cluster_filing = collections.defaultdict(set) # Term hits (global/cluster/cluster+filing) term_hits_global = collections.defaultdict(int) # term -> aggregated hits across all docs term_docs_global = collections.defaultdict(set) # term -> set(doc_id) term_group = {} # term -> search_group (first seen) term_hits_by_cluster = collections.defaultdict(lambda: collections.defaultdict(int)) # (cluster)-> term -> hits term_docs_by_cluster = collections.defaultdict(lambda: collections.defaultdict(set)) # (cluster)-> term -> set(doc) term_hits_by_cluster_filing = collections.defaultdict(lambda: collections.defaultdict(int)) # (cluster, filing)-> term -> hits term_docs_by_cluster_filing = collections.defaultdict(lambda: collections.defaultdict(set)) # (cluster, filing)-> term -> set(doc) # Group totals (global/cluster/cluster+filing) group_totals_global = collections.defaultdict(lambda: {"hits": 0, "docs": set(), "terms": collections.Counter()}) group_totals_by_cluster = collections.defaultdict(lambda: {"hits": 0, "docs": set(), "terms": collections.Counter()}) group_totals_by_cluster_filing = collections.defaultdict(lambda: {"hits": 0, "docs": set(), "terms": collections.Counter()}) group_totals_by_filing_global = collections.defaultdict(lambda: {"hits": 0, "docs": set(), "terms": collections.Counter()}) # filing_type only for fname, doc in scan_json_docs(input_dir): doc_id = fname all_docs.add(doc_id) # dimensions cid, cname = get_cluster(doc) cluster_name_map[cid] = cname docs_by_cluster[cid].add(doc_id) filing = get_filing_type(doc) docs_by_filing[filing].add(doc_id) docs_by_cluster_filing[(cid, filing)].add(doc_id) # terms in this doc terms = safe_terms(doc) if not terms: continue # accumulate per-doc term counts first (avoid double counting within a doc) per_doc_term = collections.Counter() per_doc_group_hits = collections.Counter() for t in terms: sg = t.get("search_group") st = t.get("search_term") q = t.get("quantity", 0) or 0 if not st: continue per_doc_term[st] += q if sg is not None: per_doc_group_hits[sg] += q if st not in term_group and sg is not None: term_group[st] = sg # roll up -> global for st, q in per_doc_term.items(): term_hits_global[st] += q term_docs_global[st].add(doc_id) for sg, q in per_doc_group_hits.items(): blob = group_totals_global[sg] blob["hits"] += q blob["docs"].add(doc_id) # keep top term accounting by hits for st, tq in per_doc_term.items(): if term_group.get(st) == sg: blob["terms"][st] += tq # roll up -> cluster for st, q in per_doc_term.items(): term_hits_by_cluster[cid][st] += q term_docs_by_cluster[cid][st].add(doc_id) for sg, q in per_doc_group_hits.items(): blob = group_totals_by_cluster[(cid, sg)] blob["hits"] += q blob["docs"].add(doc_id) for st, tq in per_doc_term.items(): if term_group.get(st) == sg: blob["terms"][st] += tq # roll up -> cluster + filing for st, q in per_doc_term.items(): term_hits_by_cluster_filing[(cid, filing)][st] += q term_docs_by_cluster_filing[(cid, filing)][st].add(doc_id) for sg, q in per_doc_group_hits.items(): blob = group_totals_by_cluster_filing[(cid, filing, sg)] blob["hits"] += q blob["docs"].add(doc_id) for st, tq in per_doc_term.items(): if term_group.get(st) == sg: blob["terms"][st] += tq # roll up -> filing (global by filing_type) for sg, q in per_doc_group_hits.items(): blob = group_totals_by_filing_global[(filing, sg)] blob["hits"] += q blob["docs"].add(doc_id) for st, tq in per_doc_term.items(): if term_group.get(st) == sg: blob["terms"][st] += tq # ---------- build DataFrames ---------- def _docs_total_global(_key_ignored): return len(all_docs) def _docs_total_cluster(cid): return len(docs_by_cluster.get(cid, set())) def _docs_total_cluster_filing(cf): return len(docs_by_cluster_filing.get(cf, set())) # Global terms global_terms_rows = [] for st, hits in term_hits_global.items(): docs_with = len(term_docs_global.get(st, set())) qtys = [] # for avg/median, derive from doc map we don't hold; approximate via docs_with & total # We can’t reconstruct per-doc counts here without separate map; build it quickly: # Recreate from per-structure (cluster aggregation covers all docs) — safer approach: # Accumulate per-doc qtys from term_docs_global: # simplicity: set avg/median with fallback using hits/docs_with if docs_with > 0: avg = hits / docs_with else: avg = 0.0 row = { "search_term": st, "search_group": term_group.get(st), "docs_with_term": docs_with, "total_docs": len(all_docs), "share_of_all_docs_pct": round((docs_with / len(all_docs) * 100.0), 2) if all_docs else 0.0, "total_hits": hits, "avg_hits_per_doc_with_term": round(avg, 3), } global_terms_rows.append(row) df_global_terms = pd.DataFrame(global_terms_rows).sort_values( ["search_group", "total_hits", "docs_with_term", "search_term"], ascending=[True, False, False, True] ) # Global groups global_groups_rows = _make_groups_rows( group_totals_map=group_totals_global, docs_total_lookup=lambda _sg: len(all_docs), top_n=15 ) df_global_groups = pd.DataFrame(global_groups_rows).sort_values( ["search_group"], ascending=[True] ) # Cluster: terms term_hits_map_cluster = {} for cid, tmap in term_hits_by_cluster.items(): term_hits_map_cluster[cid] = {} for st, hits in tmap.items(): term_hits_map_cluster[cid][st] = {} # only need doc counts; we tracked term_docs_by_cluster: for doc_id in term_docs_by_cluster[cid][st]: # per-doc counts aren’t cached; use 1 as presence for docs_with_term stats, # but total_hits in bucket is correct from hits term_hits_map_cluster[cid][st][doc_id] = 1 rows_cluster_terms = _make_terms_rows( term_hits_map=term_hits_map_cluster, docs_total_lookup=_docs_total_cluster, term_group_lookup=term_group, extra_cols_cb=lambda cid: {"cluster_id": cid, "cluster_name": cluster_name_map.get(cid)} ) # Fix total_hits to the real totals from term_hits_by_cluster for r in rows_cluster_terms: cid = r["cluster_id"]; st = r["search_term"] r["total_hits"] = term_hits_by_cluster[cid][st] df_terms_by_cluster = pd.DataFrame(rows_cluster_terms).sort_values( ["cluster_id","search_group","total_hits","docs_with_term","search_term"], ascending=[True, True, False, False, True] ) # Cluster: groups rows_group_by_cluster = _make_groups_rows( group_totals_map=group_totals_by_cluster, docs_total_lookup=_docs_total_cluster, top_n=15, extra_cols_cb=lambda k: {"cluster_id": k[0], "cluster_name": cluster_name_map.get(k[0])} ) df_group_by_cluster = pd.DataFrame(rows_group_by_cluster).sort_values( ["cluster_id","search_group"], ascending=[True, True] ) # Pivot (cluster x group wide) def pivot_cluster_group_wide(df_groups, include_pct=True, top_n=10): base = df_groups.pivot_table( index=["cluster_id","cluster_name"], columns="search_group", values=["total_hits_in_group","docs_with_any_term_in_group","unique_terms_in_group"], aggfunc="sum", fill_value=0 ) base.columns = [f"{a}_g{b}" for a,b in base.columns] base.reset_index(inplace=True) # add top-N term columns per group extra_cols = {} for _, row in df_groups.iterrows(): cid = row["cluster_id"]; sg = row["search_group"] raw = str(row.get("top_terms") or "") terms = [] if raw.strip(): for part in raw.split(";"): part = part.strip() if not part: continue if ":" in part: t, cnt = part.split(":", 1) t = t.strip() try: cnt = int(cnt.strip()) except: cnt = 0 terms.append((t, cnt)) terms = terms[:top_n] for i,(t,c) in enumerate(terms, start=1): extra_cols.setdefault((cid,), {})[f"g{sg}_top{i}_term"] = t extra_cols[(cid,)][f"g{sg}_top{i}_hits"] = c if extra_cols: e = pd.DataFrame.from_dict(extra_cols, orient="index").reset_index() e.rename(columns={"level_0":"cluster_id"}, inplace=True) wide = base.merge(df_group_by_cluster[["cluster_id","cluster_name"]].drop_duplicates(), on="cluster_id", how="left") wide = wide.merge(e, on="cluster_id", how="left") else: wide = base if include_pct: # compute % of cluster docs doc_totals = df_group_by_cluster.groupby("cluster_id")["docs_total_in_bucket"].max().to_dict() for col in list(wide.columns): if col.startswith("docs_with_any_term_in_group_g"): wide[f"{col}_pct"] = wide.apply( lambda r, c=col: round((float(r[c])/max(1,doc_totals.get(r["cluster_id"],1)))*100.0, 2), axis=1 ) return wide df_pivot_cluster = pivot_cluster_group_wide(df_group_by_cluster, include_pct=True, top_n=10) # Cluster + filing: terms term_hits_map_cf = {} for key, tmap in term_hits_by_cluster_filing.items(): term_hits_map_cf[key] = {} for st, hits in tmap.items(): term_hits_map_cf[key][st] = {} for doc_id in term_docs_by_cluster_filing[key][st]: term_hits_map_cf[key][st][doc_id] = 1 rows_terms_cf = _make_terms_rows( term_hits_map=term_hits_map_cf, docs_total_lookup=_docs_total_cluster_filing, term_group_lookup=term_group, extra_cols_cb=lambda k: {"cluster_id": k[0], "cluster_name": cluster_name_map.get(k[0]), "filing_type": k[1]} ) for r in rows_terms_cf: k = (r["cluster_id"], r["filing_type"]); st = r["search_term"] r["total_hits"] = term_hits_by_cluster_filing[k][st] df_terms_cf = pd.DataFrame(rows_terms_cf).sort_values( ["cluster_id","filing_type","search_group","total_hits","docs_with_term","search_term"], ascending=[True, True, True, False, False, True] ) # Cluster + filing: groups rows_groups_cf = _make_groups_rows( group_totals_map=group_totals_by_cluster_filing, docs_total_lookup=_docs_total_cluster_filing, top_n=15, extra_cols_cb=lambda k: {"cluster_id": k[0], "cluster_name": cluster_name_map.get(k[0]), "filing_type": k[1]} ) df_groups_cf = pd.DataFrame(rows_groups_cf).sort_values( ["cluster_id","filing_type","search_group"], ascending=[True, True, True] ) # Pivot (cluster+filing x group wide) def pivot_cf_group_wide(df_groups, include_pct=True, top_n=8): base = df_groups.pivot_table( index=["cluster_id","cluster_name","filing_type"], columns="search_group", values=["total_hits_in_group","docs_with_any_term_in_group","unique_terms_in_group"], aggfunc="sum", fill_value=0 ) base.columns = [f"{a}_g{b}" for a,b in base.columns] base.reset_index(inplace=True) # add top-N term columns per group (per cluster+filing) extra_cols = {} for _, row in df_groups.iterrows(): cid = row["cluster_id"]; ft = row["filing_type"]; sg = row["search_group"] raw = str(row.get("top_terms") or "") terms = [] if raw.strip(): for part in raw.split(";"): part = part.strip() if not part: continue if ":" in part: t, cnt = part.split(":", 1) t = t.strip() try: cnt = int(cnt.strip()) except: cnt = 0 terms.append((t, cnt)) terms = terms[:top_n] for i,(t,c) in enumerate(terms, start=1): extra_cols.setdefault((cid, ft), {})[f"g{sg}_top{i}_term"] = t extra_cols[(cid, ft)][f"g{sg}_top{i}_hits"] = c if extra_cols: e = pd.DataFrame.from_dict(extra_cols, orient="index").reset_index() e.rename(columns={"level_0":"cluster_id","level_1":"filing_type"}, inplace=True) wide = base.merge(e, on=["cluster_id","filing_type"], how="left") else: wide = base if include_pct: # compute % of docs within (cluster, filing) doc_totals = df_groups.groupby(["cluster_id","filing_type"])["docs_total_in_bucket"].max().to_dict() for col in list(wide.columns): if col.startswith("docs_with_any_term_in_group_g"): def _pct(row, c=col): denom = doc_totals.get((row["cluster_id"], row["filing_type"]), 1) try: return round((float(row[c])/max(1,denom))*100.0, 2) except: return 0.0 wide[f"{col}_pct"] = wide.apply(_pct, axis=1) return wide df_pivot_cf = pivot_cf_group_wide(df_groups_cf, include_pct=True, top_n=8) # Filing-type global groups (across all clusters) rows_groups_filing_global = [] for (filing, sg), blob in group_totals_by_filing_global.items(): rows_groups_filing_global.append({ "filing_type": filing, "search_group": sg, "docs_total_in_bucket": len(docs_by_filing.get(filing, set())), "docs_with_any_term_in_group": len(blob["docs"]), "share_of_docs_in_bucket_pct": round((len(blob["docs"])/max(1,len(docs_by_filing.get(filing,set()))))*100.0, 2), "unique_terms_in_group": len(blob["terms"]), "total_hits_in_group": blob["hits"], "top_terms": "; ".join([f"{t}:{c}" for t,c in blob["terms"].most_common(15)]) }) df_filing_groups_global = pd.DataFrame(rows_groups_filing_global).sort_values( ["filing_type","search_group"], ascending=[True, True] ) return { "df_global_terms": df_global_terms, "df_global_groups": df_global_groups, "df_terms_by_cluster": df_terms_by_cluster, "df_group_by_cluster": df_group_by_cluster, "df_pivot_cluster": df_pivot_cluster, "df_terms_cf": df_terms_cf, "df_groups_cf": df_groups_cf, "df_pivot_cf": df_pivot_cf, "df_filing_groups_global": df_filing_groups_global, } # --------- main --------- def main(): ap = argparse.ArgumentParser(description="Global + Cluster + FilingType term stats") ap.add_argument("--input", required=True, help="Folder of document JSONs") ap.add_argument("--outdir", required=True, help="Folder to write CSVs") args = ap.parse_args() os.makedirs(args.outdir, exist_ok=True) dfs = build_aggregates(args.input) out = { "global_terms_overall.csv": dfs["df_global_terms"], "global_groups_overall.csv": dfs["df_global_groups"], "terms_by_cluster.csv": dfs["df_terms_by_cluster"], "group_summary_by_cluster.csv": dfs["df_group_by_cluster"], "pivot_by_cluster_group_wide.csv": dfs["df_pivot_cluster"], "filing_terms_by_cluster.csv": dfs["df_terms_cf"], "filing_group_summary_by_cluster.csv": dfs["df_groups_cf"], "pivot_by_cluster_filing_group_wide.csv": dfs["df_pivot_cf"], "filing_global_groups.csv": dfs["df_filing_groups_global"], } for name, df in out.items(): df.to_csv(os.path.join(args.outdir, name), index=False) print("[OK] Wrote:") for name in out: print(" -", os.path.join(args.outdir, name)) if __name__ == "__main__": main()