"""Fill-prediction pipeline for mm_model_dataset.csv.

Run:  python research/mm_modeling.py research/mm_model_dataset.csv

FIXES OVER THE FIRST DRAFT
--------------------------
1. The count model now actually uses legs_filled. The draft computed `yfill`
   and then passed `y_train` (which is `complete`, a 0/1) to Poisson and
   NegativeBinomial. The giveaway was NB scoring a WORSE loglike than Poisson
   (-368 vs -166), which is impossible for nested models: Poisson is NB at
   alpha=0, so NB can never fit worse.

2. arm and league are one-hot encoded instead of silently dropped. The draft
   ran them through pd.to_numeric, which turned both to all-NaN, then to 0 via
   fillna, then dropped them as zero-variance. arm carries the largest effect
   in the settled data (tieload strands 50% of its books).

3. The decision threshold is picked on a VALIDATION split and applied to an
   untouched test split. The draft chose the F1-maximising cutoff on y_test and
   then reported P/R/F1 on that same y_test, which is the best achievable score
   on the split rather than an out-of-sample one.

4. cv_auc keeps its intercept. np.std(add_constant(X), axis=0) > 1e-12 is False
   for the column of ones, so the draft's CV silently fit a NO-INTERCEPT model,
   and without the L1 penalty its reported model used. The 0.819 vs 0.733 gap
   was two different models, not an overfitting gap.

5. Imputation and scaling happen INSIDE each fold via a Pipeline, so no fold
   sees statistics computed from its own held-out rows.

WHAT THIS CAN AND CANNOT TELL YOU
---------------------------------
`size` and `arm` are frozen as of 2026-09-07 (SIZE=1, latch3 the only arm).
They vary in this historical sample and carry large coefficients, but they are
CONSTANTS going forward, so a model leaning on them cannot transfer. Both specs
are reported: FULL (historical, for understanding) and CURRENT (frozen
predictors removed, for anything forward-looking).
"""
import sys
import numpy as np
import pandas as pd
import statsmodels.api as sm
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (roc_auc_score, confusion_matrix,
                             precision_recall_curve,
                             precision_recall_fscore_support)

NUMERIC = ['hrs_to_ko_open', 'gross_c_open', 'ask_max_open', 'ask_min_open',
           'spread_c_mean', 'queue_ahead_min', 'queue_ahead_max',
           'book_volume_open', 'leg_volume_min', 'open_interest_sum',
           'is_xlow_any']
FROZEN = ['size']            # pinned 2026-09-07; varies historically only
CATEGORICAL = ['arm']        # league is 36 levels on 517 rows -- too sparse


def build_X(df, include_frozen):
    cols = NUMERIC + (FROZEN if include_frozen else [])
    X = df[cols].apply(pd.to_numeric, errors='coerce')
    X = X.replace([np.inf, -np.inf], np.nan)
    if include_frozen:
        cats = pd.get_dummies(df[CATEGORICAL].astype(str),
                              drop_first=True, dtype=float)
        X = pd.concat([X, cats], axis=1)
    return X.loc[:, X.nunique(dropna=False) > 1]


def pipe():
    return Pipeline([('imp', SimpleImputer(strategy='median')),
                     ('sc', StandardScaler()),
                     ('lr', LogisticRegression(penalty='l1', C=1.0,
                                               solver='liblinear',
                                               max_iter=500))])


def cv_auc(X, y, folds=5, seed=0):
    skf = StratifiedKFold(folds, shuffle=True, random_state=seed)
    out = []
    for tr, te in skf.split(X, y):
        m = pipe().fit(X.iloc[tr], y.iloc[tr])
        out.append(roc_auc_score(y.iloc[te], m.predict_proba(X.iloc[te])[:, 1]))
    return float(np.mean(out)), float(np.std(out))


def run(df, label, include_frozen):
    X, y = build_X(df, include_frozen), df['complete']
    print('')
    print('=' * 68)
    print('%s' % label)
    print('  n=%d, %d features' % (len(X), X.shape[1]))
    print('=' * 68)

    mu, sd = cv_auc(X, y)
    print('  5-fold CV AUC: %.3f (sd %.3f)' % (mu, sd))

    # three-way split so the threshold never touches the test rows
    X_tr, X_tmp, y_tr, y_tmp = train_test_split(
        X, y, test_size=.40, random_state=0, stratify=y)
    X_va, X_te, y_va, y_te = train_test_split(
        X_tmp, y_tmp, test_size=.50, random_state=0, stratify=y_tmp)

    m = pipe().fit(X_tr, y_tr)
    p_va = m.predict_proba(X_va)[:, 1]
    p_te = m.predict_proba(X_te)[:, 1]

    prec, rec, thr = precision_recall_curve(y_va, p_va)
    f1 = 2 * prec * rec / (prec + rec + 1e-9)
    cut = thr[int(np.argmax(f1[:-1]))]

    pred = (p_te > cut).astype(int)
    pr = precision_recall_fscore_support(y_te, pred, average='binary',
                                         zero_division=0)
    base = max(y_te.mean(), 1 - y_te.mean())
    print('  threshold %.3f chosen on validation (n=%d), applied to test (n=%d)'
          % (cut, len(y_va), len(y_te)))
    print('  TEST AUC %.3f | acc %.3f vs base rate %.3f'
          % (roc_auc_score(y_te, p_te), (pred == y_te).mean(), base))
    print('  P %.3f  R %.3f  F1 %.3f   confusion %s'
          % (pr[0], pr[1], pr[2], confusion_matrix(y_te, pred).tolist()))

    coef = pd.Series(m.named_steps['lr'].coef_[0], index=X.columns)
    order = coef.abs().sort_values(ascending=False).index
    print('  coefficients (standardised log-odds, L1; zeros omitted):')
    for k in order:
        if abs(coef[k]) > 1e-6:
            print('     %-24s %+0.3f' % (k, coef[k]))


def count_model(df):
    print('')
    print('=' * 68)
    print('COUNT MODEL on legs_filled (0-3) -- the target the draft never passed')
    print('=' * 68)
    X = build_X(df, include_frozen=True)
    y = df['legs_filled']
    Xn = X.fillna(X.median())
    Xn = (Xn - Xn.mean()) / Xn.std().replace(0, 1)
    Xc = sm.add_constant(Xn)
    ps = sm.Poisson(y, Xc).fit(disp=False, maxiter=300)
    nb = sm.NegativeBinomial(y, Xc).fit(disp=False, maxiter=300)
    print('  Poisson loglike %+0.3f' % ps.llf)
    print('  NB      loglike %+0.3f   alpha %.4f'
          % (nb.llf, nb.params.get('alpha', float('nan'))))
    print('  higher: %s  (NB >= Poisson must hold -- Poisson is NB at alpha=0)'
          % ('NB' if nb.llf > ps.llf else 'Poisson'))
    print('  dispersion: mean %.3f  var %.3f  -> %s'
          % (y.mean(), y.var(),
             'over-dispersed' if y.var() > y.mean() else 'under-dispersed'))
    print('  CAVEAT: legs_filled is bounded 0-3 and 68% of it is 0. A count')
    print('          family is the wrong shape here -- an ordinal or hurdle')
    print('          model fits it properly. Reported for comparability only.')


def main(path):
    df = pd.read_csv(path).dropna(subset=['complete'])
    print('rows %d | complete rate %.3f | legs_filled %s'
          % (len(df), df['complete'].mean(),
             df['legs_filled'].value_counts().sort_index().to_dict()))
    run(df, 'FULL SPEC (historical: includes size and arm)', True)
    run(df, 'CURRENT SPEC (size and arm frozen 2026-09-07, excluded)', False)
    count_model(df)


if __name__ == '__main__':
    main(sys.argv[1] if len(sys.argv) > 1 else 'research/mm_model_dataset.csv')