# tests/factor/test_fundamental_neutralize.py
"""合成层 v2b: cs_neutralize 截面中性化算子数学等值/边界 + fund_*_neu 6 因子注册 + weighting 三源表."""
import datetime as dt
import json
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "vnpy_v4.4.0")))

import polars as pl
import pytest

from vnpy.alpha.dataset.utility import DataProxy, EXPRESSION_FUNCTIONS, calculate_by_expression

from sanguo_factor import composite_library, composite_weighting, fundamental_neutralize
from sanguo_factor.fundamental_neutralize import (
    FUND_NEU_SOURCES, SIZE_EXPRESSION, cs_neutralize, build_neutralized_expression,
)
from sanguo_factor.composite_weighting import WeightProfile, register_composite_from_profile
from sanguo_factor.fundamental_library import (
    FUNDAMENTAL_FACTORS, FUNDAMENTAL_P1_FACTORS, FUNDAMENTAL_P1B_FACTORS,
)
from sanguo_factor.registry import get_factor, _REGISTRY


@pytest.fixture(autouse=True)
def _ensure_v2b_registered():
    """其它测试模块清空 _REGISTRY 后幂等重挂(与 test_composite_weighting 同模式)."""
    fundamental_neutralize._register_all()
    composite_weighting._register_all()


def _proxy(rows: list[tuple[str, str, float | None]]) -> DataProxy:
    """(date, symbol, value) 行列表 → DataProxy(datetime/vt_symbol/data)."""
    df = pl.DataFrame({
        "datetime": [dt.date.fromisoformat(d) for d, _s, _v in rows],
        "vt_symbol": [s for _d, s, _v in rows],
        "v": [v for _d, _s, v in rows],
    })
    return DataProxy(df)


# ==================== 算子数学等值(手算 OLS) ====================

def test_cs_neutralize_matches_hand_computed_ols_positive_slope():
    """正斜率日手算: x=[1,2,3,4], y=[1,3,2,5] → b=1.1, a=0, resid=[-0.1,0.8,-1.3,0.6]."""
    x = _proxy([("2024-01-02", f"s{i}", v) for i, v in enumerate([1.0, 2.0, 3.0, 4.0])])
    y = _proxy([("2024-01-02", f"s{i}", v) for i, v in enumerate([1.0, 3.0, 2.0, 5.0])])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got == pytest.approx([-0.1, 0.8, -1.3, 0.6], abs=1e-12)


def test_cs_neutralize_matches_hand_computed_ols_negative_slope():
    """负斜率日手算: x=[1,2,3,4], y=[5,4,2,1] → b=-1.4, a=6.5, resid=[-0.1,0.3,-0.3,0.1]."""
    x = _proxy([("2024-01-03", f"s{i}", v) for i, v in enumerate([1.0, 2.0, 3.0, 4.0])])
    y = _proxy([("2024-01-03", f"s{i}", v) for i, v in enumerate([5.0, 4.0, 2.0, 1.0])])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got == pytest.approx([-0.1, 0.3, -0.3, 0.1], abs=1e-12)


def test_cs_neutralize_days_are_independent():
    """逐日独立回归: 同一调用内两日用各自斜率(拼 day1+day2,两组手算值各归各)."""
    x = _proxy([(d, f"s{i}", v) for d, vals in (("2024-01-02", [1.0, 2.0, 3.0, 4.0]),
                                                ("2024-01-03", [1.0, 2.0, 3.0, 4.0]))
                for i, v in enumerate(vals)])
    y = _proxy([(d, f"s{i}", v) for d, vals in (("2024-01-02", [1.0, 3.0, 2.0, 5.0]),
                                                ("2024-01-03", [5.0, 4.0, 2.0, 1.0]))
                for i, v in enumerate(vals)])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got == pytest.approx([-0.1, 0.8, -1.3, 0.6, -0.1, 0.3, -0.3, 0.1], abs=1e-12)


# ==================== 算子边界 ====================

def test_cs_neutralize_null_rows_propagate():
    """null 传播: x/y 任一 null 的行残差 null;其余行按有效配对样本回归."""
    # 5 股: 3 个有效配对(n=3 达阈值) + 1 个 y-null + 1 个 x-null
    x = _proxy([("2024-01-02", "s0", 1.0), ("2024-01-02", "s1", 2.0), ("2024-01-02", "s2", 3.0),
                ("2024-01-02", "s3", 4.0), ("2024-01-02", "s4", None)])
    y = _proxy([("2024-01-02", "s0", 1.0), ("2024-01-02", "s1", 3.0), ("2024-01-02", "s2", 2.0),
                ("2024-01-02", "s3", None), ("2024-01-02", "s4", 9.0)])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got[3] is None and got[4] is None        # null 行残差 null
    # 前 3 行 = x=[1,2,3], y=[1,3,2] 的手算 OLS: b=0.5, a=1.0 → resid=[-0.5,1.0,-0.5]
    assert got[:3] == pytest.approx([-0.5, 1.0, -0.5], abs=1e-12)


def test_cs_neutralize_small_group_returns_null():
    """当日有效样本 < 3 → 该日全部残差 null(含有效行)."""
    x = _proxy([("2024-01-02", "s0", 1.0), ("2024-01-02", "s1", 2.0),
                ("2024-01-02", "s2", None), ("2024-01-02", "s3", None)])
    y = _proxy([("2024-01-02", "s0", 1.0), ("2024-01-02", "s1", 3.0),
                ("2024-01-02", "s2", 7.0), ("2024-01-02", "s3", 9.0)])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got == [None, None, None, None]


def test_cs_neutralize_zero_variance_returns_null():
    """当日 var(x)=0(x 全同) → 该日全部残差 null."""
    x = _proxy([("2024-01-02", f"s{i}", 5.0) for i in range(4)])
    y = _proxy([("2024-01-02", f"s{i}", v) for i, v in enumerate([1.0, 2.0, 3.0, 4.0])])
    got = cs_neutralize(y, x).df["data"].to_list()
    assert got == [None, None, None, None]


def test_cs_neutralize_residual_mean_zero_per_day():
    """截面性质: 残差逐日均值 ≈ 0(OLS 回归性质)."""
    rows_x, rows_y = [], []
    for day, xvals, yvals in (
        ("2024-01-02", [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2.0, 1.5, 4.0, 3.5, 6.5, 5.0]),
        ("2024-01-03", [10.0, 20.0, 30.0, 40.0, 50.0, 60.0], [8.0, 9.0, 2.0, 3.0, 1.0, 7.0]),
        ("2024-01-04", [-3.0, -1.0, 0.0, 2.0, 4.0, 9.0], [5.0, -2.0, 3.0, -4.0, 1.0, 0.5]),
    ):
        for i, (xv, yv) in enumerate(zip(xvals, yvals)):
            rows_x.append((day, f"s{i}", xv))
            rows_y.append((day, f"s{i}", yv))
    out = cs_neutralize(_proxy(rows_y), _proxy(rows_x)).df
    per_day = out.group_by("datetime").agg(pl.col("data").mean().alias("mean_resid"))
    for row in per_day.iter_rows(named=True):
        assert row["mean_resid"] == pytest.approx(0.0, abs=1e-10)


# ==================== 引擎注册与表达式求值 ====================

def test_cs_neutralize_registered_in_expression_functions():
    """沿 fast_ops 官方扩展点: cs_neutralize 已注册进 vnpy EXPRESSION_FUNCTIONS."""
    assert EXPRESSION_FUNCTIONS.get("cs_neutralize") is cs_neutralize


def test_cs_neutralize_evaluates_via_engine_expression():
    """引擎端到端: calculate_by_expression('cs_neutralize(y_col, x_col)') 多参数映射列,
    输出与直接函数调用逐点一致."""
    df = pl.DataFrame({
        "datetime": [dt.date(2024, 1, 2)] * 4 + [dt.date(2024, 1, 3)] * 4,
        "vt_symbol": [f"s{i}" for i in range(4)] * 2,
        "y_col": [1.0, 3.0, 2.0, 5.0, 5.0, 4.0, 2.0, 1.0],
        "x_col": [1.0, 2.0, 3.0, 4.0] * 2,
    })
    via_engine = calculate_by_expression(df, "cs_neutralize(y_col, x_col)")
    direct = cs_neutralize(DataProxy(df[["datetime", "vt_symbol", "y_col"]]),
                           DataProxy(df[["datetime", "vt_symbol", "x_col"]]))
    assert via_engine["data"].to_list() == pytest.approx(direct.df["data"].to_list(), abs=1e-12)


# ==================== FUND_NEU_SOURCES 表与 6 因子注册 ====================

def test_fund_neu_sources_table_shape():
    """表形态: 6 项 = FUND_SOURCES 同名源加 _neu 后缀,方向全 '+'(高=好语义)."""
    assert FUND_NEU_SOURCES == [(f"{name}_neu", "+") for name, _d in composite_library.FUND_SOURCES]
    assert len(FUND_NEU_SOURCES) == 6


def test_six_neu_factors_registered_with_cs_neutralize_and_size():
    """6 个 fund_*_neu 已注册(category=fundamental),表达式含 cs_neutralize 与总市值."""
    for neu_name, _direction in FUND_NEU_SOURCES:
        fac = get_factor(neu_name)
        assert fac is not None, f"{neu_name} 未注册"
        assert fac["category"] == "fundamental"
        assert fac["expression"].startswith("cs_rank(")
        assert "cs_neutralize(" in fac["expression"]
        assert SIZE_EXPRESSION in fac["expression"]


def test_neu_sign_direction_matches_source():
    """负号方向与原源一致: 原内层带负号(nsi/gdhs_chg)→ '(-1) * cs_neutralize(正指标)';
    原内层无负号(bp)→ 直接 cs_neutralize.不手抄指标式,符号从注册表表达式派生."""
    nsi = get_factor("fund_nsi_neu")["expression"]
    assert nsi == "cs_rank((-1) * cs_neutralize(nsi, close * share_capital))"
    gdhs = get_factor("fund_gdhs_chg_neu")["expression"]
    assert gdhs == "cs_rank((-1) * cs_neutralize(gdhs_chg, close * share_capital))"
    bp = get_factor("fund_bp_neu")["expression"]
    assert bp == "cs_rank(cs_neutralize(equity / (close * share_capital), close * share_capital))"
    # 全 6 源一致性: 符号形态从原源表达式剥壳推导,逐源核对
    for src_name, _d in composite_library.FUND_SOURCES:
        assert get_factor(f"{src_name}_neu")["expression"] == \
            build_neutralized_expression(get_factor(src_name)["expression"])


def test_original_fund_sources_unchanged():
    """原 6 源零回归: fund_bp 等表达式与 fundamental_library 注册表逐字一致."""
    registered = {name: expr for name, expr, _doc, _ic
                  in [*FUNDAMENTAL_FACTORS, *FUNDAMENTAL_P1_FACTORS, *FUNDAMENTAL_P1B_FACTORS]}
    for src_name, _d in composite_library.FUND_SOURCES:
        assert get_factor(src_name)["expression"] == registered[src_name]


# ==================== weighting 三源表匹配 + v2b graceful ====================

def test_weighting_accepts_fund_neu_profile():
    """纯 FUND_NEU_SOURCES 档案: 方向表匹配通过并注册,项=内嵌(已定向,不再包 cs_rank)."""
    name = "composite_fund6_neu_test"
    _REGISTRY.pop(name, None)
    sources = {n: (d, 1.0) for n, d in FUND_NEU_SOURCES}
    register_composite_from_profile(WeightProfile("t-neu", "t", "manual", sources, (), {}), name)
    fac = get_factor(name)
    assert fac is not None and fac["category"] == "composite"
    assert get_factor("fund_bp_neu")["expression"] in fac["expression"]
    assert "vma_60" not in fac["expression"]
    _REGISTRY.pop(name, None)


def test_weighting_accepts_combined_18_profile():
    """QUANT+FUND_NEU 合并 18 源档案: 匹配通过,量价=rank 项+中性化=内嵌项并存."""
    name = "composite_all18_v2b_test"
    _REGISTRY.pop(name, None)
    sources = {n: (d, 1.0) for n, d in composite_library.QUANT_SOURCES}
    sources.update({n: (d, 1.0) for n, d in FUND_NEU_SOURCES})
    register_composite_from_profile(WeightProfile("t-18", "t", "manual", sources, (), {}), name)
    expr = get_factor(name)["expression"]
    assert get_factor("vma_60")["expression"] in expr      # 量价源进合成
    assert get_factor("fund_bp_neu")["expression"] in expr  # 中性化源进合成
    _REGISTRY.pop(name, None)


def test_weighting_rejects_mixed_direction_drift():
    """方向漂移仍拒: 18 源档案翻转一源方向 → ValueError 且不注册(不被三选一吞掉)."""
    name = "composite_all18_v2b_bad"
    sources = {n: (d, 1.0) for n, d in composite_library.QUANT_SOURCES}
    sources.update({n: (d, 1.0) for n, d in FUND_NEU_SOURCES})
    first = next(iter(sources))
    d, w = sources[first]
    sources[first] = ("-" if d == "+" else "+", w)
    with pytest.raises(ValueError, match="QUANT_SOURCES"):
        register_composite_from_profile(
            WeightProfile("t-bad", "t", "manual", sources, (), {}), name)
    assert get_factor(name) is None


def test_v2b_profile_absent_skips_registration(tmp_path, monkeypatch):
    """v2b 档案不存在(真档案已上线,monkeypatch 指向不存在路径模拟缺席):
    _register_all 不炸、不注册 composite_all18_v2b."""
    _REGISTRY.pop(composite_weighting._V2B_COMPOSITE_NAME, None)
    monkeypatch.setattr(composite_weighting, "_V2B_PROFILE_PATH",
                        tmp_path / "absent_v2b.json")
    composite_weighting._register_all()  # 不应 raise
    assert get_factor(composite_weighting._V2B_COMPOSITE_NAME) is None


def test_v2b_profile_present_registers(tmp_path, monkeypatch):
    """v2b 档案存在(tmp 档案+monkeypatch 路径): _register_all 自动注册 composite_all18_v2b."""
    _REGISTRY.pop(composite_weighting._V2B_COMPOSITE_NAME, None)
    sources = {n: (d, 1.0) for n, d in composite_library.QUANT_SOURCES}
    sources.update({n: (d, 1.0) for n, d in FUND_NEU_SOURCES})
    profile = tmp_path / "all18_v2b_v1.json"
    profile.write_text(json.dumps({
        "profile_id": "all18-v2b-v1", "name": "t", "method": "manual",
        "sources": {n: {"direction": d, "weight": w} for n, (d, w) in sources.items()},
    }, ensure_ascii=False), encoding="utf-8")
    monkeypatch.setattr(composite_weighting, "_V2B_PROFILE_PATH", profile)
    composite_weighting._register_all()
    fac = get_factor(composite_weighting._V2B_COMPOSITE_NAME)
    assert fac is not None and fac["category"] == "composite"
    assert get_factor("fund_bp_neu")["expression"] in fac["expression"]
    _REGISTRY.pop(composite_weighting._V2B_COMPOSITE_NAME, None)  # 还原全局注册表