#!/usr/bin/env python3
"""Generate a learning-to-rank training script (LightGBM lambdarank and/or XGBoost rank:ndcg).

The script it writes:
  - loads the CSV, sorts by query id (rankers take group sizes in row order), keeps the chosen features
  - drops queries whose documents all share one label from TRAINING (no pairwise signal), keeps them for
    evaluation bookkeeping
  - runs GroupKFold by query (no query in both train and validation), with early stopping on NDCG@k of an
    inner validation split of each training fold, and collects out-of-fold scores
  - LightGBM: LGBMRanker(objective="lambdarank", label_gain = 2^label - 1, lambdarank_truncation_level),
    XGBoost: XGBRanker(objective="rank:ndcg", lambdarank_pair_method="topk")
  - writes oof_scores.csv (qid, label, [doc id], one column per model plus the raw features) so that
    rank_eval.py compares the models with single-feature baselines on identical queries, prints NDCG@k per
    model, refits each model on all data with the mean best iteration and saves it, and prints feature
    importance (gain)
Needs pandas, numpy, scikit-learn, lightgbm and/or xgboost, joblib.

Usage:
    python gen_ltr_script.py --data ltr.csv --qid-col qid --label-col label --features bm25 freshness clicks -o ltr.py
    python gen_ltr_script.py --data train.csv --qid-col query --label-col rel --exclude doc_id --models lgbm --k 5
    python gen_ltr_script.py --data ltr.csv --features bm25 clicks --models lgbm xgb --folds 5 --k 10

Exit codes: 0 written, 2 usage error.
"""

import argparse
import os
import sys

TEMPLATE = '''"""Learning-to-rank models for {data!r}. Generated by gen_ltr_script.py.

Run:  python {script}
Then: python rank_eval.py oof_scores.csv --qid-col {qid} --label-col {label} --scores {score_cols} --baseline {baseline} --k {k}
"""
import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import GroupKFold, GroupShuffleSplit

DATA, QID, LABEL, FEATURES, EXCLUDE = {data!r}, {qid!r}, {label!r}, {features!r}, {exclude!r}
MODELS, K, FOLDS, SEED = {models!r}, {k}, {folds}, 0

df = pd.read_csv(DATA)
df = df.sort_values(QID, kind="stable").reset_index(drop=True)       # rankers need contiguous query blocks
feats = FEATURES or [c for c in df.columns if c not in [QID, LABEL] + EXCLUDE and pd.api.types.is_numeric_dtype(df[c])]
df[LABEL] = df[LABEL].astype(int)
informative = df.groupby(QID)[LABEL].transform("nunique") > 1
print(f"{{df[QID].nunique()}} queries, {{len(df)}} documents, features: {{feats}}; "
      f"{{(~informative).groupby(df[QID]).first().sum()}} single-label queries excluded from training")


def groups_of(frame):
    return frame.groupby(QID, sort=False).size().to_numpy()


def make(name, n_estimators=2000):
    if name == "lgbm":
        from lightgbm import LGBMRanker
        max_label = int(df[LABEL].max())
        return LGBMRanker(objective="lambdarank", n_estimators=n_estimators, learning_rate=0.05, num_leaves=15,
                          min_child_samples=10, subsample=0.8, subsample_freq=1, colsample_bytree=0.8,
                          label_gain=[2 ** i - 1 for i in range(max_label + 1)], lambdarank_truncation_level=K,
                          random_state=SEED, verbose=-1)
    from xgboost import XGBRanker
    return XGBRanker(objective="rank:ndcg", n_estimators=n_estimators, learning_rate=0.05, max_depth=4,
                     subsample=0.8, colsample_bytree=0.8, lambdarank_pair_method="topk",
                     lambdarank_num_pair_per_sample=K, eval_metric=f"ndcg@{{K}}", early_stopping_rounds=100,
                     random_state=SEED)


def fit(name, model, tr, va):
    if name == "lgbm":
        import lightgbm as lgb
        model.fit(tr[feats], tr[LABEL], group=groups_of(tr), eval_set=[(va[feats], va[LABEL])],
                  eval_group=[groups_of(va)], eval_at=[K], callbacks=[lgb.early_stopping(100, verbose=False)])
        return model.best_iteration_ or model.n_estimators
    model.fit(tr[feats], tr[LABEL], qid=tr[QID].astype("category").cat.codes, eval_set=[(va[feats], va[LABEL])],
              eval_qid=[va[QID].astype("category").cat.codes], verbose=False)
    return (model.best_iteration or model.n_estimators - 1) + 1


oof = df[[QID, LABEL] + [c for c in df.columns if c not in [QID, LABEL] and c not in feats] + feats].copy()
best_iters = {{m: [] for m in MODELS}}
gkf = GroupKFold(n_splits=FOLDS)
for fold, (tr_idx, te_idx) in enumerate(gkf.split(df, groups=df[QID])):
    train_part = df.iloc[tr_idx]
    train_part = train_part[informative.iloc[tr_idx].to_numpy()]
    inner = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=SEED)
    a_idx, v_idx = next(inner.split(train_part, groups=train_part[QID]))
    tr = train_part.iloc[a_idx].sort_values(QID, kind="stable")
    va = train_part.iloc[v_idx].sort_values(QID, kind="stable")
    for m in MODELS:
        model = make(m)
        best_iters[m].append(fit(m, model, tr, va))
        oof.loc[df.index[te_idx], m] = model.predict(df.iloc[te_idx][feats])


def ndcg_at(frame, col):
    vals = []
    for _, g in frame.groupby(QID):
        # ties broken pessimistically (lower label first), as in rank_eval.py: tied scores must not look good
        lab = g.sort_values([col, LABEL], ascending=[False, True], kind="stable")[LABEL].to_numpy()[:K]
        ideal = np.sort(g[LABEL].to_numpy())[::-1][:K]
        disc = 1 / np.log2(np.arange(2, len(lab) + 2))
        idcg = ((2.0 ** ideal - 1) * disc[:len(ideal)]).sum()
        if idcg > 0:
            vals.append(((2.0 ** lab - 1) * disc).sum() / idcg)
    return float(np.mean(vals))


oof.to_csv("oof_scores.csv", index=False)
for m in MODELS:
    print(f"{{m}}: out-of-fold NDCG@{{K}} {{ndcg_at(oof, m):.4f}} (best iterations {{best_iters[m]}})")
for f in feats:
    print(f"single feature {{f}}: NDCG@{{K}} {{ndcg_at(oof, f):.4f}}")
full = df[informative.to_numpy()]
for m in MODELS:
    n_iter = max(10, int(np.mean(best_iters[m]) * 1.1))
    model = make(m, n_estimators=n_iter)
    if m == "lgbm":
        model.fit(full[feats], full[LABEL], group=groups_of(full))
        imp = dict(zip(feats, model.booster_.feature_importance(importance_type="gain")))
    else:
        model.set_params(early_stopping_rounds=None)
        model.fit(full[feats], full[LABEL], qid=full[QID].astype("category").cat.codes)
        score = model.get_booster().get_score(importance_type="total_gain")
        imp = {{f: score.get(f, 0.0) for f in feats}}
    tot = sum(imp.values()) or 1.0
    print(f"{{m}} gain importance: " + ", ".join(f"{{f}} {{v / tot:.0%}}" for f, v in sorted(imp.items(), key=lambda kv: -kv[1])))
    joblib.dump(model, f"ranker_{{m}}.joblib")
print("wrote oof_scores.csv and ranker_<model>.joblib")
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a learning-to-rank training script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data", required=True)
    ap.add_argument("--qid-col", default="qid")
    ap.add_argument("--label-col", default="label")
    ap.add_argument("--features", nargs="*", default=[])
    ap.add_argument("--exclude", nargs="*", default=[])
    ap.add_argument("--models", nargs="+", choices=["lgbm", "xgb"], default=["lgbm"])
    ap.add_argument("--k", type=int, default=10)
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    if a.k < 1 or a.folds < 2:
        print("error: --k must be >= 1 and --folds >= 2", file=sys.stderr)
        return 2
    baseline = a.features[0] if a.features else "<feature>"
    text = TEMPLATE.format(data=a.data, qid=a.qid_col, label=a.label_col, features=a.features, exclude=a.exclude,
                           models=a.models, k=a.k, folds=a.folds, score_cols=" ".join(a.models), baseline=baseline,
                           script=os.path.basename(a.output or "ltr.py"))
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {', '.join(a.models)}, NDCG@{a.k}, {a.folds} grouped folds", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())