ggx/src/hdiv/universe/daily.py

"""每日动态股票池筛选器（``backtest --mode daily`` 的选股环节）。

**它做什么**：从 ``start`` 起，对**每一个交易日**按当日可见数据重建股票池。
判定逻辑一行不改 —— 仍然调用 :class:`~hdiv.universe.selector.UniverseSelector`
与四个 ``Filter``；本模块只负责两件工程上的事：

1. **取数**：用 :class:`~hdiv.universe.pit.PitRepo` 按区块批量预载，
   逐日切片在内存完成（单次筛选从 10~18 秒降到秒级）；
2. **候选集预剪枝**：见下。

------------------------------------------------------------
候选集预剪枝：为什么是「精确」的，而不是「近似」
------------------------------------------------------------

市场滤网要逐行判断 5000 余只股票的交易所/板块/上市年限/市值/流动性，
这一段的 Python 开销与**候选数**成正比，是每日循环里最大的一项。

预剪枝只剔除「在整个回测区间内**不可能**通过市场滤网」的股票，
判据都是**上界**：

- **交易所 / 板块**：与日期无关，不在配置名单里的股票永远不可能通过；
- **上市年限**：当 ``list_date + min_listing_years`` 晚于区间**最后一天**时，
  该股在区间内任何一天都不满足 ``listed_years >= min_listing_years``；
- **市值**：当该股在区间内的 ``MAX(total_mv)``（换算为元）仍低于
  ``min_market_cap`` 时，任何一天都不满足市值下限。取不到市值（NULL）时
  **保留**，不剪。

被剪掉的股票在原流程里**必然**在第一个滤网（market）就被淘汰，因此：
最终入选集合逐只相同，`hd_daily_universe` 的内容也相同。
差别只在于「被剪掉的股票没有留下逐滤网的原因」—— 而每日选股模式
**不落库逐股淘汰原因**（只落库每日入选成员），所以这个差别不可观测。

正确性由 ``tests/test_daily.py::test_prune_does_not_change_selection`` 锁定：
同一批交易日，开/关剪枝必须选出**完全相同**的成员。

**流动性没有做预剪枝**：``stock_daily`` 的量价单位在 2015-2019 是「手/千元」、
2020 起是「股/元」，用 ``MAX(amount)`` 做上界会在早年低估 1000 倍，
误剪掉本该通过的股票。宁可少一项优化，也不接受一个会改变结果的上界。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import date, timedelta
from typing import Any

import numpy as np
import pandas as pd

from hdiv.core.config import UniverseConfig
from hdiv.core.errors import HdivError
from hdiv.data import db
from hdiv.universe.pit import PitRepo
from hdiv.universe.selector import UniverseSelector

__all__ = ["DailyUniverseScreener", "PruneReport", "ScreenDay"]

#: daily_basic 的市值列以**万元**存放（见 data/units.py）
_WAN = 1e4


@dataclass
class PruneReport:
    """预剪枝的规模，用于回答「为什么候选从 5000 变成 1000」。"""

    total: int = 0
    pruned_exchange: int = 0
    pruned_board: int = 0
    pruned_listing: int = 0
    pruned_market_cap: int = 0
    kept: int = 0

    def as_dict(self) -> dict[str, int]:
        return dict(self.__dict__)


@dataclass
class ScreenDay:
    """某一天的选股结果。"""

    trade_date: date
    candidate_count: int
    member_count: int
    symbols: list[str]
    members: pd.DataFrame
    stats: dict[str, int] = field(default_factory=dict)
    #: 当日**市场候选数**（未预剪枝）。与 ``candidate_count``（已预剪枝）分开记录 ——
    #: 预剪枝是纯性能开关，不该让页面上的「候选」含义随开关变化。
    listed_count: int = 0
    #: 入选股票在**决策日的因子取值**（来自 ``UniverseSelector.run`` 的 ``values``）。
    #: 必须带上它：``selected`` 里没有 ``dividend_yield`` 列（它在滤网的 values 里），
    #: 只从 ``selected`` 取列会让落库的股息率**整列为 NULL**。
    values: dict[str, dict[str, Any]] = field(default_factory=dict)


class DailyUniverseScreener:
    """逐日重建股票池（PIT），复用既有滤网，不改判定口径。"""

    def __init__(
        self,
        config: UniverseConfig,
        repo: PitRepo,
        *,
        verbose: bool = True,
    ) -> None:
        self.config = config
        self.repo = repo
        self.selector = UniverseSelector(config, repo=repo)
        self.verbose = verbose
        self.prune = PruneReport()
        self._allowed: set[str] | None = None

    # 刻意**没有** ``from_strategy(registry, strategy, ...)``：
    # ``registry.resolved_universe`` 只含策略 + universe.yml，拿不到
    # backtest.yml 的 ``universe_exclusions``（行业黑名单）。保留这样一个
    # 便捷构造器，等于给「逐日选股绕过行业排除」留了一条谁都看不出来的路。
    # 调用方应先用 ``BacktestConfig.resolved_universe`` 合并出**生效后**的
    # UniverseConfig（见 DailyRunner.__init__），再用普通构造函数传入。

    # ------------------------------------------------------------------
    # 预剪枝
    # ------------------------------------------------------------------

    def build_prune_set(
        self, window_start: date, window_end: date, *, use_market_cap: bool = True
    ) -> set[str]:
        """计算「在 ``[window_start, window_end]`` 内不可能通过市场滤网」的补集。

        返回**允许保留**的 symbol 集合。
        """
        cfg = self.config.market
        master = self.repo.stock_master().copy()
        rep = PruneReport(total=len(master))
        keep = pd.Series(True, index=master.index)

        def _drop(mask: pd.Series, counter: str) -> None:
            nonlocal keep
            hit = keep & mask
            setattr(rep, counter, getattr(rep, counter) + int(hit.sum()))
            keep = keep & ~mask

        # --- 交易所 / 板块（与日期无关）---
        if cfg.exchanges:
            allowed = {str(x) for x in cfg.exchanges}
            _drop(~master["exchange"].astype(str).isin(allowed), "pruned_exchange")
        if cfg.markets:
            allowed_m = {str(x) for x in cfg.markets}
            _drop(~master["market"].astype(str).isin(allowed_m), "pruned_board")

        # --- 上市年限：区间最后一天仍不足，则区间内永远不足 ---
        if cfg.min_listing_years and cfg.min_listing_years > 0:
            ld = pd.to_datetime(master["list_date"], errors="coerce")
            need_days = cfg.min_listing_years * 365.25
            can_pass = (window_end - ld.dt.date).apply(
                lambda x: x.days if pd.notna(x) else -1
            ) >= need_days
            _drop(~can_pass, "pruned_listing")

        # --- 市值：区间内 MAX(total_mv) 仍低于下限 ---
        if use_market_cap:
            cap = self._max_market_cap(window_start, window_end)
            if cap:
                for col, limit, counter in (
                    ("total_mv", cfg.min_market_cap, "pruned_market_cap"),
                    ("circ_mv", cfg.min_float_market_cap, "pruned_market_cap"),
                ):
                    if limit is None:
                        continue
                    mx = master["symbol"].map(cap.get(col, {}))
                    # 取不到市值时不剪（保守）：只有**确知**上限低于阈值才剔除
                    too_small = mx.notna() & (mx < float(limit))
                    _drop(too_small, counter)

        allowed = set(master.loc[keep, "symbol"].astype(str).tolist())
        rep.kept = len(allowed)
        self._allowed = allowed
        self.prune = rep
        self.repo.set_candidate_scope(allowed)
        return allowed

    def _max_market_cap(
        self, start: date, end: date
    ) -> dict[str, dict[str, float]]:
        """区间内逐股 ``MAX(total_mv)`` / ``MAX(circ_mv)``，单位为**元**。

        返回 ``{"total_mv": {symbol: 元}, "circ_mv": {symbol: 元}}``。

        单位：``daily_basic.total_mv`` / ``circ_mv`` 以**万元**存放（Tushare 口径），
        因此换算为元后再与配置里的元阈值比较。与 ``normalize_market_panel``
        的 ×1e4 是同一件事。
        """
        cfg = self.repo.cfg
        if not db.table_exists("daily_basic", cfg):
            return {}
        try:
            df = db.read_sql(
                "SELECT symbol, MAX(total_mv) AS mx_total_mv, "
                "       MAX(circ_mv) AS mx_circ_mv "
                "FROM daily_basic WHERE trade_date BETWEEN :s AND :e "
                "GROUP BY symbol",
                {"s": start, "e": end}, cfg=cfg,
            )
        except Exception:  # pragma: no cover - 取不到就退化为不剪枝
            return {}
        if df.empty:
            return {}
        out: dict[str, dict[str, float]] = {"total_mv": {}, "circ_mv": {}}
        syms = df["symbol"].astype(str).tolist()
        total = pd.to_numeric(df["mx_total_mv"], errors="coerce") * _WAN
        circ = pd.to_numeric(df["mx_circ_mv"], errors="coerce") * _WAN
        for sym, t, c in zip(syms, total, circ, strict=False):
            if pd.notna(t):
                out["total_mv"][sym] = float(t)
            if pd.notna(c):
                out["circ_mv"][sym] = float(c)
        return out

    # ------------------------------------------------------------------
    # 逐日筛选
    # ------------------------------------------------------------------

    def screen_day(self, day: date) -> ScreenDay:
        """筛选单个交易日。"""

        def _hook(stage: str, live: pd.DataFrame) -> None:
            # 把取数范围收窄到本阶段真正要评估的股票（纯性能开关）
            self.repo.restrict_to(live["symbol"].tolist())

        try:
            res = self.selector.run(
                asof=day, persist=False, verbose=False, on_stage=_hook
            )
        finally:
            self.repo.restrict_to(None)
        members = res["selected"]
        symbols = [str(s) for s in members["symbol"].tolist()]
        listed_total, screened = self.repo.listed_counts(res["asof_date"])
        vals = res.get("values") or {}
        return ScreenDay(
            trade_date=res["asof_date"],
            candidate_count=screened,
            member_count=int(res["member_count"]),
            symbols=symbols,
            members=members,
            stats=dict(res["stats"]),
            listed_count=listed_total,
            # 只为入选股票保留因子取值（全市场 4000 余只 × 1600 天会白占内存）
            values={s: dict(vals.get(s) or {}) for s in symbols},
        )

    def screen(self, days: list[date]) -> dict[date, set[str]]:
        """对 ``days`` 逐日筛选，返回 ``{交易日: 入选代码集合}``。"""
        if self._allowed is None:
            raise HdivError(
                "DailyUniverseScreener 必须先 build_prune_set(...) 再 screen(...)。\n"
                "  预剪枝是可选的性能优化；若不想剪枝，请显式传 use_market_cap=False\n"
                "  并把候选范围设为「全部上市股票」。"
            )
        out: dict[date, set[str]] = {}
        for day in days:
            sd = self.screen_day(day)
            out[sd.trade_date] = set(sd.symbols)
        return out

    # ------------------------------------------------------------------
    # 落库行
    # ------------------------------------------------------------------

    @staticmethod
    def member_rows(
        run_id: str, screens: list[ScreenDay], *, created_at: Any
    ) -> list[dict[str, Any]]:
        """把逐日选股结果摊平成 ``hd_daily_universe`` 的行。

        取值优先级：**滤网的 ``values`` → ``selected`` 的列**。
        股息率、支付率、FCF 覆盖等只在 ``values`` 里（它们是滤网算出来的），
        ``selected`` 只有行情/年报均值那几列；反过来 ``total_mv``/``roe_avg``
        只在列里。只取其中一边都会让某些列整列为 NULL。
        """
        rows: list[dict[str, Any]] = []
        for sd in screens:
            m = sd.members
            if m is None or m.empty:
                continue
            for rec in m.to_dict("records"):
                sym = str(rec.get("symbol"))
                vals = dict(sd.values.get(sym) or {})
                merged = {**{k: rec.get(k) for k in rec}, **vals}

                def pick(key: str) -> Any:
                    v = vals.get(key)
                    if v is None or (isinstance(v, float) and v != v):
                        v = rec.get(key)
                    return v

                rows.append({
                    "run_id": run_id,
                    "trade_date": sd.trade_date,
                    "symbol": sym,
                    "name": _s(rec.get("name")),
                    "industry": _s(rec.get("industry")),
                    # 股息率：筛选口径（自算优先）在 values 里；dv_ttm 在列里
                    "dividend_yield": _f(pick("dividend_yield")),
                    "total_mv": _f(pick("total_mv")),
                    "roe_avg": _f(pick("roe_avg")),
                    "listed_count": int(sd.listed_count) if sd.listed_count else None,
                    "candidate_count": int(sd.candidate_count),
                    "values_json": _values_json(merged),
                    "created_at": created_at,
                })
        return rows


def _s(v: Any) -> str | None:
    if v is None or (isinstance(v, float) and v != v):
        return None
    return str(v)[:64]


def _f(v: Any) -> float | None:
    if v is None:
        return None
    try:
        x = float(v)
    except (TypeError, ValueError):
        return None
    return None if (x != x or np.isinf(x)) else x


def _values_json(rec: dict[str, Any]) -> str:
    """入选时的关键因子快照（供「为什么是这只」复核）。"""
    import json

    keys = (
        "dividend_yield", "dividend_yield_computed", "dv_ttm", "ttm_dps",
        "pe_ttm", "pb", "ps_ttm", "total_mv", "circ_mv", "avg_amount_20d",
        "dividend_continuity_years", "dividend_years_in_window", "payout_ratio",
        "fcf_dividend_cover", "dps_cagr_5y", "roe", "roe_avg", "roic", "roic_avg",
        "debt_ratio", "ocf_to_netprofit", "ocf_to_profit_avg", "fin_years_count",
    )
    out: dict[str, Any] = {}
    for k in keys:
        if k not in rec:
            continue
        v = _f(rec.get(k))
        if v is not None:
            out[k] = v
    return json.dumps(out, ensure_ascii=False)