#!/usr/bin/env python3
"""Fit a logistic regression on king_features.csv's raw per-defect
breakdown against the mate/control label, to see whether _EvalKing's
defects should be weighted differently than the hand-picked constants
currently baked into KING_SAFETY_BY_COUNTER / KING_COUNTER_BY_ATTACK_PATTERN
/ EmptyAttackedSquare etc. See CLAUDE.md's king-safety discussion for the
full methodology this is built on.

Deliberately excludes `countertotal` and `safetyvalue` from the feature
set -- those are the CURRENT engine's own derived output (the thing we're
trying to second-guess), not a raw defect. Using them as inputs would be
circular. `initloc` is kept as a feature (a compact king-square exposure
signal) since, unlike the other two, it doesn't already fold in this
sample's own defect counts.

Attacker material confound (conversation context: _EvalKing scales its
final score by REDUCED_MATERIAL_DOWN_SCALER[armyscaler], but the raw
defect counts here are unscaled, and a real mate is naturally rarer under
a weak attacker regardless of what the defects say) is handled two ways:
included as a plain covariate (attackerarmyscaler), and separately
validated by refitting on the top-tercile-army-strength subset and
diffing coefficients against the full-population fit.

Ring-square bucketing: slots 3/4 (the KingSafetyDeltas -1/+1 deltas -- the
two SAME-RANK side squares, not "front" despite an earlier mislabeling in
conversation) are kept individually resolved since they carried real,
stable signal in initial runs. The other six slots (the diagonal +
vertical squares one rank away) are collapsed into "front"/"back" bucket
sums, color-dependent per EmptyAttackedSquare's own existing asymmetry
(WHITE's front is slots 0/1/2, BLACK's is 5/6/7 -- see
KingRingBucket-equivalent logic below) -- these six raw per-slot
coefficients were the noisiest, most sign-flipping entries in every
full-vs-top-tercile drift table, and bucketing trades a little resolution
for a lot more stability at the sample sizes actually available here.
`filestorm` itself is dropped as a feature in favor of its two already-
independent sub-components (filestormopenfile, filestormpawn) which map
directly onto separate hand-picked tables in eval.c (KingFileDefects /
KingStormingPawnDefects) -- no reason to force a regression to re-derive
their relative importance from one pre-summed number.
"""
import argparse
import csv
import sys

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score, accuracy_score

PIECE_NAMES = ("pawn", "minor", "rook", "queen")
PER_SQUARE_KEYS = ["offboard"] + [f"atk_{p}" for p in PIECE_NAMES] + \
    ["undef", "def", "xray", "flight"]

SIDE_SLOTS = (3, 4)  # kept individually resolved
# Color-dependent diagonal/vertical bucket membership, mirroring
# EmptyAttackedSquare's own asymmetry (chess.h: WHITE=1, BLACK=0).
BUCKET_SLOTS = {
    "white": {"front": (0, 1, 2), "back": (5, 6, 7)},
    "black": {"front": (5, 6, 7), "back": (0, 1, 2)},
}

FEATURE_COLUMNS = []
for _sq in SIDE_SLOTS:
    for _k in PER_SQUARE_KEYS:
        FEATURE_COLUMNS.append(f"sq{_sq}_{_k}")
for _bucket in ("front", "back"):
    for _k in PER_SQUARE_KEYS:
        FEATURE_COLUMNS.append(f"{_bucket}_{_k}")
FEATURE_COLUMNS += ["filestormopenfile", "filestormpawn", "queenproximity",
                    "initloc", "castled", "minorsathome",
                    "misscastleshort", "misscastlelong", "attackerarmyscaler"]

# Reference-only columns: the engine's OWN current derived output for this
# position, kept for comparison but never fed into the model as inputs.
REFERENCE_COLUMNS = ["countertotal", "safetyvalue", "filestorm"]


def _row_to_features(r):
    color = r["mated_color_or_sampled_color"]
    bucket_slots = BUCKET_SLOTS[color]
    out = []
    for sq in SIDE_SLOTS:
        for k in PER_SQUARE_KEYS:
            out.append(float(r[f"sq{sq}_{k}"]))
    for bucket_name in ("front", "back"):
        slots = bucket_slots[bucket_name]
        for k in PER_SQUARE_KEYS:
            # Sum across the 3 slots in this bucket -- offboard included,
            # since a king near the edge legitimately has fewer on-board
            # ring squares and that's itself signal, not noise to hide.
            out.append(sum(float(r[f"sq{s}_{k}"]) for s in slots))
    for c in ["filestormopenfile", "filestormpawn", "queenproximity",
              "initloc", "castled", "minorsathome",
              "misscastleshort", "misscastlelong", "attackerarmyscaler"]:
        out.append(float(r[c]))
    return out


def load_data(csv_path):
    with open(csv_path, newline="") as f:
        reader = csv.DictReader(f)
        rows = list(reader)
    y = np.array([int(r["label"]) for r in rows], dtype=int)
    X = np.array([_row_to_features(r) for r in rows], dtype=float)
    ref = {c: np.array([float(r[c]) for r in rows]) for c in REFERENCE_COLUMNS}
    return X, y, ref


def fit_and_report(X, y, label, top_n=25):
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y, test_size=0.2, random_state=42, stratify=y)

    clf = LogisticRegression(max_iter=2000, C=1.0)
    clf.fit(X_train, y_train)

    y_pred = clf.predict(X_test)
    y_prob = clf.predict_proba(X_test)[:, 1]
    acc = accuracy_score(y_test, y_pred)
    auc = roc_auc_score(y_test, y_prob)

    print(f"\n=== {label} (n={len(y)}, train={len(y_train)}, test={len(y_test)}) ===")
    print(f"Test accuracy: {acc:.4f}  Test AUC: {auc:.4f}")

    # Standardized coefficients are directly comparable to each other
    # (same units: change in log-odds per one std-dev of that feature)
    # regardless of the feature's raw scale (a 0/1 flag vs. armyscaler's
    # 0-62 range) -- that comparability is the entire point of this
    # exercise, so report these, not raw unscaled coefficients.
    coefs = clf.coef_[0]
    order = np.argsort(-np.abs(coefs))
    print(f"\nTop {top_n} features by |standardized coefficient| "
          f"(positive = more mate-predictive):")
    for i in order[:top_n]:
        print(f"  {FEATURE_COLUMNS[i]:28s} {coefs[i]:+.4f}")

    return dict(zip(FEATURE_COLUMNS, coefs)), scaler, clf


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("features_csv")
    ap.add_argument("--top-n", type=int, default=25)
    args = ap.parse_args()

    X, y, ref = load_data(args.features_csv)
    print(f"Loaded {len(y)} rows, {X.shape[1]} features, "
          f"{y.sum()} positives / {(y == 0).sum()} controls", file=sys.stderr)

    print("\n--- Reference: current engine's own derived score ---")
    for c in REFERENCE_COLUMNS:
        pos_mean = ref[c][y == 1].mean()
        ctl_mean = ref[c][y == 0].mean()
        print(f"  {c}: mean(mate)={pos_mean:.1f}  mean(control)={ctl_mean:.1f}  "
              f"ratio={pos_mean / max(ctl_mean, 1e-9):.2f}x")

    full_coefs, _, _ = fit_and_report(X, y, "Full population", top_n=args.top_n)

    # Attacker-material confound check: refit on the top tercile of
    # attacker army strength only, where REDUCED_MATERIAL_DOWN_SCALER is
    # near its ceiling and roughly constant across the subset -- if
    # coefficients move a lot vs. the full-population fit, the raw defect
    # weights aren't material-invariant and the interaction should be
    # modeled explicitly rather than left as a plain covariate.
    army_idx = FEATURE_COLUMNS.index("attackerarmyscaler")
    army_vals = X[:, army_idx]
    top_tercile_cutoff = np.percentile(army_vals, 66.7)
    mask = army_vals >= top_tercile_cutoff
    print(f"\nTop-tercile army strength cutoff: attackerarmyscaler >= "
          f"{top_tercile_cutoff:.0f} ({mask.sum()} rows)", file=sys.stderr)
    strong_coefs, _, _ = fit_and_report(
        X[mask], y[mask], "Top-tercile attacker army strength", top_n=args.top_n)

    print(f"\n=== Coefficient drift: full population vs. top-tercile army strength ===")
    print(f"(large drift on a feature means its true weight is NOT "
          f"material-invariant -- REDUCED_MATERIAL_DOWN_SCALER alone may "
          f"not be enough to capture that interaction)")
    drift = [(c, full_coefs[c], strong_coefs.get(c, 0.0),
              abs(full_coefs[c] - strong_coefs.get(c, 0.0)))
             for c in FEATURE_COLUMNS if c != "attackerarmyscaler"]
    drift.sort(key=lambda t: -t[3])
    for c, fv, sv, d in drift[:args.top_n]:
        print(f"  {c:28s} full={fv:+.4f}  top_tercile={sv:+.4f}  drift={d:.4f}")


if __name__ == "__main__":
    main()