"""结果分析器

训练后的分析和结果处理：
1. 特征重要性分析（Top N、零贡献特征）
2. 结果组装（生成每日 Top N）
3. 结果保存
"""

from typing import Any, Dict, List, Optional
import os
import polars as pl
import pandas as pd
import numpy as np


class ResultAnalyzer:
    """结果分析器

    分析训练结果，生成报告并保存。
    """

    def analyze_feature_importance(
        self,
        model,
        feature_cols: List[str],
        top_n: int = 20,
        verbose: bool = True,
    ) -> Dict[str, Any]:
        """分析特征重要性

        Args:
            model: 训练好的模型
            feature_cols: 特征列名列表
            top_n: 显示 Top N 特征
            verbose: 是否打印信息

        Returns:
            分析结果字典
        """
        importance = model.feature_importance()

        if importance is None:
            if verbose:
                print("[警告] 无法获取特征重要性")
            return {}

        # 按重要性排序
        importance_sorted = importance.sort_values(ascending=False)

        # 计算百分比
        total_importance = importance_sorted.sum()
        importance_pct = (importance_sorted / total_importance * 100).round(2)

        # 识别零贡献特征
        zero_importance_features = importance_sorted[
            importance_sorted == 0
        ].index.tolist()

        if verbose:
            print("\n" + "=" * 80)
            print("特征重要性分析")
            print("=" * 80)

            # 打印 Top N
            print(f"\nTop {top_n} 特征:")
            print("-" * 80)
            print(f"{'排名':<6}{'特征名':<35}{'重要性':<15}{'占比':<10}")
            print("-" * 80)

            for i, (feature, score) in enumerate(
                importance_sorted.head(top_n).items(), 1
            ):
                pct = importance_pct[feature]
                if pct >= 10:
                    marker = " [高贡献]"
                elif pct >= 1:
                    marker = " [中贡献]"
                else:
                    marker = " [低贡献]"
                print(f"{i:<6}{feature:<35}{score:<15.2f}{pct:<8.2f}%{marker}")

            # 打印零贡献特征
            if zero_importance_features:
                print("\n" + "-" * 80)
                print(f"[警告] 贡献为0的特征（共 {len(zero_importance_features)} 个）:")
                for i, feature in enumerate(zero_importance_features, 1):
                    print(f"'{feature}',")

            # 统计摘要
            print("\n" + "=" * 80)
            print("统计摘要:")
            print("-" * 80)
            print(f"  特征总数: {len(importance_sorted)}")
            print(
                f"  有贡献特征数: {len(importance_sorted) - len(zero_importance_features)}"
            )
            print(f"  零贡献特征数: {len(zero_importance_features)}")
            if len(importance_sorted) > 0:
                print(
                    f"  零贡献占比: {len(zero_importance_features) / len(importance_sorted) * 100:.1f}%"
                )
            print(f"  Top {top_n} 累计占比: {importance_pct.head(top_n).sum():.1f}%")
            print("=" * 80)

        return {
            "importance": importance_sorted,
            "importance_pct": importance_pct,
            "zero_importance_features": zero_importance_features,
            "top_n": importance_sorted.head(top_n),
        }

    def assemble_results(
        self,
        test_data: Dict[str, Any],
        predictions: np.ndarray,
        top_n: int = 50,
        verbose: bool = True,
    ) -> pl.DataFrame:
        """组装结果

        生成每日 Top N 股票推荐列表。

        Args:
            test_data: 测试数据字典
            predictions: 预测结果数组
            top_n: 每日选择的股票数
            verbose: 是否打印信息

        Returns:
            结果数据框
        """
        # 添加预测列
        raw_data = test_data["raw_data"]
        results = raw_data.with_columns([pl.Series("prediction", predictions)])

        # 按日期分组取 Top N
        unique_dates = results["trade_date"].unique().sort()
        topn_by_date = []

        for date in unique_dates:
            day_data = results.filter(results["trade_date"] == date)
            topn = day_data.sort("prediction", descending=True).head(top_n)
            topn_by_date.append(topn)

        # 合并所有日期的 Top N
        topn_results = pl.concat(topn_by_date)

        if verbose:
            print(f"\n生成每日 Top {top_n} 股票列表:")
            print(f"  交易日数: {len(unique_dates)}")
            print(f"  总推荐数: {len(topn_results)}")

        return topn_results

    def save_results(
        self,
        results: pl.DataFrame,
        output_path: str,
        verbose: bool = True,
    ) -> None:
        """保存结果

        Args:
            results: 结果数据框
            output_path: 输出路径
            verbose: 是否打印信息
        """
        # 格式化日期并调整列顺序
        formatted = results.select(
            [
                (
                    pl.col("trade_date").str.slice(0, 4)
                    + "-"
                    + pl.col("trade_date").str.slice(4, 2)
                    + "-"
                    + pl.col("trade_date").str.slice(6, 2)
                ).alias("date"),
                pl.col("prediction").alias("score"),
                pl.col("ts_code"),
            ]
        )

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

        # 保存 CSV
        formatted.write_csv(output_path, include_header=True)

        if verbose:
            print(f"  保存路径: {output_path}")
            print(f"  保存行数: {len(formatted)}")
