"""Non-neural baselines. A published neural checkpoint has to beat these."""
from __future__ import annotations

import json
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from common import NUMERIC, report, split

df = pd.read_parquet("data/rows.parquet")
df["network_cat"] = df.network.astype("category")
results: dict = {}
Path("runs").mkdir(exist_ok=True)

for h in (7, 14):
    lab = f"paid_{h}d"
    tr, va, te = split(df, h)
    print(f"\n== paid_{h}d  train {len(tr):,}  val {len(va):,}  test {len(te):,} "
          f"(test hosts {te.host.nunique()}, unseen in train)")
    res = {}

    # 1. Recency rule: the fresher the last paid call, the likelier the next one.
    rec = te.days_since_last_call.fillna(99).to_numpy()
    iso = IsotonicRegression(out_of_bounds="clip", increasing=False).fit(
        va.days_since_last_call.fillna(99), va[lab])
    res["recency_rule"] = report(te, iso.predict(rec), h)

    # 2. Logistic regression on numeric features.
    X = lambda d: d[NUMERIC].fillna(-1).to_numpy()
    lr = make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000)).fit(X(tr), tr[lab])
    res["logreg"] = report(te, lr.predict_proba(X(te))[:, 1], h)

    # 3. LightGBM on numeric + network.
    feats = NUMERIC + ["network_cat"]
    gbm = lgb.LGBMClassifier(n_estimators=600, learning_rate=0.03, num_leaves=31,
                             min_child_samples=50, subsample=0.8, subsample_freq=1,
                             colsample_bytree=0.8, verbose=-1)
    gbm.fit(tr[feats], tr[lab], eval_set=[(va[feats], va[lab])],
            callbacks=[lgb.early_stopping(50, verbose=False)])
    p = gbm.predict_proba(te[feats])[:, 1]
    res["lightgbm"] = report(te, p, h)
    imp = sorted(zip(feats, gbm.booster_.feature_importance("gain")), key=lambda x: -x[1])
    res["lightgbm_top_features"] = [(f, round(float(g), 1)) for f, g in imp[:10]]
    if h == 7:
        gbm.booster_.save_model("runs/lightgbm_paid7d.txt")
        np.save("runs/lightgbm_paid7d_test_pred.npy", p)

    for k, v in res.items():
        if k.endswith("features"):
            print(" ", k, v[:6])
            continue
        print(f"  {k:14s}", "  ".join(f"{seg}: auc {m['auc']} ap {m['ap']} ece {m['ece']}"
                                       for seg, m in v.items()))
    results[f"paid_{h}d"] = res

Path("runs/baselines.json").write_text(json.dumps(results, indent=1))