#!/usr/bin/env python3
"""
MCRO Language Analytics Report (v3 — comparative + detailed, patched)
----------------------------------------------------------------------
Patch notes:
- Robustly coerces merge keys `search_group` and `search_term` to STRING dtype
  across *all* intermediate DataFrames (terms, instances, t, t_focus_terms, links, inst)
  to prevent dtype-mismatch errors like:
    ValueError: You are trying to merge on int64 and string columns for key 'search_group'.

Functionality is identical to the prior v3.
"""

import argparse
import json
from pathlib import Path
import sys
import pandas as pd

def as_list(x):
    if x is None:
        return []
    if isinstance(x, list):
        return x
    return [x]

def list_join(vals):
    vals = [v for v in vals if pd.notna(v)]
    return "; ".join(sorted(set(map(str, vals))))

def ensure_string_cols(df, cols):
    for c in cols:
        if c in df.columns:
            df[c] = df[c].astype("string")
    return df

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, help="Directory with JSON files")
    ap.add_argument("--outdir", default="report_language_v3", help="Output folder")
    ap.add_argument("--pattern", default="*.json", help="Glob for JSON files")
    ap.add_argument("--recurse", action="store_true", help="Recurse into subdirs")
    ap.add_argument("--focus_case", default="27-CR-23-1886", help="Focal case for comparative analysis")
    args = ap.parse_args()

    in_dir = Path(args.input)
    files = sorted(in_dir.rglob(args.pattern) if args.recurse else in_dir.glob(args.pattern))
    if not files:
        print(f"[warn] No JSON files matched {args.pattern} in {in_dir}")
        sys.exit(0)

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    # ---- Extract minimal doc, case, language tables ----
    docs_rows, terms_rows, inst_rows = [], [], []

    for fp in files:
        try:
            data = json.loads(fp.read_text(encoding="utf-8"))
        except Exception:
            continue

        doc_sha256 = data.get("doc_sha256") or data.get("sha256") or ""
        filename = data.get("filename", "")
        filing_type = data.get("filing_type", "")
        filing_date = data.get("filing_date", "")
        case = data.get("case") or {}
        case_id = case.get("case_id")
        cluster_id = case.get("cluster_id")
        cluster_name = case.get("cluster_name")

        docs_rows.append({
            "doc_sha256": doc_sha256,
            "filename": filename,
            "filing_type": filing_type,
            "filing_date": filing_date,
            "case_id": case_id,
            "cluster_id": cluster_id,
            "cluster_name": cluster_name,
        })

        lang = data.get("language") or {}
        for t in as_list(lang.get("terms")):
            terms_rows.append({
                "doc_sha256": doc_sha256,
                "search_group": t.get("search_group"),
                "search_term": t.get("search_term"),
                "quantity": t.get("quantity"),
            })
        for inst in as_list(lang.get("instances")):
            inst_rows.append({
                "doc_sha256": doc_sha256,
                "search_group": inst.get("search_group"),
                "search_term": inst.get("search_term"),
                "page_num": inst.get("page_num"),
                "num_pages": inst.get("num_pages"),
                "full_text_row": inst.get("full_text_row"),
                "search_term_url": inst.get("search_term_url"),
            })

    docs = pd.DataFrame(docs_rows)
    terms = pd.DataFrame(terms_rows)
    instances = pd.DataFrame(inst_rows)

    if terms.empty:
        (outdir/"README.txt").write_text("No language terms found.\n", encoding="utf-8")
        sys.exit(0)

    # ---- Coerce key columns to string early ----
    terms = ensure_string_cols(terms, ["search_group","search_term"])
    instances = ensure_string_cols(instances, ["search_group","search_term"])

    # Normalize numeric
    terms["quantity"] = pd.to_numeric(terms["quantity"], errors="coerce").fillna(0).astype(int)

    # Join doc context
    t = terms.merge(docs, on="doc_sha256", how="left")
    t = ensure_string_cols(t, ["search_group","search_term"])

    # Context tables
    terms_by_case = (
        t.groupby(["case_id","cluster_id","cluster_name","search_group","search_term"], dropna=False, as_index=False)
         .agg(total_quantity=("quantity","sum"),
              doc_count=("doc_sha256","nunique"),
              filenames=("filename", list_join))
    )
    terms_by_cluster = (
        t.groupby(["cluster_id","cluster_name","search_group","search_term"], dropna=False, as_index=False)
         .agg(total_quantity=("quantity","sum"),
              doc_count=("doc_sha256","nunique"),
              case_count=("case_id","nunique"),
              filenames=("filename", list_join))
    )

    # ----- FOCAL SETUP -----
    focus = args.focus_case
    t_focus = t[t["case_id"] == focus].copy()

    if t_focus.empty:
        (outdir/"README.txt").write_text(f"No terms found for focus case {focus}.\n", encoding="utf-8")
        sys.exit(0)

    focal_terms = (
        t_focus.groupby(["search_group","search_term"], as_index=False)
               .agg(in_case_quantity=("quantity","sum"),
                    in_case_doc_count=("doc_sha256","nunique"),
                    focal_filenames=("filename", list_join),
                    cluster_id=("cluster_id", lambda s: s.dropna().iloc[0] if len(s.dropna()) else pd.NA),
                    cluster_name=("cluster_name", lambda s: s.dropna().iloc[0] if len(s.dropna()) else pd.NA))
    )
    focal_terms = ensure_string_cols(focal_terms, ["search_group","search_term"])
    focal_terms["case_id"] = focus  # carry focus id

    # Compute others counts
    t_others = t[t["case_id"] != focus].copy()
    t_others = ensure_string_cols(t_others, ["search_group","search_term"])

    others_case_counts = (
        t_others.groupby(["search_group","search_term"], as_index=False)
                .agg(others_case_count=("case_id","nunique"))
    )
    others_cluster_counts = (
        t_others.groupby(["search_group","search_term"], as_index=False)
                .agg(others_cluster_count=("cluster_id","nunique"))
    )

    focal_cmp = (focal_terms
                 .merge(others_case_counts, on=["search_group","search_term"], how="left")
                 .merge(others_cluster_counts, on=["search_group","search_term"], how="left"))
    focal_cmp["others_case_count"] = focal_cmp["others_case_count"].fillna(0).astype(int)
    focal_cmp["others_cluster_count"] = focal_cmp["others_cluster_count"].fillna(0).astype(int)

    # Keep ONLY terms with at least one other case
    focal_cmp = focal_cmp[focal_cmp["others_case_count"] >= 1].copy()

    # Term-level weights
    focal_cmp["weight_cases"] = 1.0 / focal_cmp["others_case_count"]
    focal_cmp["weight_clusters"] = focal_cmp.apply(
        lambda r: (1.0 / r["others_cluster_count"]) if r["others_cluster_count"] >= 1 else 0.0, axis=1)

    # ----- DETAILED OTHER-CASE TABLES -----
    t_focus_terms = focal_cmp[["search_group","search_term"]].drop_duplicates()
    t_focus_terms = ensure_string_cols(t_focus_terms, ["search_group","search_term"])

    links = t_others.merge(t_focus_terms, on=["search_group","search_term"], how="inner")
    links = ensure_string_cols(links, ["search_group","search_term"])

    def agg_types(x):
        vals = [v for v in x if pd.notna(v)]
        return "; ".join(sorted(set(vals)))

    focal_links_by_case = (
        links.groupby(["search_group","search_term","case_id","cluster_id","cluster_name"], as_index=False)
             .agg(other_case_quantity=("quantity","sum"),
                  other_case_doc_count=("doc_sha256","nunique"),
                  other_filenames=("filename", list_join),
                  other_filing_types=("filing_type", agg_types),
                  first_seen_date=("filing_date", lambda s: s.dropna().min() if len(s.dropna()) else pd.NA),
                  last_seen_date=("filing_date", lambda s: s.dropna().max() if len(s.dropna()) else pd.NA))
    )

    focal_links_by_case_docs = (
        links.groupby(["search_group","search_term","doc_sha256","filename","filing_type","filing_date","case_id","cluster_id","cluster_name"], as_index=False)
             .agg(other_doc_quantity=("quantity","sum"))
    )

    if not instances.empty:
        inst = instances.merge(docs[["doc_sha256","filename","filing_type","filing_date","case_id","cluster_id","cluster_name"]],
                               on="doc_sha256", how="left")
        inst = ensure_string_cols(inst, ["search_group","search_term"])
        inst = inst.merge(t_focus_terms, on=["search_group","search_term"], how="inner")
        focal_links_by_case_instances = inst[inst["case_id"] != focus].copy()
        focal_links_by_case_instances = focal_links_by_case_instances[[
            "search_group","search_term","case_id","cluster_id","cluster_name","filename","page_num","num_pages","full_text_row","search_term_url"
        ]].sort_values(["search_group","search_term","case_id","filename","page_num"])
    else:
        focal_links_by_case_instances = pd.DataFrame(columns=[
            "search_group","search_term","case_id","cluster_id","cluster_name","filename","page_num","num_pages","full_text_row","search_term_url"
        ])

    focal_links_by_cluster = (
        links.groupby(["search_group","search_term","cluster_id","cluster_name"], as_index=False)
             .agg(other_cluster_quantity=("quantity","sum"),
                  other_cluster_doc_count=("doc_sha256","nunique"),
                  other_case_ids=("case_id", list_join))
    )

    # ----- MATRICES -----
    case_qty = focal_links_by_case.pivot_table(
        index=["search_group","search_term"],
        columns="case_id",
        values="other_case_quantity",
        fill_value=0,
        aggfunc="sum"
    ).reset_index()
    case_presence = case_qty.copy()
    case_presence.iloc[:, 2:] = (case_presence.iloc[:, 2:] > 0).astype(int)

    term_weights = focal_cmp.set_index(["search_group","search_term"])["weight_cases"]
    def weight_apply(df):
        core = df.set_index(["search_group","search_term"])
        mask = core > 0
        weights = pd.DataFrame(index=core.index, columns=core.columns, data=0.0)
        w = term_weights.reindex(core.index).fillna(0).astype(float)
        for col in core.columns:
            weights[col] = w.where(mask[col], 0.0)
        weights = weights.reset_index()
        return weights
    case_weight = weight_apply(case_qty)

    cluster_qty = focal_links_by_cluster.pivot_table(
        index=["search_group","search_term"],
        columns="cluster_id",
        values="other_cluster_quantity",
        fill_value=0,
        aggfunc="sum"
    ).reset_index()
    cluster_presence = cluster_qty.copy()
    cluster_presence.iloc[:, 2:] = (cluster_presence.iloc[:, 2:] > 0).astype(int)

    term_cluster_weights = focal_cmp.set_index(["search_group","search_term"])["weight_clusters"]
    def cluster_weight_apply(df):
        core = df.set_index(["search_group","search_term"])
        mask = core > 0
        weights = pd.DataFrame(index=core.index, columns=core.columns, data=0.0)
        w = term_cluster_weights.reindex(core.index).fillna(0).astype(float)
        for col in core.columns:
            weights[col] = w.where(mask[col], 0.0)
        weights = weights.reset_index()
        return weights
    cluster_weight = cluster_weight_apply(cluster_qty)

    # ----- SUMMARY & WRITE -----
    metrics = pd.DataFrame([{
        "focus_case": focus,
        "focus_terms_with_others": int(focal_cmp.shape[0]),
        "distinct_other_cases_linked": int(focal_links_by_case["case_id"].nunique()) if not focal_links_by_case.empty else 0,
        "distinct_other_clusters_linked": int(focal_links_by_cluster["cluster_id"].nunique()) if not focal_links_by_cluster.empty else 0,
        "matrix_cases_terms": int(case_qty.shape[0]),
        "matrix_cases_cols": int(case_qty.shape[1]-2 if case_qty.shape[1]>=2 else 0),
        "matrix_clusters_terms": int(cluster_qty.shape[0]),
        "matrix_clusters_cols": int(cluster_qty.shape[1]-2 if cluster_qty.shape[1]>=2 else 0),
    }])

    outdir.mkdir(parents=True, exist_ok=True)
    terms_by_case.to_csv(outdir/"terms_by_case.csv", index=False)
    terms_by_cluster.to_csv(outdir/"terms_by_cluster.csv", index=False)
    focal_cmp.to_csv(outdir/"focal_term_compare_summary.csv", index=False)
    focal_links_by_case.to_csv(outdir/"focal_term_links_by_case.csv", index=False)
    focal_links_by_case_docs.to_csv(outdir/"focal_term_links_by_case_docs.csv", index=False)
    focal_links_by_case_instances.to_csv(outdir/"focal_term_links_by_case_instances.csv", index=False)
    focal_links_by_cluster.to_csv(outdir/"focal_term_links_by_cluster.csv", index=False)
    case_qty.to_csv(outdir/"focal_term_by_case_matrix_quantity.csv", index=False)
    case_presence.to_csv(outdir/"focal_term_by_case_matrix_presence.csv", index=False)
    case_weight.to_csv(outdir/"focal_term_by_case_matrix_weight.csv", index=False)
    cluster_qty.to_csv(outdir/"focal_term_by_cluster_matrix_quantity.csv", index=False)
    cluster_presence.to_csv(outdir/"focal_term_by_cluster_matrix_presence.csv", index=False)
    cluster_weight.to_csv(outdir/"focal_term_by_cluster_matrix_weight.csv", index=False)
    metrics.to_csv(outdir/"overview_metrics.csv", index=False)

    readme = f"""
    MCRO Language Analytics (v3 — comparative + detailed, patched)
    =============================================================
    Focal case: {args.focus_case}

    This report excludes focal-only terms, and provides:
      - Detailed other-case tables (case rollups, doc granularity, page-level instances)
      - Wide matrices (terms x other cases/clusters) with quantity, presence, and weights

    Type coercion patch:
      - All merge/pivot keys (`search_group`, `search_term`) are coerced to STRING
        to avoid dtype mismatch errors observed in some corpora.
    """.strip() + "\n"
    (outdir/"README.txt").write_text(readme, encoding="utf-8")

    print(f"[ok] Wrote v3 comparative language reports to: {outdir.resolve()}")

if __name__ == "__main__":
    main()