"""TabM回归任务

TabM模型的回归训练任务实现。
"""

from typing import Dict, Any, Optional
from pathlib import Path

import numpy as np
import polars as pl
import matplotlib.pyplot as plt

from src.training.tasks.base import BaseTask
from src.training.components.models.tabm_model import TabMModel

# Type alias for model type
TabMModelType = TabMModel


class TabMRegressionTask(BaseTask):
    """TabM回归任务

    使用TabM模型进行回归训练，支持:
    - 内置集成训练（ensemble_size）
    - 早停机制
    - 训练曲线绘制

    Attributes:
        model_params: 模型参数字典
        label_name: 目标列名称
        model: TabMModel实例
    """

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

        Args:
            model_params: TabM模型参数，包含:
                - n_blocks: MLP层数
                - d_block: 每层神经元数
                - dropout: Dropout率
                - ensemble_size: 集成大小
                - batch_size: 批次大小
                - learning_rate: 学习率
                - weight_decay: 权重衰减
                - epochs: 训练轮数
                - early_stopping_patience: 早停耐心值
            label_name: 目标列名称
        """
        super().__init__(model_params, label_name)
        self.model_params = model_params
        self.label_name = label_name
        self.model = None  # type: Optional[TabMModelType]

    def prepare_labels(self, data: Dict[str, Dict]) -> Dict[str, Dict]:
        """准备标签

        回归任务不需要转换标签，直接返回原始数据。

        Args:
            data: 数据字典

        Returns:
            未修改的数据字典
        """
        # 回归任务：标签已经是连续值，无需转换
        return data

    def fit(
        self,
        train_data: Dict[str, Any],
        val_data: Dict[str, Any],
    ) -> None:
        """训练TabM模型

        Args:
            train_data: 训练数据字典，包含:
                - X: 特征DataFrame
                - y: 标签Series
            val_data: 验证数据字典，包含:
                - X: 特征DataFrame
                - y: 标签Series
        """
        print("\n[TabMRegressionTask] 开始训练...")

        # 创建模型实例
        self.model = TabMModel(self.model_params)

        # 提取训练数据
        X_train = train_data["X"]
        y_train = train_data["y"]
        X_val = val_data["X"]
        y_val = val_data["y"]

        # 支持 date_col 参数
        date_col = self.model_params.get("date_col", None)

        # 训练模型
        self.model.fit(X=X_train, y=y_train, eval_set=(X_val, y_val), date_col=date_col)

        print("[TabMRegressionTask] 训练完成")

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

        Args:
            test_data: 测试数据字典，包含:
                - X: 特征DataFrame

        Returns:
            预测结果数组
        """
        if self.model is None:
            raise RuntimeError("模型未训练，请先调用fit()")

        X_test = test_data["X"]
        return self.model.predict(X_test)

    def get_model(self) -> Any:
        """获取训练好的模型

        Returns:
            TabMModel实例或None
        """
        return self.model

    def plot_training_metrics(self, output_path: Optional[str] = None) -> None:
        """绘制训练指标

        绘制训练和验证损失曲线。

        Args:
            output_path: 图表保存路径，None则显示图表
        """
        if self.model is None or not self.model.training_history_["train_loss"]:
            print("[TabMRegressionTask] 无训练历史可绘制")
            return

        history = self.model.training_history_

        fig, ax = plt.subplots(figsize=(10, 6))

        epochs = range(1, len(history["train_loss"]) + 1)
        ax.plot(epochs, history["train_loss"], "b-", label="Train Loss", linewidth=2)

        if history["val_loss"]:
            ax.plot(epochs, history["val_loss"], "r-", label="Val Loss", linewidth=2)

        ax.set_xlabel("Epoch", fontsize=12)
        ax.set_ylabel("Loss (MSE)", fontsize=12)
        ax.set_title("TabM Training History", fontsize=14)
        ax.legend(fontsize=10)
        ax.grid(True, alpha=0.3)

        plt.tight_layout()

        if output_path:
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_path, dpi=150, bbox_inches="tight")
            print(f"[TabMRegressionTask] 训练曲线保存到: {output_path}")
        else:
            plt.show()

        plt.close()
