"""
Walk-forward signal combination.

Features are cross-sectional percentile ranks of every factor among PIT
members on each date; the target is the cross-sectional rank of the forward
open-to-open return over ``horizon`` days. Each calendar year is predicted by
a model trained only on samples whose labels ended before that year starts
(a ``horizon``-day purge), so every score is out-of-sample.

    from research.lab import ml
    feats = ml.feature_ranks(panel, names)
    score = ml.walk_forward(panel, feats, model="lgbm", horizon=20)
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from research.lab.factors import FACTORS
from research.lab.report import forward_returns


def feature_ranks(p, names: list[str]) -> dict[str, pd.DataFrame]:
    out = {}
    for n in names:
        f = FACTORS[n](p).reindex(index=p.dates, columns=p.close.columns)
        out[n] = f.where(p.member).rank(axis=1, pct=True)
    return out


def _stack(feats: dict[str, pd.DataFrame], rows: pd.DatetimeIndex) -> pd.DataFrame:
    cols = {n: f.loc[rows].stack(future_stack=True) for n, f in feats.items()}
    return pd.DataFrame(cols)


def walk_forward(p, feats: dict[str, pd.DataFrame], model: str = "lgbm", horizon: int = 20,
                 sample_every: int = 5, first_test_year: int = 2010, min_train_years: int = 4,
                 params: dict | None = None, seed: int = 0) -> pd.DataFrame:
    fwd = forward_returns(p, horizon).where(p.member)
    target = fwd.rank(axis=1, pct=True)
    dates = p.dates
    sample_rows = dates[::sample_every]
    X_all = _stack(feats, sample_rows)
    y_all = target.loc[sample_rows].stack(future_stack=True).reindex(X_all.index)
    ok = y_all.notna() & (X_all.notna().sum(axis=1) >= len(feats) // 2)
    X_all, y_all = X_all[ok], y_all[ok]
    sample_dates = X_all.index.get_level_values(0)

    score = pd.DataFrame(np.nan, index=dates, columns=p.close.columns)
    years = sorted({d.year for d in dates if d.year >= first_test_year})
    for yr in years:
        test_start = pd.Timestamp(f"{yr}-01-01")
        pos = dates.searchsorted(test_start)
        if pos >= len(dates):
            break
        # purge: the label of a sample at d uses opens through d + horizon + 1
        cutoff = dates[max(pos - horizon - 2, 0)]
        tr = sample_dates <= cutoff
        if (sample_dates[tr].max() - sample_dates[tr].min()).days < 365 * min_train_years:
            continue
        Xtr, ytr = X_all[tr], y_all[tr]
        test_rows = dates[(dates >= test_start) & (dates < pd.Timestamp(f"{yr + 1}-01-01"))]
        Xte = _stack(feats, test_rows)
        pred = _fit_predict(model, Xtr, ytr, Xte, params, seed)
        s = pd.Series(pred, index=Xte.index).unstack()
        score.loc[s.index, s.columns] = s.values
    return score


def _fit_predict(model, Xtr, ytr, Xte, params, seed):
    if model == "lgbm":
        import lightgbm as lgb
        prm = dict(objective="regression", learning_rate=0.03, num_leaves=31,
                   min_data_in_leaf=1000, bagging_fraction=0.7, bagging_freq=1,
                   feature_fraction=0.7, lambda_l2=1.0, seed=seed, verbose=-1, num_threads=8)
        prm.update(params or {})
        rounds = prm.pop("n_estimators", 300)
        booster = lgb.train(prm, lgb.Dataset(Xtr.values, ytr.values), num_boost_round=rounds)
        return booster.predict(Xte.values)
    if model == "ridge":
        Xf = Xtr.fillna(0.5).values - 0.5
        lam = (params or {}).get("lam", 10.0) * len(Xf)
        beta = np.linalg.solve(Xf.T @ Xf + lam * np.eye(Xf.shape[1]), Xf.T @ (ytr.values - 0.5))
        return (Xte.fillna(0.5).values - 0.5) @ beta
    if model == "ic_weight":
        # sign-and-size each feature by its training-period mean rank IC
        d = Xtr.index.get_level_values(0)
        ics = {}
        for c in Xtr.columns:
            df = pd.DataFrame({"x": Xtr[c], "y": ytr, "d": d}).dropna()
            ics[c] = df.groupby("d").apply(lambda g: g["x"].corr(g["y"])).mean()
        w = pd.Series(ics).fillna(0.0)
        return (Xte.fillna(0.5) - 0.5).values @ w.values
    raise ValueError(model)