"""数据处理器实现

包含标准化、缩尾、缺失值填充等数据处理器。
"""

import math
from typing import List, Literal, Optional, Union

import polars as pl

from src.training.components.base import BaseProcessor
from src.training.registry import register_processor


@register_processor("null_filler")
class NullFiller(BaseProcessor):
    """缺失值填充处理器

    支持多种填充策略：固定值、0、均值、中值。
    可以全局填充或使用当天截面统计量填充。

    填充策略：
        - "zero": 填充0
        - "mean": 填充均值（全局或当天截面）
        - "median": 填充中值（全局或当天截面）
        - "value": 填充指定数值

    Attributes:
        strategy: 填充策略，可选 "zero", "mean", "median", "value"
        fill_value: 当 strategy="value" 时使用的填充值
        by_date: 是否按日期独立计算统计量（仅对 mean/median 有效）
        date_col: 日期列名
        feature_cols: 参与填充的特征列名列表
        stats_: 存储学习到的统计量（全局模式）
    """

    name = "null_filler"

    def __init__(
        self,
        feature_cols: List[str],
        strategy: Literal["zero", "mean", "median", "value"] = "zero",
        fill_value: Optional[float] = None,
        by_date: bool = True,
        date_col: str = "trade_date",
    ):
        """初始化缺失值填充处理器

        Args:
            feature_cols: 参与填充的特征列名列表
            strategy: 填充策略，默认 "zero"
                - "zero": 填充0
                - "mean": 填充均值
                - "median": 填充中值
                - "value": 填充指定数值（需配合 fill_value）
            fill_value: 当 strategy="value" 时的填充值，默认为 None
            by_date: 是否每天独立计算统计量，默认 True（截面统计量）
            date_col: 日期列名，默认 "trade_date"

        Raises:
            ValueError: 策略无效或 fill_value 未提供时
        """
        if strategy not in ("zero", "mean", "median", "value"):
            raise ValueError(
                f"无效的填充策略: {strategy}，必须是 'zero', 'mean', 'median', 'value' 之一"
            )

        if strategy == "value" and fill_value is None:
            raise ValueError("当 strategy='value' 时，必须提供 fill_value")

        self.feature_cols = feature_cols
        self.strategy = strategy
        self.fill_value = fill_value
        self.by_date = by_date
        self.date_col = date_col
        self.stats_ = {}

    def fit(self, X: pl.DataFrame) -> "NullFiller":
        """学习统计量（仅在全局模式下）

        在全局模式下，计算每列的均值或中值作为填充值。
        在截面模式下（by_date=True），不需要 fit，每天独立计算。

        Args:
            X: 训练数据

        Returns:
            self
        """
        if not self.by_date and self.strategy in ("mean", "median"):
            for col in self.feature_cols:
                if col in X.columns and (
                    X[col].dtype.is_numeric() or X[col].dtype == pl.Boolean
                ):
                    if self.strategy == "mean":
                        self.stats_[col] = X[col].mean() or 0.0
                    else:  # median
                        self.stats_[col] = X[col].median() or 0.0

        return self

    def transform(self, X: pl.DataFrame) -> pl.DataFrame:
        """填充缺失值

        Args:
            X: 待转换数据

        Returns:
            填充后的数据
        """
        if self.strategy == "zero":
            return self._fill_with_zero(X)
        elif self.strategy == "value":
            return self._fill_with_value(X)
        elif self.strategy in ("mean", "median"):
            if self.by_date:
                return self._fill_by_date(X)
            else:
                return self._fill_global(X)
        else:
            # 不应该到达这里，因为 __init__ 已经验证
            raise ValueError(f"未知的填充策略: {self.strategy}")

    def _fill_with_zero(self, X: pl.DataFrame) -> pl.DataFrame:
        """使用0填充缺失值（同时处理 NaN 和 null）"""
        expressions = []
        for col in X.columns:
            if col in self.feature_cols and (
                X[col].dtype.is_numeric() or X[col].dtype == pl.Boolean
            ):
                # 先 fill_nan 再 fill_null，确保两种缺失值都被处理
                expr = pl.col(col).fill_nan(0).fill_null(0).alias(col)
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))

        return X.select(expressions)

    def _fill_with_value(self, X: pl.DataFrame) -> pl.DataFrame:
        """使用指定值填充缺失值（同时处理 NaN 和 null）"""
        expressions = []
        for col in X.columns:
            if col in self.feature_cols and (
                X[col].dtype.is_numeric() or X[col].dtype == pl.Boolean
            ):
                # 先 fill_nan 再 fill_null
                expr = (
                    pl.col(col)
                    .fill_nan(self.fill_value)
                    .fill_null(self.fill_value)
                    .alias(col)
                )
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))

        return X.select(expressions)

    def _fill_global(self, X: pl.DataFrame) -> pl.DataFrame:
        """使用全局统计量填充（训练集学到的统计量，同时处理 NaN 和 null）"""
        expressions = []
        for col in X.columns:
            if col in self.stats_:
                fill_val = self.stats_[col]
                # 先 fill_nan 再 fill_null
                expr = pl.col(col).fill_nan(fill_val).fill_null(fill_val).alias(col)
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))

        return X.select(expressions)

    def _fill_by_date(self, X: pl.DataFrame) -> pl.DataFrame:
        """使用每天截面统计量填充（同时处理 NaN 和 null）"""
        # 确定需要处理的列（仅 numeric 类型，排除 boolean）
        # 注意：boolean 类型没有 NaN 概念，fill_nan 会报错
        target_cols = [
            col
            for col in self.feature_cols
            if col in X.columns and X[col].dtype.is_numeric()
        ]

        # 计算每天的统计量
        stat_exprs = []
        for col in target_cols:
            if self.strategy == "mean":
                stat_exprs.append(
                    pl.col(col).mean().over(self.date_col).alias(f"{col}_stat")
                )
            else:  # median
                stat_exprs.append(
                    pl.col(col).median().over(self.date_col).alias(f"{col}_stat")
                )

        # 添加统计量列
        result = X.with_columns(stat_exprs)

        # 使用统计量填充缺失值
        # 注意：如果某天某列全为null，统计量也会为null，所以需要链式填充
        # 同时处理 NaN 和 null
        fill_exprs = []
        for col in X.columns:
            if col in target_cols:
                # 先用当天统计量填充 NaN 和 null，如果统计量也是null则用0填充
                expr = (
                    pl.col(col)
                    .fill_nan(pl.col(f"{col}_stat"))
                    .fill_null(pl.col(f"{col}_stat"))
                    .fill_nan(0)  # 如果统计量是 NaN，再用 0 填充
                    .fill_null(0)  # 如果统计量是 null，再用 0 填充
                    .alias(col)
                )
                fill_exprs.append(expr)
            else:
                fill_exprs.append(pl.col(col))

        result = result.select(fill_exprs)

        return result


@register_processor("standard_scaler")
class StandardScaler(BaseProcessor):
    """标准化处理器（全局标准化）

    在整个训练集上学习均值和标准差，
    然后应用到训练集和测试集。

    适用于需要全局统计量的场景。

    Attributes:
        feature_cols: 参与标准化的特征列名列表
        mean_: 学习到的均值字典 {列名: 均值}
        std_: 学习到的标准差字典 {列名: 标准差}
    """

    name = "standard_scaler"

    def __init__(self, feature_cols: List[str]):
        """初始化标准化处理器

        Args:
            feature_cols: 参与标准化的特征列名列表
        """
        self.feature_cols = feature_cols
        self.mean_ = {}
        self.std_ = {}

    def fit(self, X: pl.DataFrame) -> "StandardScaler":
        """计算均值和标准差（仅在训练集上）

        Args:
            X: 训练数据

        Returns:
            self
        """
        for col in self.feature_cols:
            # 仅处理数值类型，排除布尔类型（标准化布尔类型语义不明确）
            if col in X.columns and X[col].dtype.is_numeric():
                col_mean = X[col].mean()
                col_std = X[col].std()
                # 关键修复：检查是否为 None 且不是 NaN
                # 注意：使用 try-except 处理类型转换，避免 LSP 类型检查错误
                try:
                    mean_is_valid = (
                        col_mean is not None
                        and isinstance(col_mean, (int, float))
                        and not math.isnan(col_mean)
                    )
                    std_is_valid = (
                        col_std is not None
                        and isinstance(col_std, (int, float))
                        and not math.isnan(col_std)
                    )
                except (TypeError, ValueError):
                    mean_is_valid = False
                    std_is_valid = False

                if mean_is_valid and std_is_valid:
                    self.mean_[col] = col_mean
                    self.std_[col] = col_std
                else:
                    # 如果统计量无效，使用默认值（mean=0, std=1）
                    # 防止 transform 时产生更多 NaN
                    self.mean_[col] = 0.0
                    self.std_[col] = 1.0

        return self

    def transform(self, X: pl.DataFrame) -> pl.DataFrame:
        """标准化（使用训练集学到的参数）

        Args:
            X: 待转换数据

        Returns:
            标准化后的数据
        """
        expressions = []
        for col in X.columns:
            if col in self.mean_ and col in self.std_:
                # 避免除以0
                std_val = self.std_[col] if self.std_[col] != 0 else 1.0
                expr = ((pl.col(col) - self.mean_[col]) / std_val).alias(col)
                expressions.append(expr)
            elif col in self.feature_cols:
                # 对于应该被处理但未学习到统计量的列，统一转换为float
                expr = pl.col(col).cast(pl.Float64).alias(col)
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))

        return X.select(expressions)


@register_processor("cs_standard_scaler")
class CrossSectionalStandardScaler(BaseProcessor):
    """截面标准化处理器

    每天独立进行标准化：对当天所有股票的某一因子进行标准化。

    特点：
        - 不需要 fit，每天独立计算当天的均值和标准差
        - 适用于截面因子，消除市值等行业差异
        - 公式：z = (x - mean_today) / std_today

    Attributes:
        feature_cols: 参与标准化的特征列名列表
        date_col: 日期列名
    """

    name = "cs_standard_scaler"

    def __init__(
        self,
        feature_cols: List[str],
        date_col: str = "trade_date",
    ):
        """初始化截面标准化处理器

        Args:
            feature_cols: 参与标准化的特征列名列表
            date_col: 日期列名
        """
        self.feature_cols = feature_cols
        self.date_col = date_col

    def transform(self, X: pl.DataFrame) -> pl.DataFrame:
        """截面标准化

        按日期分组，每天独立计算均值和标准差并标准化。
        不需要 fit，因为每天使用当天的统计量。

        Args:
            X: 待转换数据

        Returns:
            标准化后的数据
        """
        # 构建表达式列表
        expressions = []
        for col in X.columns:
            # 仅处理数值类型，排除布尔类型（标准化布尔类型语义不明确）
            if col in self.feature_cols and X[col].dtype.is_numeric():
                # 截面标准化：每天独立计算均值和标准差
                # 避免除以0，当std为0时设为1
                expr = (
                    (pl.col(col) - pl.col(col).mean().over(self.date_col))
                    / (pl.col(col).std().over(self.date_col) + 1e-10)
                ).alias(col)
                expressions.append(expr)
            elif col in self.feature_cols:
                # 对于应该被处理但类型不匹配的列，转换为float
                expr = pl.col(col).cast(pl.Float64).alias(col)
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))

        return X.select(expressions)


@register_processor("winsorizer")
class Winsorizer(BaseProcessor):
    """缩尾处理器

    对每一列的极端值进行截断处理。
    可以全局截断（在整个训练集上学习分位数），
    也可以截面截断（每天独立处理）。

    Attributes:
        feature_cols: 参与缩尾的特征列名列表
        lower: 下分位数（如0.01表示1%分位数）
        upper: 上分位数（如0.99表示99%分位数）
        by_date: True=每天独立缩尾, False=全局缩尾
        date_col: 日期列名
        bounds_: 存储分位数边界（全局模式）
    """

    name = "winsorizer"

    def __init__(
        self,
        feature_cols: List[str],
        lower: float = 0.01,
        upper: float = 0.99,
        by_date: bool = False,
        date_col: str = "trade_date",
    ):
        """初始化缩尾处理器

        Args:
            feature_cols: 参与缩尾的特征列名列表
            lower: 下分位数，默认0.01
            upper: 上分位数，默认0.99
            by_date: 每天独立缩尾，默认False（全局缩尾）
            date_col: 日期列名

        Raises:
            ValueError: 分位数参数无效
        """
        if not 0 <= lower < upper <= 1:
            raise ValueError(
                f"lower ({lower}) 必须小于 upper ({upper})，且都在 [0, 1] 范围内"
            )

        self.feature_cols = feature_cols
        self.lower = lower
        self.upper = upper
        self.by_date = by_date
        self.date_col = date_col
        self.bounds_ = {}

    def fit(self, X: pl.DataFrame) -> "Winsorizer":
        """学习分位数边界（仅在全局模式下）

        Args:
            X: 训练数据

        Returns:
            self
        """
        if not self.by_date:
            for col in self.feature_cols:
                # 仅处理数值类型，排除布尔类型（quantile 不支持布尔类型）
                if col in X.columns and X[col].dtype.is_numeric():
                    self.bounds_[col] = {
                        "lower": X[col].quantile(self.lower),
                        "upper": X[col].quantile(self.upper),
                    }
        return self

    def transform(self, X: pl.DataFrame) -> pl.DataFrame:
        """缩尾处理

        Args:
            X: 待转换数据

        Returns:
            缩尾处理后的数据
        """
        if self.by_date:
            return self._transform_by_date(X)
        else:
            return self._transform_global(X)

    def _transform_global(self, X: pl.DataFrame) -> pl.DataFrame:
        """全局缩尾（使用训练集学到的边界）"""
        expressions = []
        for col in X.columns:
            if col in self.bounds_:
                lower = self.bounds_[col]["lower"]
                upper = self.bounds_[col]["upper"]
                expr = pl.col(col).clip(lower, upper).alias(col)
                expressions.append(expr)
            elif col in self.feature_cols:
                # 对于应该被处理但未学习到边界的列（如全为NaN、布尔列等）
                # 统一转换为float
                expr = pl.col(col).cast(pl.Float64).alias(col)
                expressions.append(expr)
            else:
                expressions.append(pl.col(col))
        return X.select(expressions)

    def _transform_by_date(self, X: pl.DataFrame) -> pl.DataFrame:
        """每日独立缩尾"""
        # 确定需要处理的列（仅 numeric 类型，排除 boolean）
        # 注意：quantile 操作不支持布尔类型
        target_cols = [
            col
            for col in self.feature_cols
            if col in X.columns and X[col].dtype.is_numeric()
        ]

        # 先计算每天的分位数
        lower_exprs = [
            pl.col(col).quantile(self.lower).over(self.date_col).alias(f"{col}_lower")
            for col in target_cols
        ]
        upper_exprs = [
            pl.col(col).quantile(self.upper).over(self.date_col).alias(f"{col}_upper")
            for col in target_cols
        ]

        # 添加分位数列
        result = X.with_columns(lower_exprs + upper_exprs)

        # 执行缩尾
        clip_exprs = []
        for col in X.columns:
            if col in target_cols:
                clipped = (
                    pl.col(col)
                    .clip(pl.col(f"{col}_lower"), pl.col(f"{col}_upper"))
                    .alias(col)
                )
                clip_exprs.append(clipped)
            else:
                clip_exprs.append(pl.col(col))

        result = result.select(clip_exprs)

        return result
