"""回归任务实现

实现回归任务的训练流程：
- Label 无需转换（保持连续值）
- 使用 LightGBM 回归模型
- 支持 MAE/RMSE 评估
"""

from typing import Any, Dict, Optional
import numpy as np
import polars as pl

from src.training.tasks.base import BaseTask
from src.training.components.models.lightgbm import LightGBMModel


class RegressionTask(BaseTask):
    """回归任务

    使用 LightGBM 进行回归训练，支持早停和训练曲线。
    """

    def __init__(
        self,
        model_params: Dict[str, Any],
        label_name: str = "future_return_5",
    ):
        """初始化回归任务

        Args:
            model_params: LightGBM 参数字典
            label_name: Label 列名
        """
        super().__init__(model_params, label_name)
        self.evals_result: Optional[Dict] = None

    def prepare_labels(self, data: Dict[str, Dict]) -> Dict[str, Dict]:
        """准备标签（回归任务无需转换）

        Args:
            data: 数据字典

        Returns:
            原样返回数据字典
        """
        # 回归任务不需要转换 Label
        return data

    def fit(self, train_data: Dict, val_data: Dict) -> None:
        """训练回归模型

        Args:
            train_data: 训练数据 {"X": DataFrame, "y": Series}
            val_data: 验证数据
        """
        self.model = LightGBMModel(params=self.model_params)

        X_train = train_data["X"]
        y_train = train_data["y"]
        X_val = val_data["X"]
        y_val = val_data["y"]

        self.model.fit(
            X_train, y_train, eval_set=(X_val, y_val) if X_val is not None else None
        )

    def predict(self, test_data: Dict) -> np.ndarray:
        """生成预测

        Args:
            test_data: 测试数据

        Returns:
            预测结果
        """
        return self.model.predict(test_data["X"])

    def plot_training_metrics(self) -> None:
        """绘制训练指标曲线"""
        if self.model and hasattr(self.model, "get_evals_result"):
            try:
                import lightgbm as lgb

                evals_result = self.model.get_evals_result()
                if evals_result:
                    lgb.plot_metric(evals_result)
                else:
                    print("[警告] 没有训练指标数据可供绘制")
            except Exception as e:
                print(f"[警告] 无法绘制训练曲线: {e}")
