"""Model training on GP-mined or qlib features."""

from __future__ import annotations

from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error


def prepare_ml_matrix(df: pd.DataFrame, feature_prefix: str = "factor_", target_col: str = "target_return"):
    feature_cols = [c for c in df.columns if c.startswith(feature_prefix)]
    work = df.copy()
    for col in feature_cols:
        work[col] = work.groupby("date")[col].transform(lambda x: x.rank(pct=True))
    work["target_rank"] = work.groupby("date")[target_col].transform(lambda x: x.rank(pct=True))
    work = work.replace([np.inf, -np.inf], np.nan).dropna()
    return work, feature_cols


def train_lightgbm(
    train_df: pd.DataFrame,
    valid_df: pd.DataFrame,
    feature_cols: list[str],
    target_col: str = "target_rank",
    model_path: Path | None = None,
) -> lgb.Booster:
    train_set = lgb.Dataset(train_df[feature_cols], label=train_df[target_col])
    valid_set = lgb.Dataset(valid_df[feature_cols], label=valid_df[target_col], reference=train_set)

    params = {
        "objective": "regression",
        "metric": "mse",
        "learning_rate": 0.05,
        "num_leaves": 31,
        "feature_fraction": 0.8,
        "bagging_fraction": 0.8,
        "bagging_freq": 5,
        "verbose": -1,
        "seed": 42,
    }

    model = lgb.train(
        params,
        train_set,
        num_boost_round=500,
        valid_sets=[valid_set],
        callbacks=[lgb.early_stopping(50, verbose=False)],
    )

    if model_path is not None:
        model_path.parent.mkdir(parents=True, exist_ok=True)
        model.save_model(str(model_path))

    valid_pred = model.predict(valid_df[feature_cols])
    rmse = mean_squared_error(valid_df[target_col], valid_pred, squared=False)
    print(f"LightGBM valid RMSE: {rmse:.6f}")
    return model