"""LocalFactorEvaluator - FactorEngine 执行封装。

封装本地 FactorEngine，提供与 FactorMiner compute_tree_signals 兼容的输出接口，
用于在评估管线中计算因子信号。

Features:
    - 封装 FactorEngine，内建数据路由读取 pro_bar 表
    - 输入 (name, formula) 列表，输出 {name: (M,T) np.ndarray}
    - 支持批量计算和单个因子计算
    - 自动计算收益率矩阵用于 IC 分析
"""

from __future__ import annotations

import re
from typing import Dict, List, Optional, Tuple

import numpy as np
import polars as pl

from src.factors import FactorEngine
from src.factorminer.evaluation.stock_pool_registry import StockPoolRegistry


class LocalFactorEvaluator:
    """本地因子评估器 - FactorEngine 封装。

    封装 FactorEngine，提供与 FactorMiner 评估管线兼容的接口，
    直接利用 FactorEngine 内建的数据路由读取 pro_bar 表，
    无需外部数据加载器。

    Attributes:
        start_date: 计算开始日期 (YYYYMMDD)
        end_date: 计算结束日期 (YYYYMMDD)
        stock_codes: 股票代码列表，None 表示全量
        engine: FactorEngine 实例
    """

    def __init__(
        self,
        start_date: str,
        end_date: str,
        stock_codes: Optional[List[str]] = None,
        stock_pool_registry: Optional[StockPoolRegistry] = None,
    ) -> None:
        """初始化评估器。

        Args:
            start_date: 计算开始日期，YYYYMMDD 格式
            end_date: 计算结束日期，YYYYMMDD 格式
            stock_codes: 可选的股票代码列表，None 表示全量
            stock_pool_registry: 股票池注册表，用于多股票池评估
        """
        self.start_date = start_date
        self.end_date = end_date
        self.stock_codes = stock_codes
        self.engine = FactorEngine()
        self.stock_pool_registry = stock_pool_registry
        self._asset_codes: Optional[List[str]] = None

    def evaluate(
        self,
        specs: List[Tuple[str, str]],
    ) -> Dict[str, np.ndarray]:
        """批量计算并返回 {name: (M, T) 矩阵}。

        Args:
            specs: (因子名, 本地 DSL 公式) 列表

        Returns:
            每个因子对应的 (asset, time) numpy 矩阵，缺失值填充 np.nan
        """
        if not specs:
            return {}

        if len(specs) > 1:
            print(f"[local_engine] 开始批量计算 {len(specs)} 个因子...")

        # 注册所有因子（防御性清洗：去除 LLM 可能返回的编号前缀如 "8. ")
        numbered_prefix_pattern = re.compile(r"^\s*\d+[\.\)]\s*")
        for name, formula in specs:
            cleaned_formula = numbered_prefix_pattern.sub("", formula)
            # 将 returns / $returns 替换为本地可用的 DSL 表达式
            cleaned_formula = re.sub(
                r"(?<!\w)\$returns(?!\w)",
                "(close / ts_delay(close, 1) - 1)",
                cleaned_formula,
            )
            cleaned_formula = re.sub(
                r"(?<!\w)returns(?!\w)",
                "(close / ts_delay(close, 1) - 1)",
                cleaned_formula,
            )
            try:
                self.engine.add_factor(name, cleaned_formula)
            except Exception as e:
                print(f"[ERROR] 注册因子 {name} 失败: {e}")
                raise

        # 批量计算
        factor_names = [name for name, _ in specs]
        try:
            result_df = self.engine.compute(
                factor_names=factor_names,
                start_date=self.start_date,
                end_date=self.end_date,
                stock_codes=self.stock_codes,
            )
        except Exception as e:
            print(f"[ERROR] 因子计算失败: {e}")
            raise

        # 转换为 (M, T) 矩阵
        signals_dict = self._pivot_to_matrix(result_df, factor_names)

        # 清理注册的因子
        self.engine.clear()

        if len(specs) > 1:
            print(f"[local_engine] 批量计算完成，返回 {len(signals_dict)} 个因子")
        return signals_dict

    def evaluate_single(
        self,
        name: str,
        formula: str,
    ) -> np.ndarray:
        """计算单个因子。

        Args:
            name: 因子名称
            formula: 本地 DSL 公式

        Returns:
            (M, T) 的因子信号矩阵
        """
        result = self.evaluate([(name, formula)])
        if name in result:
            return result[name]
        raise ValueError(f"因子 {name} 计算失败或返回为空")

    def evaluate_returns(
        self,
        periods: int = 1,
    ) -> np.ndarray:
        """计算收益率矩阵，用于后续 IC / quintile 分析。

        Args:
            periods: 计算 N 日后的收益率

        Returns:
            (M, T) 的 forward returns 矩阵
        """
        # 使用 DSL 计算收益率
        formula = f"close / ts_delay(close, {periods}) - 1"

        try:
            self.engine.add_factor("__returns_tmp", formula)
            result_df = self.engine.compute(
                factor_names=["__returns_tmp"],
                start_date=self.start_date,
                end_date=self.end_date,
                stock_codes=self.stock_codes,
            )
            self.engine.clear()

            # 转换为矩阵
            returns_dict = self._pivot_to_matrix(result_df, ["__returns_tmp"])
            return returns_dict.get("__returns_tmp", np.array([]))
        except Exception as e:
            print(f"[ERROR] 计算收益率矩阵失败: {e}")
            raise

    def get_asset_codes(self) -> List[str]:
        """返回最近一次 evaluate/evaluate_returns 中的资产代码列表。

        Returns:
            按字母序排列的股票代码列表 (M,)。
        """
        if self._asset_codes is None:
            raise RuntimeError("请先调用 evaluate() 或 evaluate_returns()")
        return self._asset_codes

    def _get_metadata_df(self, columns: List[str]) -> Optional[pl.DataFrame]:
        """拉取构建股票池所需的元数据（默认取 end_date 最新截面）。

        Args:
            columns: 需要的额外列名列表。

        Returns:
            包含 ts_code 和 columns 的 DataFrame，若失败则返回 None。
        """
        if not columns:
            return None
        try:
            df = self.engine.router._load_table(
                table_name="daily_basic",
                columns=columns,
                start_date=self.end_date,
                end_date=self.end_date,
                stock_codes=self.stock_codes,
            )
            if "trade_date" in df.columns:
                df = df.sort("trade_date", descending=True)
            df = df.unique(subset=["ts_code"], maintain_order=True)
            return df.select(["ts_code"] + columns)
        except Exception as exc:
            print(f"[WARN] 拉取股票池元数据失败: {exc}")
            return None

    def evaluate_returns_by_pool(self, periods: int = 1) -> Dict[str, np.ndarray]:
        """计算各股票池的收益率矩阵。

        先计算全市场收益率矩阵，然后根据注册的股票池掩码生成子池矩阵。
        未入选该池的资产在所有时间上的收益率标记为 NaN。

        Args:
            periods: 计算 N 日后的收益率，默认 1。

        Returns:
            {pool_name: (M, T) np.ndarray} 字典，包含所有注册的股票池和 "all"。
        """
        returns_all = self.evaluate_returns(periods=periods)
        result: Dict[str, np.ndarray] = {"all": returns_all}

        if self.stock_pool_registry is None:
            return result

        codes = self.get_asset_codes()
        req_cols = self.stock_pool_registry.get_required_columns()
        metadata = self._get_metadata_df(req_cols)
        self.stock_pool_registry.build_masks(codes, metadata_df=metadata)

        for name in self.stock_pool_registry.get_pool_names():
            if name == "all":
                continue
            mask = self.stock_pool_registry.masks[name]
            pool_returns = np.full_like(returns_all, np.nan, dtype=np.float64)
            pool_returns[mask, :] = returns_all[mask, :]
            result[name] = pool_returns

        return result

    def _pivot_to_matrix(
        self,
        df: pl.DataFrame,
        factor_names: List[str],
    ) -> Dict[str, np.ndarray]:
        """将 Polars DataFrame 透视为 {name: (M, T)} 字典。

        按 ts_code 字母序和 trade_date 时间序排列。

        Args:
            df: 包含 ts_code, trade_date 和因子列的 DataFrame
            factor_names: 要提取的因子名称列表

        Returns:
            {因子名: (M, T) numpy 矩阵} 字典
        """
        if len(df) == 0:
            return {name: np.array([]) for name in factor_names}

        # 确保日期排序
        df = df.sort(["trade_date", "ts_code"])

        # 获取时间戳和股票代码的唯一值（已排序）
        timestamps = df["trade_date"].unique().sort()
        asset_codes = df["ts_code"].unique().sort()
        if self._asset_codes is None:
            self._asset_codes = asset_codes.to_list()

        n_assets = len(asset_codes)
        n_times = len(timestamps)

        # 创建 timestamp 到索引的映射
        ts_to_idx = {ts: i for i, ts in enumerate(timestamps)}
        asset_to_idx = {code: i for i, code in enumerate(asset_codes)}

        # 初始化结果字典
        result: Dict[str, np.ndarray] = {}

        for name in factor_names:
            if name not in df.columns:
                result[name] = np.full((n_assets, n_times), np.nan)
                continue

            # 创建 (M, T) 矩阵并填充 NaN
            matrix = np.full((n_assets, n_times), np.nan)

            # 使用 Polars 的 pivot 操作
            try:
                pivot_df = df.select(["ts_code", "trade_date", name]).pivot(
                    values=name,
                    on="trade_date",
                    index="ts_code",
                    aggregate_function="first",
                )

                # 获取透视后的列顺序
                pivoted_timestamps = [c for c in pivot_df.columns if c != "ts_code"]

                for ts in pivoted_timestamps:
                    if ts in ts_to_idx:
                        col_idx = ts_to_idx[ts]
                        for row_idx, code in enumerate(pivot_df["ts_code"]):
                            if code in asset_to_idx:
                                asset_idx = asset_to_idx[code]
                                val = pivot_df[row_idx, ts]
                                if val is not None and not (
                                    isinstance(val, float) and np.isnan(val)
                                ):
                                    matrix[asset_idx, col_idx] = val
            except Exception:
                # fallback: 逐行遍历
                for row in df.iter_rows(named=True):
                    code = row["ts_code"]
                    ts = row["trade_date"]
                    if code in asset_to_idx and ts in ts_to_idx:
                        asset_idx = asset_to_idx[code]
                        time_idx = ts_to_idx[ts]
                        val = row.get(name)
                        if val is not None and not (
                            isinstance(val, float) and np.isnan(val)
                        ):
                            matrix[asset_idx, time_idx] = val

            result[name] = matrix

        return result
