"""
问题2.3数值显著差异因子分析测试

目标：分析并修复以下因子在不同LOOKBACK_DAYS下的一致性问题
- GTJA_alpha016: cs_rank嵌套
- GTJA_alpha032: ts_sum嵌套cs_rank
- GTJA_alpha077: cs_rank+ts_decay_linear
- GTJA_alpha091: cs_rank嵌套max_
- GTJA_alpha121: ts_rank嵌套ts_corr
- GTJA_alpha130: ts_rank+ts_decay_linear
- GTJA_alpha141: cs_rank+ts_corr
"""

import json
from datetime import datetime, timedelta
from typing import Any, Dict, List

import numpy as np
import polars as pl
import pytest

from src.factors import FactorEngine
from src.factors.metadata import FactorManager


# =============================================================================
# 测试配置
# =============================================================================
PREDICT_START = "20250101"
PREDICT_END = "20250131"
LOOKBACK_2Y = 365 * 3  # 1095天
LOOKBACK_3Y = 365 * 4  # 1460天


# 问题2.3因子列表
ISSUE23_FACTORS = [
    "GTJA_alpha016",
    "GTJA_alpha032",
    "GTJA_alpha077",
    "GTJA_alpha091",
    "GTJA_alpha121",
    "GTJA_alpha130",
    "GTJA_alpha141",
]


def get_lookback_start_date(start_date: str, lookback_days: int) -> str:
    """计算考虑回看窗口后的实际开始日期。"""
    start_dt = datetime.strptime(start_date, "%Y%m%d")
    lookback_dt = start_dt - timedelta(days=lookback_days)
    return lookback_dt.strftime("%Y%m%d")


def compute_factors_with_lookback(
    lookback_days: int,
    factor_names: List[str],
) -> pl.DataFrame:
    """使用指定的回看窗口计算因子。"""
    actual_start = get_lookback_start_date(PREDICT_START, lookback_days)

    print(f"\n{'=' * 60}")
    print(f"LOOKBACK_DAYS = {lookback_days} ({lookback_days // 365}年)")
    print(f"实际加载数据范围: {actual_start} - {PREDICT_END}")

    # 创建 FactorEngine
    engine = FactorEngine()

    # 注册因子
    for name in factor_names:
        engine.add_factor(name)

    # 计算因子
    data = engine.compute(
        factor_names=factor_names,
        start_date=actual_start,
        end_date=PREDICT_END,
    )

    # 过滤到预测日期范围
    data = data.filter(data["trade_date"] >= PREDICT_START)

    print(f"计算完成: {data.shape}")
    return data


def compare_factor_values(
    data_2y: pl.DataFrame,
    data_3y: pl.DataFrame,
    factor_names: List[str],
) -> Dict[str, Any]:
    """比较两种回看窗口设置下的因子值。"""
    results = {
        "factors": {},
    }

    for factor_name in factor_names:
        if factor_name not in data_2y.columns or factor_name not in data_3y.columns:
            print(f"[跳过] {factor_name}: 因子不存在于两个数据集中")
            continue

        values_2y = data_2y[factor_name].to_numpy()
        values_3y = data_3y[factor_name].to_numpy()

        values_2y = np.asarray(values_2y, dtype=np.float64)
        values_3y = np.asarray(values_3y, dtype=np.float64)

        mask_2y = ~np.isnan(values_2y)
        mask_3y = ~np.isnan(values_3y)
        valid_mask = mask_2y & mask_3y

        if np.sum(valid_mask) == 0:
            continue

        valid_2y = values_2y[valid_mask]
        valid_3y = values_3y[valid_mask]

        # 计算差异
        diff = np.abs(valid_2y - valid_3y)
        max_diff = np.max(diff)
        mean_diff = np.mean(diff)
        count_diff = np.sum(diff > 1e-10)

        # 分类差异
        if max_diff >= 0.1 or np.isinf(max_diff):
            severity = "极高"
        elif max_diff >= 0.01:
            severity = "高"
        elif max_diff >= 1e-6:
            severity = "中"
        else:
            severity = "低"

        results["factors"][factor_name] = {
            "max_diff": max_diff,
            "mean_diff": mean_diff,
            "count_diff": count_diff,
            "severity": severity,
        }

        print(
            f"[{severity}] {factor_name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6e}, count={count_diff}"
        )

    return results


def test_issue23_factor_differences():
    """测试问题2.3因子的差异情况。"""
    print("\n" + "=" * 80)
    print("问题2.3数值显著差异因子分析")
    print("=" * 80)

    # 获取因子DSL定义
    manager = FactorManager()
    factor_defs = {}
    for name in ISSUE23_FACTORS:
        df = manager.get_factors_by_name(name)
        if df is None or len(df) == 0:
            print(f"[警告] 找不到因子 {name} 的定义")
            continue
        row = df.row(0)
        factor_defs[name] = {
            "dsl": df["dsl"][0],
            "factor_id": df["factor_id"][0],
        }
        print(f"\n{name}:")
        print(f"  DSL: {df['dsl'][0]}")

    # 计算两种回看期下的因子值
    data_2y = compute_factors_with_lookback(LOOKBACK_2Y, ISSUE23_FACTORS)
    data_3y = compute_factors_with_lookback(LOOKBACK_3Y, ISSUE23_FACTORS)

    # 排序以便比较
    data_2y = data_2y.sort(["trade_date", "ts_code"])
    data_3y = data_3y.sort(["trade_date", "ts_code"])

    # 比较因子值
    results = compare_factor_values(data_2y, data_3y, ISSUE23_FACTORS)

    # 汇总结果
    print(f"\n{'=' * 80}")
    print("问题2.3因子差异汇总")
    print("=" * 80)

    high_severity = []
    for name, stats in results["factors"].items():
        if stats["severity"] in ["高", "极高"]:
            high_severity.append(name)
        print(f"  {name}: {stats['severity']} (max_diff={stats['max_diff']:.6f})")

    print(f"\n高严重度因子数量: {len(high_severity)}")
    print(f"高严重度因子: {high_severity}")

    # 保存结果
    results_file = "tests/debug/fix_lookback_issue23/issue23_results.json"
    with open(results_file, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False, default=str)
    print(f"\n结果已保存到: {results_file}")

    return results


def test_cs_rank_behavior():
    """测试 cs_rank 在不同条件下的行为。"""
    print("\n" + "=" * 80)
    print("cs_rank 行为分析")
    print("=" * 80)

    # 创建测试数据：同一截面内有无NA的情况
    df_with_na = pl.DataFrame(
        {
            "trade_date": ["20250101"] * 5,
            "ts_code": ["001", "002", "003", "004", "005"],
            "value": [1.0, 2.0, 3.0, None, 5.0],  # 004有NA
        }
    )

    df_without_na = pl.DataFrame(
        {
            "trade_date": ["20250101"] * 5,
            "ts_code": ["001", "002", "003", "004", "005"],
            "value": [1.0, 2.0, 3.0, 4.0, 5.0],  # 无NA
        }
    )

    # 测试 rank() / count() 的结果
    for i, df in enumerate([df_with_na, df_without_na]):
        name = "有NA" if i == 0 else "无NA"
        print(f"\n{name}截面数据:")
        print(df["value"])

        result = (
            df.lazy()
            .with_columns(
                [
                    pl.col("value").rank().alias("rank"),
                    pl.col("value").count().over("trade_date").alias("count"),
                ]
            )
            .with_columns([(pl.col("rank") / pl.col("count")).alias("normalized_rank")])
            .collect()
        )
        print(f"排名结果:")
        print(result.select(["ts_code", "value", "rank", "count", "normalized_rank"]))


def test_analyze_factor_formula():
    """分析问题因子的公式结构。"""
    print("\n" + "=" * 80)
    print("问题2.3因子公式结构分析")
    print("=" * 80)

    manager = FactorManager()
    formulas = {}

    for name in ISSUE23_FACTORS:
        df = manager.get_factors_by_name(name)
        if df is not None and len(df) > 0:
            formulas[name] = df["dsl"][0]

    # 分析每个因子
    for name, dsl in formulas.items():
        print(f"\n{name}:")
        print(f"  DSL: {dsl}")

        # 分析涉及的函数
        functions = []
        if "cs_rank" in dsl:
            functions.append("cs_rank")
        if "ts_rank" in dsl:
            functions.append("ts_rank")
        if "ts_corr" in dsl:
            functions.append("ts_corr")
        if "ts_sum" in dsl:
            functions.append("ts_sum")
        if "ts_decay_linear" in dsl:
            functions.append("ts_decay_linear")
        if "ts_max" in dsl or "ts_min" in dsl:
            functions.append("ts_max/ts_min")
        if "max_" in dsl or "min_" in dsl:
            functions.append("max_/min_")

        print(f"  涉及函数: {functions}")


if __name__ == "__main__":
    # 运行测试
    test_issue23_factor_differences()
    test_cs_rank_behavior()
    test_analyze_factor_formula()
