#!/usr/bin/env python3
# (See header docstring for details.)
import csv, json, re, sys
from pathlib import Path
from collections import defaultdict, Counter
from typing import Dict, Tuple, List, Any

SUFFIX_RE = re.compile(
    r",\s*(?:jr\.?|sr\.?|ii|iii|iv|v|vi|vii|viii|ix|x|[A-Za-z]{2}\.?)$",
    re.IGNORECASE
)

def normalize_defendant(raw: str, lowercase_key: bool = True, strict_two_token: bool = False):
    s = (raw or "").strip()
    s = SUFFIX_RE.sub("", s)
    s = re.sub(r"\s+", " ", s)
    toks = s.split() if s else []
    toks_key = [t.lower() for t in toks] if lowercase_key else toks
    if len(toks_key) >= 2:
        key = (toks_key[0], toks_key[-1])
    elif len(toks_key) == 1:
        key = (toks_key[0],) if not strict_two_token else tuple()
    else:
        key = tuple()
    return s, key

def read_csv(path: Path):
    with path.open("r", encoding="utf-8-sig", newline="") as f:
        r = csv.DictReader(f)
        headers = r.fieldnames or []
        rows = [dict(x) for x in r]
    return headers, rows

def write_csv(path: Path, headers: List[str], rows: List[Dict[str, Any]]):
    with path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=headers, extrasaction="ignore")
        w.writeheader()
        for row in rows:
            w.writerow(row)

def main(argv: List[str]):
    import argparse
    ap = argparse.ArgumentParser(description="Cluster defendants by normalized name and rank by unique case count.")
    ap.add_argument("input_csv")
    ap.add_argument("output_csv")
    ap.add_argument("--json", dest="json_out")
    ap.add_argument("--defendant-col", default="Defendant")
    ap.add_argument("--case-col", default="Case Number")
    ap.add_argument("--index-col", default=None)
    ap.add_argument("--keep-blank", action="store_true")
    ap.add_argument("--lowercase-key", action="store_true", default=True)
    ap.add_argument("--strict-two-token", action="store_true")
    args = ap.parse_args(argv[1:])

    in_path = Path(args.input_csv); out_path = Path(args.output_csv)
    headers, rows = read_csv(in_path)
    if args.defendant_col not in headers:
        sys.exit(f'Missing defendant column: "{args.defendant_col}"')
    if args.case_col not in headers:
        sys.exit(f'Missing case column: "{args.case_col}"')

    cluster_data: Dict[Tuple[str, ...], Dict[str, Any]] = defaultdict(lambda: {
        "variants": Counter(),
        "case_numbers": set(),
        "row_indices": []
    })

    for i, r in enumerate(rows):
        raw_def = (r.get(args.defendant_col, "") or "").strip()
        case_no = (r.get(args.case_col, "") or "").strip()
        cleaned, key = normalize_defendant(raw_def, lowercase_key=args.lowercase_key, strict_two_token=args.strict_two_token)
        if not key and not args.keep_blank:
            continue
        cluster_data[key]["variants"][cleaned] += 1
        if case_no:
            cluster_data[key]["case_numbers"].add(case_no)
        cluster_data[key]["row_indices"].append(i)

    def key_to_str(k: Tuple[str, ...]) -> str:
        return " ".join(k)

    ranking = sorted(
        cluster_data.items(),
        key=lambda kv: (-len(kv[1]["case_numbers"]), key_to_str(kv[0]))
    )

    cluster_id_by_key = {k: rank for rank, (k, _) in enumerate(ranking, start=1)}

    out_headers = list(headers)
    for col in ["ClusterID", "ClusterKey", "ClusterSize", "VariantCount", "CaseCount"]:
        if col not in out_headers:
            out_headers.append(col)

    for k, data in ranking:
        cid = cluster_id_by_key[k]
        cluster_size = len(data["row_indices"])
        case_count  = len(data["case_numbers"])
        variant_cnt = len(data["variants"])
        key_str = key_to_str(k)
        for idx in data["row_indices"]:
            rows[idx]["ClusterID"]    = str(cid)
            rows[idx]["ClusterKey"]   = key_str
            rows[idx]["ClusterSize"]  = str(cluster_size)
            rows[idx]["VariantCount"] = str(variant_cnt)
            rows[idx]["CaseCount"]    = str(case_count)

    write_csv(out_path, out_headers, rows)

    if args.json_out:
        clusters_list = []
        for k, data in ranking:
            cid = cluster_id_by_key[k]
            clusters_list.append({
                "cluster_id": cid,
                "cluster_key": key_to_str(k),
                "case_count": len(data["case_numbers"]),
                "row_count": len(data["row_indices"]),
                "variant_count": len(data["variants"]),
                "variants": [{"name": n, "count": c} for n, c in sorted(data["variants"].items(), key=lambda x: (-x[1], x[0]))],
                "case_numbers": sorted(data["case_numbers"])
            })
        payload = {
            "source_csv": in_path.name,
            "defendant_col": args.defendant_col,
            "case_col": args.case_col,
            "cluster_count": len(clusters_list),
            "clusters": clusters_list
        }
        Path(args.json_out).write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"Wrote clusters JSON: {args.json_out}")

    print(f"Wrote augmented CSV: {out_path}  (clusters: {len(cluster_data)})")

if __name__ == "__main__":
    main(sys.argv)