"""训练模块工具函数

提供数据质量检查、验证等通用工具函数。
"""

from typing import Dict, List, Optional, Union

import polars as pl


def check_data_quality(
    df: pl.DataFrame,
    feature_cols: List[str],
    date_col: str = "trade_date",
    check_all_null: bool = True,
    check_all_zero: bool = True,
    check_all_nan: bool = True,
    raise_on_error: bool = True,
) -> Dict[str, List[Dict[str, Union[str, int]]]]:
    """检查数据质量，检测某天某个因子是否全部为空、0或NaN

    此检查必须在 fillna、标准化等处理之前执行，否则错误会被掩盖。

    Args:
        df: 待检查的数据
        feature_cols: 特征列名列表
        date_col: 日期列名，默认 "trade_date"
        check_all_null: 是否检查全空，默认 True
        check_all_zero: 是否检查全零，默认 True
        check_all_nan: 是否检查全NaN，默认 True
        raise_on_error: 发现问题时是否抛出异常，默认 True

    Returns:
        检查结果字典，格式：
        {
            "all_null": [{"date": "20240101", "factor": "factor_name", "count": 500}],
            "all_zero": [{"date": "20240101", "factor": "factor_name", "count": 500}],
            "all_nan": [{"date": "20240101", "factor": "factor_name", "count": 500}],
        }

    Raises:
        ValueError: 当发现质量问题且 raise_on_error=True 时

    Examples:
        >>> import polars as pl
        >>> df = pl.DataFrame({
        ...     "trade_date": ["20240101", "20240101", "20240102"],
        ...     "ts_code": ["000001.SZ", "000002.SZ", "000001.SZ"],
        ...     "factor1": [1.0, 2.0, None],
        ...     "factor2": [0.0, 0.0, 1.0],
        ... })
        >>> result = check_data_quality(df, ["factor1", "factor2"])
    """
    issues = {
        "all_null": [],
        "all_zero": [],
        "all_nan": [],
    }

    # 获取实际存在的特征列
    existing_cols = [col for col in feature_cols if col in df.columns]
    if not existing_cols:
        return issues

    # 按日期分组检查
    for date in df[date_col].unique():
        day_data = df.filter(pl.col(date_col) == date)
        day_str = str(date)

        for col in existing_cols:
            if not day_data[col].dtype.is_numeric():
                continue

            col_data = day_data[col]
            non_null_count = col_data.count()

            if non_null_count == 0:
                # 该日期该因子完全没有有效数据
                if check_all_null:
                    issues["all_null"].append(
                        {
                            "date": day_str,
                            "factor": col,
                            "count": len(day_data),
                        }
                    )
                continue

            # 检查是否全为零
            if check_all_zero:
                abs_sum = col_data.abs().sum()
                if abs_sum == 0:
                    issues["all_zero"].append(
                        {
                            "date": day_str,
                            "factor": col,
                            "count": non_null_count,
                        }
                    )

            # 检查是否全为NaN（在Polars中表现为null）
            if check_all_nan:
                null_count = col_data.null_count()
                if null_count == non_null_count:
                    issues["all_nan"].append(
                        {
                            "date": day_str,
                            "factor": col,
                            "count": non_null_count,
                        }
                    )

    # 生成报告
    total_issues = sum(len(v) for v in issues.values())

    if total_issues > 0:
        report_lines = ["\n" + "=" * 80, "数据质量检查报告", "=" * 80]

        if issues["all_null"]:
            report_lines.append(f"\n[严重] 发现 {len(issues['all_null'])} 个全空因子:")
            report_lines.append(
                "  (某天的某个因子所有值都是 null，可能是数据缺失或计算错误)"
            )
            for issue in issues["all_null"][:10]:  # 最多显示10条
                msg = f"    - 日期 {issue['date']}: {issue['factor']} (样本数: {issue['count']})"
                report_lines.append(msg)
            if len(issues["all_null"]) > 10:
                report_lines.append(f"    ... 还有 {len(issues['all_null']) - 10} 个")

        if issues["all_zero"]:
            report_lines.append(f"\n[警告] 发现 {len(issues['all_zero'])} 个全零因子:")
            report_lines.append(
                "  (某天的某个因子所有值都是 0，可能是计算错误或数据源问题)"
            )
            for issue in issues["all_zero"][:10]:
                msg = f"    - 日期 {issue['date']}: {issue['factor']} (样本数: {issue['count']})"
                report_lines.append(msg)
            if len(issues["all_zero"]) > 10:
                report_lines.append(f"    ... 还有 {len(issues['all_zero']) - 10} 个")

        if issues["all_nan"]:
            report_lines.append(f"\n[警告] 发现 {len(issues['all_nan'])} 个全NaN因子:")
            report_lines.append("  (某天的某个因子所有值都是 NaN，可能是数值计算错误)")
            for issue in issues["all_nan"][:10]:
                msg = f"    - 日期 {issue['date']}: {issue['factor']} (样本数: {issue['count']})"
                report_lines.append(msg)
            if len(issues["all_nan"]) > 10:
                report_lines.append(f"    ... 还有 {len(issues['all_nan']) - 10} 个")

        report_lines.extend(
            [
                "\n" + "-" * 80,
                "建议处理方式:",
                "  1. 检查因子定义和数据源，确认计算逻辑是否正确",
                "  2. 如果是预期内的缺失（如新股无历史数据），考虑调整因子计算窗口",
                "  3. 如果是数据同步问题，重新同步相关数据",
                "  4. 可以使用 filter 排除问题日期或因子",
                "=" * 80,
            ]
        )

        report = "\n".join(report_lines)
        print(report)

        if raise_on_error:
            raise ValueError(
                f"数据质量检查失败: 发现 {total_issues} 个问题，"
                f"详见上方报告。如需忽略，请设置 raise_on_error=False"
            )

    return issues
