#!/usr/bin/env python3
"""Train LightGBM on GP-mined features."""

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import pandas as pd

from models.lightgbm_model import prepare_ml_matrix, train_lightgbm


def main():
    feature_path = ROOT / "outputs" / "gp_mining" / "qlib_gp_run_0" / "ML_Features_qlib.csv"
    if not feature_path.exists():
        feature_path = feature_path.with_suffix(".parquet")
    if not feature_path.exists():
        raise FileNotFoundError("Run GP mining first: python scripts/run_gp_mining.py")

    df = pd.read_parquet(feature_path) if feature_path.suffix == ".parquet" else pd.read_csv(feature_path)
    df["date"] = pd.to_datetime(df["date"])

    work, feature_cols = prepare_ml_matrix(df)
    train_df = work[work["date"] < "2018-01-01"]
    valid_df = work[(work["date"] >= "2018-01-01") & (work["date"] < "2019-04-01")]
    test_df = work[work["date"] >= "2019-04-01"]

    model = train_lightgbm(
        train_df,
        valid_df,
        feature_cols,
        model_path=ROOT / "outputs" / "models" / "lightgbm_gp.txt",
    )

    test_pred = model.predict(test_df[feature_cols])
    test_df = test_df.copy()
    test_df["pred_score"] = test_pred
    test_df[["date", "symbol", "pred_score", "target_return"]].to_csv(
        ROOT / "outputs" / "models" / "test_predictions.csv",
        index=False,
    )
    print("Test predictions saved.")


if __name__ == "__main__":
    main()