"""训练调度引擎

协调 FactorManager、DataPipeline、Task 和 ResultAnalyzer 完成训练流程。
"""

from typing import Any, Callable, Dict, List, Optional, Tuple
import os
from datetime import datetime

import polars as pl

from src.factors import FactorEngine
from src.training.pipeline import DataPipeline
from src.training.tasks.base import BaseTask
from src.training.result_analyzer import ResultAnalyzer
from src.training.data_quality_analyzer import DataQualityAnalyzer


class Trainer:
    """训练调度引擎

    协调各个组件执行完整训练流程：
    1. 准备数据（DataPipeline）
    2. 处理标签（Task）
    3. 训练模型（Task）
    4. 绘制指标（Task）
    5. 生成预测（Task）
    6. 分析结果（ResultAnalyzer）
    7. 保存结果

    Attributes:
        data_pipeline: 数据流水线
        task: 任务实例（RegressionTask/RankTask）
        analyzer: 结果分析器
        output_config: 输出配置
        verbose: 是否打印详细信息
        results: 训练结果
    """

    def __init__(
        self,
        data_pipeline: DataPipeline,
        task: BaseTask,
        analyzer: Optional[ResultAnalyzer] = None,
        output_config: Optional[Dict[str, Any]] = None,
        verbose: bool = True,
    ):
        """初始化训练器

        Args:
            data_pipeline: 数据流水线实例
            task: 任务实例（RegressionTask 或 RankTask）
            analyzer: 结果分析器（可选，默认创建新实例）
            output_config: 输出配置字典
            verbose: 是否打印详细信息
        """
        self.data_pipeline = data_pipeline
        self.task = task
        self.analyzer = analyzer or ResultAnalyzer()
        self.output_config = output_config or {}
        self.verbose = verbose
        self.results: Optional[pl.DataFrame] = None

    def run(
        self,
        engine: FactorEngine,
        date_range: Dict[str, Tuple[str, str]],
    ) -> pl.DataFrame:
        """执行完整训练流程

        Args:
            engine: FactorEngine 实例
            date_range: 日期范围字典
                {
                    "train": (start_date, end_date),
                    "val": (start_date, end_date),
                    "test": (start_date, end_date),
                }

        Returns:
            训练结果数据框
        """
        if self.verbose:
            print("\n" + "=" * 80)
            print(f"开始训练: {self.task.__class__.__name__}")
            print("=" * 80)

        # Step 1: 准备数据
        if self.verbose:
            print("\n[Step 1/7] 准备数据...")

        data = self.data_pipeline.prepare_data(
            engine=engine,
            date_range=date_range,
            label_name=self.task.label_name,
            verbose=self.verbose,
        )

        # Step 1.5: 数据质量分析
        if self.verbose:
            print("\n[Step 1.5/7] 数据质量分析...")

        try:
            # 获取特征列名（从训练集）
            feature_cols = data["train"].get("feature_cols", [])
            label_name = self.task.label_name

            analyzer = DataQualityAnalyzer(
                feature_cols=feature_cols,
                label_col=label_name,
                verbose=self.verbose,
            )
            analyzer.analyze(data)
        except Exception as e:
            if self.verbose:
                print(f"  [警告] 数据质量分析失败: {e}")

        # Step 2: 处理标签
        if self.verbose:
            print("\n[Step 2/7] 处理标签...")

        data = self.task.prepare_labels(data)

        # Step 3: 训练模型
        if self.verbose:
            print("\n[Step 3/7] 训练模型...")

        self.task.fit(data["train"], data["val"])

        # Step 4: 绘制训练指标
        if self.verbose:
            print("\n[Step 4/7] 绘制训练指标...")

        self.task.plot_training_metrics()

        # Step 5: 生成预测
        if self.verbose:
            print("\n[Step 5/7] 生成预测...")

        predictions = self.task.predict(data["test"])

        # Step 6: 分析结果
        if self.verbose:
            print("\n[Step 6/7] 分析结果...")

        # 特征重要性
        self.analyzer.analyze_feature_importance(
            model=self.task.get_model(),
            feature_cols=data["test"]["feature_cols"],
            top_n=20,
            verbose=self.verbose,
        )

        # NDCG 评估（排序任务特有）
        if hasattr(self.task, "evaluate_ndcg"):
            ndcg_scores = self.task.evaluate_ndcg(data["test"])
            if self.verbose:
                print("\nNDCG 评估结果:")
                for metric, score in ndcg_scores.items():
                    print(f"  {metric}: {score:.4f}")

        # 组装结果
        self.results = self.analyzer.assemble_results(
            test_data=data["test"],
            predictions=predictions,
            top_n=self.output_config.get("top_n", 50),
            verbose=self.verbose,
        )

        # Step 7: 保存结果
        if self.verbose:
            print("\n[Step 7/7] 保存结果...")

        if self.output_config.get("save_predictions", True):
            self._save_predictions()

        if self.output_config.get("save_model", False):
            self._save_model()

        if self.verbose:
            print("\n" + "=" * 80)
            print("训练完成！")
            print("=" * 80)

        return self.results

    def _save_predictions(self) -> None:
        """保存预测结果"""
        output_dir = self.output_config.get("output_dir", "experiment/output")
        output_filename = self.output_config.get("output_filename", "output.csv")
        output_path = os.path.join(output_dir, output_filename)

        self.analyzer.save_results(
            results=self.results,
            output_path=output_path,
            verbose=self.verbose,
        )

    def _save_model(self) -> None:
        """保存模型"""
        model_save_path = self.output_config.get("model_save_path")
        if not model_save_path:
            return

        # 确保目录存在
        os.makedirs(os.path.dirname(model_save_path), exist_ok=True)

        # 获取模型和相关信息
        model = self.task.get_model()

        # 保存模型
        model.save(model_save_path)

        if self.verbose:
            print(f"  模型保存路径: {model_save_path}")

    def get_results(self) -> Optional[pl.DataFrame]:
        """获取训练结果

        Returns:
            训练结果数据框，如果尚未训练则返回 None
        """
        return self.results

    def get_task(self) -> BaseTask:
        """获取任务实例

        Returns:
            任务实例
        """
        return self.task
