train_chair.py

"""T: train the empty-chair concealment model.

Feature view `empty_chair_fv` = company_registry JOIN psc_shape on company_number
(label `is_revealed` already lives in company_registry). v4 recipe from the
autoresearch jul14 round (autoresearch/REPORT.md): 10-seed LightGBM soft-vote,
unweighted, uncalibrated (the app presents rank, never probability), interaction
features and grouped out-of-fold target encoding of post_area / sic_section /
sic_code. Derivations live in chair_features.derive_features and the frozen
encoding maps ship inside the artifact (te_maps.json), so scoring cannot skew.

Every honesty control the design demands runs here and its number is printed and
saved, so a reader can see whether the signal is concealment shape or an artifact:
  - grouped split by office_group: no formation mill straddles train and test.
  - blind rule baseline: flag if silent / corporate-only / foreign-corporate PSC.
  - demographics-only control: year + sic_section + region only. If it matches the
    full model, the signal is population bias, not concealment.
  - shuffle-label control (full recipe incl. TE): must collapse to chance.
Headline = PR-AUC and precision@k lift over the blind rule (PU labels => lower bound).

Registers `empty_chair` with eval JSON + PR / calibration / importance PNGs.
"""

from __future__ import annotations

import argparse
import json
import os
import tempfile

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import HistGradientBoostingClassifier, VotingClassifier
from sklearn.metrics import (average_precision_score, brier_score_loss,
                             precision_recall_curve, roc_auc_score)
from sklearn.model_selection import GroupKFold, GroupShuffleSplit
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OrdinalEncoder

import hopsworks
from chair_features import CAT_FEATURES as CAT
from chair_features import NUM_FEATURES as NUM
from chair_features import DERIVED_NUM, TE_COLS, derive_features

DEMOG = ["incorporation_year", "sic_section", "post_area", "country"]
# the legitimate-structure confounds the bias audit flagged (property/holding SPV shape)
STRUCTURE = ["sic_section", "is_holding_sic", "is_dormant_sic", "is_mill_address",
             "office_company_count"]
LABEL = "is_revealed"
GROUP = "office_group"


def get_training_frame():
    """Read THROUGH the feature view so a training dataset version materializes and
    the registered model carries provenance (fv -> td -> model), instead of a raw
    query.read() that leaves the feature view decorative."""
    project = hopsworks.login()
    fs = project.get_feature_store()
    reg = fs.get_feature_group("company_registry", version=1)
    psc = fs.get_feature_group("psc_shape", version=1)
    query = reg.select_all().join(psc.select_except(["company_number"]), on=["company_number"])
    fv = fs.get_or_create_feature_view(
        name="empty_chair_fv", version=1, query=query, labels=[LABEL],
        description="Registry + PSC concealment shape for UK companies; label = later-revealed hidden owner",
    )
    X, y = fv.training_data(description="empty-chair full frame, grouped holdout split in-code")
    df = pd.concat([X, y], axis=1)  # office_group rides in X; label joined back
    td_version = max(t.version for t in fv.get_training_datasets())
    print(f"training dataset v{td_version} ({len(df)} rows) via {fv.name} v{fv.version}")
    return project, fv, td_version, df


def _pre(cols_cat, cols_num):
    return ColumnTransformer(
        [("cat", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1,
                                encoded_missing_value=-1), cols_cat),
         ("num", "passthrough", cols_num)])


LGBM_PARAMS = dict(
    n_estimators=1500, learning_rate=0.02, num_leaves=31,
    reg_lambda=5.0, scale_pos_weight=1.0, min_child_samples=80,
    colsample_bytree=0.8, subsample=0.8, subsample_freq=1)


def make_pipeline(cols_cat, cols_num, seeds=10, params=None):
    """The v4 winner: soft-vote over seed-varied LightGBMs, unweighted, uncalibrated."""
    from lightgbm import LGBMClassifier

    def lgbm(seed):
        return LGBMClassifier(**(LGBM_PARAMS | (params or {})),
                              random_state=seed, n_jobs=-1, verbose=-1)
    clf = VotingClassifier([(f"s{s}", lgbm(s)) for s in range(seeds)], voting="soft")
    return Pipeline([("pre", _pre(cols_cat, cols_num)), ("clf", clf)])


def make_control_pipeline(cols_cat, cols_num):
    """HGB, for the demographics and shuffle controls (cheap, single fit)."""
    clf = HistGradientBoostingClassifier(
        max_iter=400, learning_rate=0.05, max_leaf_nodes=31,
        l2_regularization=1.0, early_stopping=True, validation_fraction=0.15,
        class_weight="balanced", random_state=0)
    return Pipeline([("pre", _pre(cols_cat, cols_num)), ("clf", clf)])


TE_SMOOTH = 20


def fit_target_encoding(Xtr, ytr, groups_tr, Xte):
    """Grouped out-of-fold target encoding on the train slice (no leak into the
    trees), full-train maps applied to the holdout and frozen for the artifact."""
    prior = float(ytr.mean())
    maps = {"__prior__": prior}
    for c in TE_COLS:
        oof = np.full(len(Xtr), prior)
        for fi, vi in GroupKFold(n_splits=5).split(Xtr, ytr, groups_tr):
            agg = pd.Series(ytr[fi]).groupby(Xtr[c].astype(str).iloc[fi].values).agg(["sum", "count"])
            smooth = (agg["sum"] + TE_SMOOTH * prior) / (agg["count"] + TE_SMOOTH)
            oof[vi] = Xtr[c].astype(str).iloc[vi].map(smooth).fillna(prior).values
        Xtr[c + "_te"] = oof
        agg = pd.Series(ytr).groupby(Xtr[c].astype(str).values).agg(["sum", "count"])
        smooth = (agg["sum"] + TE_SMOOTH * prior) / (agg["count"] + TE_SMOOTH)
        Xte[c + "_te"] = Xte[c].astype(str).map(smooth).fillna(prior).values
        maps[c] = {str(k): float(v) for k, v in smooth.items()}
    return maps


def precision_at_k(y, score, k):
    order = np.argsort(score)[::-1][:k]
    return float(y[order].mean())


def blind_rule(df):
    return ((df["psc_silence"] == 1) | (df["psc_corporate_only"] == 1) |
            (df["psc_foreign_corporate"] > 0)).astype(int).values


def evaluate(name, y, score):
    ap = average_precision_score(y, score)
    roc = roc_auc_score(y, score)
    base = float(y.mean())
    out = {
        "model": name, "pr_auc": round(ap, 4), "roc_auc": round(roc, 4),
        "base_rate": round(base, 4), "pr_auc_lift": round(ap / base, 2),
        "precision_at_100": round(precision_at_k(y, score, 100), 4),
        "precision_at_1000": round(precision_at_k(y, score, 1000), 4),
    }
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--no-register", action="store_true")
    ap.add_argument("--seeds", type=int, default=10,
                    help="LGBMs in the soft-vote; 1 = the fast single-seed variant")
    ap.add_argument("--drop", default="",
                    help="comma-separated feature names to exclude from CAT/NUM")
    ap.add_argument("--params", default="",
                    help="LGBM overrides as k=v,k=v (numbers parsed)")
    ap.add_argument("--no-structure", action="store_true",
                    help="ablation: drop the property/holding/mill structure confounds")
    args = ap.parse_args()

    global CAT, NUM
    if args.no_structure:
        CAT = [c for c in CAT if c not in STRUCTURE]
        NUM = [c for c in NUM if c not in STRUCTURE]
        print(f"ABLATION: dropped structure confounds; {len(CAT+NUM)} features remain")
    if args.drop:
        gone = {c.strip() for c in args.drop.split(",") if c.strip()}
        CAT = [c for c in CAT if c not in gone]
        NUM = [c for c in NUM if c not in gone]
        print(f"dropped {sorted(gone)}; {len(CAT+NUM)} features remain")

    project, fv, td_version, df = get_training_frame()
    df = df.dropna(subset=[LABEL]).reset_index(drop=True)
    df[LABEL] = df[LABEL].astype(int)
    y = df[LABEL].values
    groups = df[GROUP].astype(str).values
    print(f"frame: {len(df)} rows, {y.sum()} positive ({y.mean():.3%})")

    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=0)
    tr, te = next(gss.split(df, y, groups))
    Xtr, Xte = df.iloc[tr], df.iloc[te]
    ytr, yte = y[tr], y[te]
    print(f"grouped split: train {len(tr)} / test {len(te)}; "
          f"train groups {len(set(groups[tr]))}, test groups {len(set(groups[te]))}, "
          f"overlap {len(set(groups[tr]) & set(groups[te]))}")

    results = {}

    # --- blind rule baseline
    results["blind_rule"] = evaluate("blind_rule", yte, blind_rule(Xte).astype(float))

    # --- full model: derived features + OOF target encoding, fit on the whole
    # train slice (the round showed a calibration slice buys nothing for ranking)
    Xtr = derive_features(df.iloc[tr].copy())
    Xte = derive_features(df.iloc[te].copy())
    te_maps = fit_target_encoding(Xtr, ytr, groups[tr], Xte)
    overrides = {}
    for kv in filter(None, args.params.split(",")):
        k, v = kv.split("=")
        overrides[k.strip()] = float(v) if "." in v else int(v)
    if overrides:
        print(f"LGBM overrides: {overrides}")

    cols_num = NUM + DERIVED_NUM + [c + "_te" for c in TE_COLS]
    pipe = make_pipeline(CAT, cols_num, seeds=args.seeds, params=overrides)
    pipe.fit(Xtr[CAT + cols_num], ytr)
    score_full = pipe.predict_proba(Xte[CAT + cols_num])[:, 1]
    results["full"] = evaluate("full", yte, score_full)
    results["full"]["brier"] = round(brier_score_loss(yte, score_full), 4)

    # --- demographics-only control
    pdem = make_control_pipeline(["sic_section", "post_area", "country"], ["incorporation_year"])
    pdem.fit(Xtr[DEMOG], ytr)
    score_dem = pdem.predict_proba(Xte[DEMOG])[:, 1]
    results["demographics_only"] = evaluate("demographics_only", yte, score_dem)

    # --- shuffle-label control on the full recipe: TE refit on the shuffled labels
    # (single seed; a leak would show regardless of the ensemble size)
    rng = np.random.RandomState(0)
    yshuf = rng.permutation(ytr)
    Xtr_s = derive_features(df.iloc[tr].copy())
    Xte_s = derive_features(df.iloc[te].copy())
    fit_target_encoding(Xtr_s, yshuf, groups[tr], Xte_s)
    pshuf = make_pipeline(CAT, cols_num, seeds=1)
    pshuf.fit(Xtr_s[CAT + cols_num], yshuf)
    score_shuf = pshuf.predict_proba(Xte_s[CAT + cols_num])[:, 1]
    results["shuffle_label"] = evaluate("shuffle_label", yte, score_shuf)

    print("\n=== RESULTS (grouped holdout) ===")
    for k, v in results.items():
        print(f"{k:20s} PR-AUC {v['pr_auc']:.3f} (lift {v['pr_auc_lift']:.1f})  "
              f"ROC {v['roc_auc']:.3f}  P@100 {v['precision_at_100']:.3f}  P@1000 {v['precision_at_1000']:.3f}")

    if args.no_register:
        return

    register(project, fv, td_version, pipe, te_maps, CAT, cols_num, results, Xte, yte, score_full)


def register(project, fv, td_version, pipe, te_maps, cols_cat, cols_num, results, Xte, yte, score_full):
    tmp = tempfile.mkdtemp()
    # PR curve
    prec, rec, _ = precision_recall_curve(yte, score_full)
    plt.figure(figsize=(5, 4)); plt.plot(rec, prec)
    plt.xlabel("recall"); plt.ylabel("precision")
    plt.title(f"PR (AP={results['full']['pr_auc']}, base={results['full']['base_rate']})")
    plt.tight_layout(); plt.savefig(f"{tmp}/pr_curve.png", dpi=120); plt.close()
    # calibration (the model is uncalibrated by design, the plot shows it honestly)
    from sklearn.calibration import calibration_curve
    frac, mean = calibration_curve(yte, score_full, n_bins=10, strategy="quantile")
    plt.figure(figsize=(5, 4)); plt.plot([0, 1], [0, 1], "--", c="gray"); plt.plot(mean, frac, "o-")
    plt.xlabel("predicted"); plt.ylabel("observed"); plt.title("calibration (uncalibrated, rank model)")
    plt.tight_layout(); plt.savefig(f"{tmp}/calibration.png", dpi=120); plt.close()
    # importance: mean LightGBM gain across the vote (permutation over the 10-model
    # ensemble costs more than it teaches here)
    names = list(pipe.named_steps["pre"].get_feature_names_out())
    gains = np.mean([m.feature_importances_ for m in
                     pipe.named_steps["clf"].estimators_], axis=0)
    order = np.argsort(gains)[::-1][:15]
    plt.figure(figsize=(6, 5))
    plt.barh([names[i].split("__")[-1] for i in order][::-1], gains[order][::-1])
    plt.title("mean LightGBM gain importance"); plt.tight_layout()
    plt.savefig(f"{tmp}/importance.png", dpi=120); plt.close()

    with open(f"{tmp}/metrics.json", "w") as f:
        json.dump(results, f, indent=2)
    # bundle the full contract WITH the model so serving cannot drift from training:
    # raw columns, derived columns, and the frozen target-encoding maps
    with open(f"{tmp}/features.json", "w") as f:
        json.dump({"cat": cols_cat, "num": cols_num, "features": cols_cat + cols_num}, f)
    with open(f"{tmp}/te_maps.json", "w") as f:
        json.dump(te_maps, f)
    import joblib
    joblib.dump(pipe, f"{tmp}/model.joblib")

    mr = project.get_model_registry()
    model = mr.python.create_model(
        name="empty_chair",
        metrics={"pr_auc": results["full"]["pr_auc"],
                 "pr_auc_lift": results["full"]["pr_auc_lift"],
                 "precision_at_100": results["full"]["precision_at_100"],
                 "roc_auc": results["full"]["roc_auc"]},
        description="Concealment-shape model: is a UK company's beneficial-ownership disclosure evasive (PU-labelled, lower bound)",
        feature_view=fv, training_dataset_version=td_version,
    )
    model.save(tmp)
    print(f"\nregistered empty_chair v{model.version} (fv {fv.name} v{fv.version}, td v{td_version})")


if __name__ == "__main__":
    main()