#!/usr/bin/env python3
import argparse, os, json, math, collections
from pathlib import Path
from statistics import mean, median
import pandas as pd

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:
            # skip unreadable/bad jsons
            continue

def safe_terms(doc):
    """
    Return list of {'search_group', 'search_term', 'quantity'} from the finalized structure.
    Primary: doc['language']['terms']; Fallback: doc['terms'] if present in some files.
    """
    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")
            qty = t.get("quantity", 0)
            if st is None:
                continue
            try:
                qty = int(qty)
            except Exception:
                qty = 0
            out.append({"search_group": sg, "search_term": st, "quantity": qty})
        return out

    # Fallback if some docs still have root 'terms' in consolidated format
    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")
            qty = t.get("quantity", 0)
            if st is None:
                continue
            try:
                qty = int(qty)
            except Exception:
                qty = 0
            out.append({"search_group": sg, "search_term": st, "quantity": qty})
    return out

def get_cluster(doc):
    case = doc.get("case", {}) or {}
    cid = case.get("cluster_id")
    cname = case.get("cluster_name")
    # normalize empties
    cid = cid if (cid not in [None, "", "NaN"]) else "UNCLUSTERED"
    return cid, cname

def build_aggregates(input_dir):
    # Per cluster: set of docs
    cluster_docs = collections.defaultdict(set)
    # Per (cluster, term): per-doc quantities list
    term_hits = collections.defaultdict(lambda: collections.defaultdict(int))
    # Per (cluster, group): accumulate totals
    group_totals = collections.defaultdict(lambda: {"hits": 0, "docs": set(), "terms": collections.Counter()})
    # We’ll also track term -> group to help labeling
    term_group = {}

    for fname, doc in scan_json_docs(input_dir):
        cid, cname = get_cluster(doc)
        cluster_docs[cid].add(fname)

        terms = safe_terms(doc)
        if not terms:
            continue

        # by term in this doc
        per_doc_term = collections.Counter()
        for t in terms:
            sg = t.get("search_group")
            st = str(t.get("search_term")).strip()
            qty = int(t.get("quantity", 0)) if t.get("quantity") is not None else 0
            if not st:
                continue
            per_doc_term[st] += qty
            term_group[st] = sg

        # roll into structures
        for st, qty in per_doc_term.items():
            term_hits[(cid, st)][fname] = qty
            sg = term_group.get(st)
            gt = group_totals[(cid, sg)]
            gt["hits"] += qty
            gt["docs"].add(fname)
            gt["terms"][st] += qty

    return cluster_docs, term_hits, group_totals, term_group

def make_terms_by_cluster_csv(cluster_docs, term_hits, term_group, outpath):
    rows = []
    for (cid, st), doc_qty_map in term_hits.items():
        docs_with = len(doc_qty_map)
        qtys = list(doc_qty_map.values())
        total_hits = sum(qtys)
        try:
            avg_hits = mean(qtys)
        except Exception:
            avg_hits = float(qtys[0]) if qtys else 0.0
        try:
            med_hits = median(qtys)
        except Exception:
            med_hits = float(qtys[0]) if qtys else 0.0
        max_hits = max(qtys) if qtys else 0
        min_hits = min(qtys) if qtys else 0
        cluster_total_docs = len(cluster_docs.get(cid, [])) or 0
        share = (docs_with / cluster_total_docs * 100.0) if cluster_total_docs else 0.0
        rows.append({
            "cluster_id": cid,
            "search_term": st,
            "search_group": term_group.get(st),
            "docs_with_term": docs_with,
            "cluster_docs_total": cluster_total_docs,
            "share_of_cluster_docs_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": max_hits,
            "min_hits": min_hits,
        })
    df = pd.DataFrame(rows)
    df.sort_values(["cluster_id","search_group","total_hits","docs_with_term","search_term"],
                   ascending=[True, True, False, False, True], inplace=True)
    df.to_csv(outpath, index=False)
    return df

def make_group_summary_by_cluster_csv(cluster_docs, group_totals, outpath, top_n=10):
    rows = []
    for (cid, sg), blob in group_totals.items():
        cluster_total_docs = len(cluster_docs.get(cid, [])) or 0
        docs_with_any = len(blob["docs"])
        share = (docs_with_any / cluster_total_docs * 100.0) if cluster_total_docs else 0.0
        # top terms by hits within this cluster+group
        top_terms = blob["terms"].most_common(top_n)
        rows.append({
            "cluster_id": cid,
            "search_group": sg,
            "cluster_docs_total": cluster_total_docs,
            "docs_with_any_term_in_group": docs_with_any,
            "share_of_cluster_docs_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 "",
        })
    df = pd.DataFrame(rows).sort_values(
        ["cluster_id","search_group"], ascending=[True, True]
    )
    df.to_csv(outpath, index=False)
    return df

def make_pivot_wide(df_groups, outpath, wide_top=10, include_pct=True):
    # Base pivot: totals per cluster/group
    base = df_groups.pivot_table(
        index="cluster_id",
        columns="search_group",
        values=["total_hits_in_group","docs_with_any_term_in_group","unique_terms_in_group"],
        aggfunc="sum",
        fill_value=0
    )
    # Flatten columns
    base.columns = [f"{a}_g{b}" for a,b in base.columns]
    base.reset_index(inplace=True)

    # Optionally add top-N terms columns per group
    # Parse df_groups['top_terms'] (already cid+sg rows)
    extra_cols = {}
    if wide_top and wide_top > 0 and "top_terms" in df_groups.columns:
        for _, row in df_groups.iterrows():
            cid = row["cluster_id"]; sg = row["search_group"]
            top_terms = []
            raw = str(row.get("top_terms") or "")
            if raw.strip():
                for part in raw.split(";"):
                    part = part.strip()
                    if not part:
                        continue
                    if ":" in part:
                        term, cnt = part.split(":", 1)
                        term = term.strip()
                        try:
                            cnt = int(cnt.strip())
                        except Exception:
                            cnt = 0
                        top_terms.append((term, cnt))
            # keep only top_n
            top_terms = top_terms[:wide_top]
            for i, (term, cnt) in enumerate(top_terms, start=1):
                key_t = f"g{sg}_top{i}_term"
                key_c = f"g{sg}_top{i}_hits"
                extra_cols.setdefault(cid, {})[key_t] = term
                extra_cols[cid][key_c] = cnt

    if extra_cols:
        extra_df = pd.DataFrame.from_dict(extra_cols, orient="index").reset_index().rename(columns={"index":"cluster_id"})
        wide = base.merge(extra_df, on="cluster_id", how="left")
    else:
        wide = base

    if include_pct:
        # add overall doc totals per cluster from groups table (max across groups)
        doc_totals = df_groups.groupby("cluster_id")["cluster_docs_total"].max().to_dict()
        for col in list(wide.columns):
            if col.startswith("docs_with_any_term_in_group_g"):
                sg = col.split("_g",1)[1]
                pct_col = f"{col}_pct"
                wide[pct_col] = wide[col].apply(
                    lambda v, cid=None: round((v / doc_totals.get(cid, 1) * 100.0), 2),
                    cid=None
                )
        # Recompute using apply with access to row:
        def _row_pct(row, col):
            cid = row["cluster_id"]
            denom = doc_totals.get(cid, 1)
            v = row[col]
            try:
                return round((float(v)/denom)*100.0, 2)
            except Exception:
                return 0.0
        pct_cols = []
        for col in wide.columns:
            if col.startswith("docs_with_any_term_in_group_g"):
                pct_name = f"{col}_pct"
                wide[pct_name] = wide.apply(lambda r, c=col: _row_pct(r, c), axis=1)
                pct_cols.append(pct_name)

    wide.to_csv(outpath, index=False)
    return wide

def main():
    ap = argparse.ArgumentParser(description="Analyze language terms by cluster")
    ap.add_argument("--input", required=True, help="Folder of document JSON")
    ap.add_argument("--outdir", required=True, help="Output directory for CSVs")
    ap.add_argument("--wide-top", type=int, default=10, help="Top-N terms per group to include in wide pivot")
    ap.add_argument("--wide-pct", action="store_true", help="Include % columns for docs_with_any_term_in_group")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    cluster_docs, term_hits, group_totals, term_group = build_aggregates(args.input)

    out_terms = os.path.join(args.outdir, "terms_by_cluster.csv")
    out_groups = os.path.join(args.outdir, "group_summary_by_cluster.csv")
    out_pivot = os.path.join(args.outdir, "pivot_by_cluster_group_wide.csv")

    df_terms = make_terms_by_cluster_csv(cluster_docs, term_hits, term_group, out_terms)
    df_groups = make_group_summary_by_cluster_csv(cluster_docs, group_totals, out_groups, top_n=args.wide_top)
    make_pivot_wide(df_groups, out_pivot, wide_top=args.wide_top, include_pct=args.wide_pct)

    print("[OK] Wrote:\n  -", out_terms, "\n  -", out_groups, "\n  -", out_pivot)

if __name__ == "__main__":
    main()