"""Minimal ranking-objective ablation loop — why lambda_rank fails here.

This repo found that with only ~50 instruments per day the rank-gradient
objectives (lambdarank / rank_xendcg) produce near-zero RankIC, while MSE
objective + RankIC early-stop is the winner. This script replays that check by
training a few LightGBM variants on the same lake split and printing RankIC.

Reference repo impl: tac-qlib/examples/run_rank_objectives.py.

    python examples/run_rank_objectives.py --universe AAPL,MSFT,QQQ
"""

from __future__ import annotations

import argparse
import os

import numpy as np
import pandas as pd

OBJECTIVES = ["mse", "lambdarank", "rank_xendcg"]


def load_frame(lake_root: str, universe: list[str], start: str, end: str) -> pd.DataFrame:
    """Stack lake bars into a qlib-like (datetime, instrument) frame."""
    from tac_qlib.data.config import LakeConfig

    lake = LakeConfig(lake_root, "US")
    frames = []
    for sym in universe:
        p = lake.bar_path("1d", sym)
        if p.exists():
            df = pd.read_parquet(p)[["t", "c"]].rename(columns={"t": "datetime", "c": "close"})
            df["instrument"] = sym
            frames.append(df)
    out = pd.concat(frames, ignore_index=True)
    out["datetime"] = pd.to_datetime(out["datetime"])
    out = out[(out["datetime"] >= start) & (out["datetime"] <= end)]
    return out.set_index(["datetime", "instrument"])


def label_5d(frame: pd.DataFrame) -> pd.Series:
    close = frame["close"].unstack()
    lbl = close.shift(-6) / close.shift(-1) - 1
    return lbl.stack().rename("label")


def train_one(lake_root: str, universe: list[str], objective: str, train: tuple, test: tuple):
    import lightgbm as lgb

    frame = load_frame(lake_root, universe, train[0], test[1])
    label = label_5d(frame)
    data = pd.concat([frame["close"], label], axis=1).dropna()

    tr = data.loc[(data.index.get_level_values(0) >= train[0]) & (data.index.get_level_values(0) <= train[1])]
    te = data.loc[(data.index.get_level_values(0) >= test[0]) & (data.index.get_level_values(0) <= test[1])]

    dtrain = lgb.Dataset(tr[["close"]], label=tr["label"])
    dtest = lgb.Dataset(te[["close"]], label=te["label"])
    params = {"objective": objective, "learning_rate": 0.05, "num_leaves": 15, "verbosity": -1}
    model = lgb.train(params, dtrain, num_boost_round=100, valid_sets=[dtest])

    pred = model.predict(te[["close"]], num_iteration=model.best_iteration)
    label_te = te["label"].to_numpy()
    # per-day RankIC
    days = te.index.get_level_values(0).unique()
    ics = []
    for d in days:
        m = te.index.get_level_values(0) == d
        if m.sum() >= 3:
            ics.append(pd.Series(pred[m]).rank().corr(pd.Series(label_te[m]).rank()))
    return float(np.nanmean(ics))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--lake-root", default=os.environ.get("TAC_LAKE_DIR", ""))
    ap.add_argument("--universe", default="AAPL,MSFT,QQQ,IVV,SMH,TLT")
    args = ap.parse_args()
    universe = [s.strip().upper() for s in args.universe.split(",")]
    train = ("2026-03-01", "2026-05-31")
    test = ("2026-07-01", "2026-08-06")
    print(f"{'objective':<14}{'test RankIC':>12}")
    for obj in OBJECTIVES:
        ic = train_one(args.lake_root, universe, obj, train, test)
        print(f"{obj:<14}{ic:>12.4f}")


if __name__ == "__main__":
    main()