#!/usr/bin/env python
"""
train_gate_xgb.py — AUTHORITATIVE trainer for the DM3 MEME-outcome gate (XGBoost path).

Why XGBoost: it is the one mainstream GBDT with a STABLE, DOCUMENTED, FIRST-PARTY serialization
(booster.save_model("*.json")). No converter chain (sklearn->skl2onnx->ONNX->custom walker), which
was the root cause of the estimator-family churn in the earlier sklearn attempts. The browser reads
the SAME JSON bytes exported here. monotone_constraints=(1,1,1) is native and exact — #151a closes.

Ships ONE artifact: meme_gate.json. No pkl, no onnx, no converter.

Label y = MEME reports >=1 site at p<=0.1. Features [num_seqs, num_sites, median_pos_dist]
(frac_p_defined dropped, #151b). median_pos_dist floored 0.001 / capped 10.
"""
import csv, numpy as np, json
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, brier_score_loss

D="/home/sweaver/programming/axomeme/dm3-gate-validation"
VAL=f"{D}/meme_validation.tsv"; OUT=f"{D}/meme_gate.json"
FLOOR,CAP=0.001,10.0
PARAMS=dict(max_depth=6,n_estimators=500,learning_rate=0.05,min_child_weight=1,
            monotone_constraints=(1,1,1),subsample=0.9,eval_metric='auc',random_state=0,verbosity=0)

def load():
    X,y=[],[]
    for r in csv.DictReader(open(VAL),delimiter="\t"):
        try: ns=int(r["nseq"]);nsit=int(r["nsites"]);bl=float(r["med_bl"]);sf=int(r["sites_found"])
        except: continue
        X.append([ns,nsit,min(max(bl,FLOOR),CAP)]); y.append(int(sf>=1))
    return np.array(X,float),np.array(y)

def check_monotone(clf):
    tot=0; bgs=[[10,100,0.005],[50,432,0.02],[200,1000,0.05],[500,3000,0.1],[3,20,0.001],[1000,8000,0.15]]
    for fi,nm,rng in [(0,"num_seqs",(3,1792)),(1,"num_sites",(6,11416)),(2,"median_pos_dist",(0.001,10))]:
        v=0
        for bg in bgs:
            B=np.tile(bg,(800,1)); B[:,fi]=np.linspace(*rng,800)
            v+=int((np.diff(clf.predict_proba(B)[:,1])<-1e-9).sum())
        print(f"  {nm}: {v} decreasing"); tot+=v
    return tot

def main():
    X,y=load()
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=0.30,random_state=7,stratify=y)
    print(f"[data] {len(y)} jobs | base {y.mean():.3f}")
    ct=xgb.XGBClassifier(**PARAMS).fit(Xtr,ytr)
    s=ct.predict_proba(Xte)[:,1]
    print(f"[held-out test] AUC {roc_auc_score(yte,s):.4f} | Brier {brier_score_loss(yte,s):.4f}")
    clf=xgb.XGBClassifier(**PARAMS).fit(X,y)
    print("[monotonicity — native monotone_constraints, want 0]")
    viol=check_monotone(clf)
    # THE deliverable: XGBoost's own JSON. Nothing else.
    clf.get_booster().save_model(OUT)
    # sidecar metadata (not needed by the browser, just provenance)
    json.dump({"features":["num_seqs","num_sites","median_pos_dist"],"label":"MEME reports >=1 site at p<=0.1",
               "bl_floor":FLOOR,"bl_cap":CAP,"objective":"binary:logistic (raw output -> sigmoid)",
               "test_auc":round(roc_auc_score(yte,s),4),"base_rate":round(float(y.mean()),3),
               "hyperparams":PARAMS,"note":"XGBoost native JSON; browser reads meme_gate.json directly"},
              open(f"{D}/meme_gate.meta.json","w"),indent=2)
    print(f"[saved] {OUT}  ({'CLEAN' if viol==0 else f'WARN {viol}'}) + meme_gate.meta.json")

if __name__=="__main__": main()