#!/usr/bin/env python3
"""Case-control matching to remove the attacker-material/development
confound found in the naive full-population regression: attackerarmyscaler
alone got AUC 0.99 distinguishing mate-precursor positions from
draw-derived controls, causing quasi-complete separation that destabilizes
every other coefficient in the joint fit. Fix: match each positive to the
nearest available control by (attackerarmyscaler, minorsathome) instead of
using controls sampled freely from the draw pool. See CLAUDE.md's
king-safety discussion for the full methodology.

Uses exact-bucket + expanding-radius matching, not a KDTree: the match
columns are low-cardinality integers, so exact-value clusters are huge and
a KDTree's k nearest neighbors are highly tied -- every positive in a big
cluster sees almost the same "nearest k" candidates, so the first few
positions processed exhaust that shared candidate set and everyone else in
the cluster fails to match (confirmed: a k=50 KDTree approach matched only
308 of 43181 positives). Bucketing by exact integer key, popping
candidates on use, and expanding search radius only when a bucket is
empty has no such collision problem and is O(1) amortized per positive.

--exact-cols vs --radius-cols, NOT a single shared radius across all match
columns: a first version used one Chebyshev radius across both
attackerarmyscaler (range 0-31) and minorsathome (range 0-4) -- "radius <=
2" is tight for army-scaler (~6% of its range) but is HALF of
minorsathome's entire range, so positives kept matching against controls
with a wildly different minorsathome just to satisfy the harder
army-scaler dimension. Confirmed by inspecting matched pairs: mean
minorsathome gap of 1.3 survived a nominally "tight" radius-2 match,
leaving minorsathome as the regression's #1 feature for reasons that were
substantially still the sampling artifact, not real signal. Fix: columns
in --exact-cols must match exactly (no radius search at all -- if a
positive's exact value has no candidates left, it's unmatched, full
stop); only --radius-cols expand outward. minorsathome defaults to exact
since its range is small enough that "close" isn't actually close.
"""
import argparse
import csv
import itertools
import random
import sys
from collections import defaultdict


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("features_csv")
    ap.add_argument("out_csv")
    ap.add_argument("--exact-cols", nargs="+", default=["minorsathome"],
                     help="Match columns that must match EXACTLY -- no "
                          "radius search. Use for low-cardinality columns "
                          "where 'close' isn't actually close.")
    ap.add_argument("--radius-cols", nargs="+", default=["attackerarmyscaler"],
                     help="Match columns allowed to expand outward by "
                          "--max-radius if no exact match is available.")
    ap.add_argument("--max-radius", type=int, default=30,
                     help="Max Chebyshev distance across --radius-cols to "
                          "search outward before giving up on a positive")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)

    with open(args.features_csv, newline="") as f:
        reader = csv.DictReader(f)
        assert reader.fieldnames is not None
        fieldnames = list(reader.fieldnames)
        rows = list(reader)

    positives = [r for r in rows if r["label"] == "1"]
    candidates = [r for r in rows if r["label"] == "0"]
    print(f"Loaded {len(positives)} positives, {len(candidates)} control "
          f"candidates", file=sys.stderr)
    print(f"Exact-match columns: {args.exact_cols}", file=sys.stderr)
    print(f"Radius-search columns (max radius {args.max_radius}): "
          f"{args.radius_cols}", file=sys.stderr)

    def exact_key_of(r):
        return tuple(int(round(float(r[c]))) for c in args.exact_cols)

    def radius_key_of(r):
        return tuple(int(round(float(r[c]))) for c in args.radius_cols)

    # Two-level index: exact key -> radius key -> list of candidate rows.
    # A positive only ever searches within its own exact-key bucket, and
    # within that, expands outward over the radius-cols only.
    buckets = defaultdict(lambda: defaultdict(list))
    for r in candidates:
        buckets[exact_key_of(r)][radius_key_of(r)].append(r)
    for ekey in buckets:
        for rkey in buckets[ekey]:
            rng.shuffle(buckets[ekey][rkey])

    ndim = len(args.radius_cols)

    def offsets_at_radius(radius):
        if radius == 0:
            return [(0,) * ndim]
        rng_vals = range(-radius, radius + 1)
        return [combo for combo in itertools.product(rng_vals, repeat=ndim)
                if max(abs(v) for v in combo) == radius]

    shells = [offsets_at_radius(r) for r in range(args.max_radius + 1)]

    matched_pairs = []
    unmatched = 0
    unmatched_no_exact_bucket = 0
    radius_hist = defaultdict(int)

    for pos_row in positives:
        ekey = exact_key_of(pos_row)
        exact_bucket = buckets.get(ekey)
        if exact_bucket is None:
            unmatched += 1
            unmatched_no_exact_bucket += 1
            continue
        pkey = radius_key_of(pos_row)
        found = None
        found_radius = None
        for radius, shell in enumerate(shells):
            for off in shell:
                cand_key = tuple(pkey[i] + off[i] for i in range(ndim))
                cand_list = exact_bucket.get(cand_key)
                if cand_list:
                    found = cand_list.pop()
                    found_radius = radius
                    break
            if found is not None:
                break
        if found is None:
            unmatched += 1
            continue
        radius_hist[found_radius] += 1
        matched_pairs.append((pos_row, found))

    print(f"Matched {len(matched_pairs)} pairs, {unmatched} positives "
          f"left unmatched ({unmatched_no_exact_bucket} had no candidates "
          f"at all with that exact-column value, {unmatched - unmatched_no_exact_bucket} "
          f"exhausted radius {args.max_radius} on the radius columns)",
          file=sys.stderr)
    print("Match radius histogram (0 = exact match on everything):",
          file=sys.stderr)
    for r in sorted(radius_hist):
        print(f"  radius {r}: {radius_hist[r]}", file=sys.stderr)

    with open(args.out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for pos_row, ctl_row in matched_pairs:
            w.writerow(pos_row)
            w.writerow(ctl_row)

    print(f"Wrote {2 * len(matched_pairs)} matched rows to {args.out_csv}",
          file=sys.stderr)

    import numpy as np
    all_cols = args.exact_cols + args.radius_cols
    def full_key_of(r):
        return tuple(int(round(float(r[c]))) for c in all_cols)
    matched_pos_vals = np.array([full_key_of(p) for p, _ in matched_pairs])
    matched_ctl_vals = np.array([full_key_of(c) for _, c in matched_pairs])
    for i, col in enumerate(all_cols):
        diff = np.abs(matched_pos_vals[:, i] - matched_ctl_vals[:, i])
        print(f"  {col}: mean|diff|={diff.mean():.2f} max|diff|={diff.max():.2f}",
              file=sys.stderr)


if __name__ == "__main__":
    main()