"""Post-execution return attribution contracts."""

from __future__ import annotations

import pandas as pd
import pytest

from quant_engine.attribution import DailyReturnAttribution
from quant_engine.execution import ExecutionConfig
from quant_engine.research_pipeline import run_factor_backtest_research


def _zero_cost_config() -> ExecutionConfig:
    return ExecutionConfig(
        commission_bps=0,
        stamp_tax_bps=0,
        slippage_bps=0,
        min_trade_amount=0,
    )


def test_daily_attribution_closes_across_rebalance_and_holding_days() -> None:
    """开盘换仓时，隔夜和日内贡献必须来自实际换仓前后持仓。"""
    dates = pd.date_range("2026-01-05", periods=4, freq="B")
    scores = pd.DataFrame(
        {"A": [2.0, 0.0], "B": [1.0, 3.0]},
        index=dates[:2],
    )
    opens = pd.DataFrame(
        {"A": [10.0, 10.0, 15.0, 15.0], "B": [20.0, 20.0, 20.0, 21.0]},
        index=dates,
    )
    closes = pd.DataFrame(
        {"A": [10.0, 12.0, 15.0, 15.0], "B": [20.0, 20.0, 18.0, 21.0]},
        index=dates,
    )
    result = run_factor_backtest_research(
        scores,
        opens,
        closes,
        top_k=1,
        execution_price_field="open",
        valuation_price_field="close",
        initial_cash=1_000.0,
        config=_zero_cost_config(),
    )

    attribution = result.return_attribution()

    assert isinstance(attribution, DailyReturnAttribution)
    assert attribution.overnight.loc[dates[2], "A"] == pytest.approx(0.25)
    assert attribution.intraday.loc[dates[2], "B"] == pytest.approx(-0.125)
    assert attribution.asset_contributions.loc[dates[3], "B"] == pytest.approx(1 / 6)
    pd.testing.assert_series_equal(
        attribution.total_return,
        result.returns.rename("total_return"),
    )
    pd.testing.assert_series_equal(
        attribution.explained_return + attribution.residual,
        attribution.total_return,
        check_names=False,
    )
    assert attribution.residual.abs().max() < 1e-12


def test_daily_attribution_reports_execution_cost_separately() -> None:
    dates = pd.date_range("2026-01-05", periods=3, freq="B")
    scores = pd.DataFrame({"A": [1.0]}, index=dates[:1])
    prices = pd.DataFrame({"A": [10.0, 10.0, 10.0]}, index=dates)
    config = ExecutionConfig(
        commission_bps=10,
        stamp_tax_bps=0,
        slippage_bps=10,
        min_trade_amount=0,
    )
    result = run_factor_backtest_research(
        scores,
        prices,
        prices,
        top_k=1,
        gross_exposure=0.5,
        execution_price_field="open",
        valuation_price_field="close",
        initial_cash=1_000.0,
        config=config,
    )

    attribution = result.return_attribution()
    execution = result.execution.daily_executions[1].executions[0]

    assert attribution.asset_contributions.loc[dates[1], "A"] == 0.0
    assert attribution.transaction_cost.loc[dates[1]] == pytest.approx(
        -execution.total_cost / 1_000.0
    )
    assert attribution.total_return.loc[dates[1]] == pytest.approx(
        attribution.transaction_cost.loc[dates[1]]
    )
    assert attribution.residual.loc[dates[1]] == pytest.approx(0.0, abs=1e-12)


def test_return_attribution_is_empty_for_empty_research_result() -> None:
    dates = pd.date_range("2026-01-05", periods=3, freq="B")
    scores = pd.DataFrame(columns=["A"], index=pd.DatetimeIndex([]), dtype=float)
    prices = pd.DataFrame({"A": [10.0, 10.0, 10.0]}, index=dates)
    result = run_factor_backtest_research(
        scores,
        prices,
        prices,
        top_k=1,
        execution_price_field="open",
        valuation_price_field="close",
    )

    attribution = result.return_attribution()

    assert attribution.overnight.empty
    assert attribution.intraday.empty
    assert attribution.total_return.empty