market_sync/app/tasks/task_market_regime.py

"""同步任务：市场情绪（market_regime_daily，衍生源）。

基于本地 kline_stock 数据聚合：advancers / decliners / turnover / advance_ratio。"""
from __future__ import annotations

from datetime import datetime
from typing import Any

import numpy as np
import pandas as pd
from sqlalchemy import text

from app.core.db import ops as db_ops
from app.core.db.orm import engine as pg_engine
from app.core.sync.base import SyncTask
from app.core.utils.logging import get_logger

logger = get_logger("sync.market_regime")

DEFAULT_START = "2018-01-01"
EXTREME_PANIC_ADV = 0.2
EXTREME_PANIC_TURN = 1.5


def _is_extreme_panic(row: pd.Series) -> int:
    ar = pd.to_numeric(row.get("advance_ratio"), errors="coerce")
    tr = pd.to_numeric(row.get("turnover_ratio_5d"), errors="coerce")
    if pd.isna(ar) or pd.isna(tr):
        return 0
    return int(ar < EXTREME_PANIC_ADV and tr > EXTREME_PANIC_TURN)


class SyncMarketRegime(SyncTask):
    dataset_id = "market_regime"

    def _run(self, *, trigger_source: str = "manual", start: str = DEFAULT_START, **kwargs) -> dict[str, Any]:
        end = datetime.now().strftime("%Y-%m-%d")
        self._progress(message=f"构建 market_regime {start} ~ {end}...")

        # 1. 取所有非退市股票代码
        codes = list(db_ops.iter_stock_codes(active_only=True))
        if not codes:
            return {"status": "error", "message": "stocks 表为空"}

        # 2. 拉所有 K 线（按股票）— 走 SA text() 走 PG
        sql = text("""
            SELECT
                stock_code,
                trade_date,
                "close",
                volume
            FROM market_data.kline_stock
            WHERE trade_date BETWEEN :start AND :end
            ORDER BY stock_code, trade_date
        """)
        with pg_engine.connect() as conn:
            rows = conn.execute(sql, {"start": start, "end": end}).mappings().all()
        if not rows:
            return {"status": "warning", "message": f"kline_stock 在 {start} ~ {end} 区间无数据"}

        df = pd.DataFrame(rows)
        df["trade_date"] = pd.to_datetime(df["trade_date"], errors="coerce")
        df["close"] = pd.to_numeric(df["close"], errors="coerce")
        df["volume"] = pd.to_numeric(df["volume"], errors="coerce").fillna(0.0)
        df = df.dropna(subset=["trade_date", "close"]).sort_values(["stock_code", "trade_date"])
        df["prev_close"] = df.groupby("stock_code")["close"].shift(1)
        df = df.dropna(subset=["prev_close"])
        df["advancers"] = (df["close"] > df["prev_close"]).astype(int)
        df["decliners"] = (df["close"] < df["prev_close"]).astype(int)
        df["turnover"] = df["close"] * df["volume"]

        daily = df.groupby("trade_date", as_index=False)[["advancers", "decliners", "turnover"]].sum()
        total = daily["advancers"] + daily["decliners"]
        daily["advance_ratio"] = np.where(total > 0, daily["advancers"] / total, np.nan)
        daily["turnover_avg_5d"] = daily["turnover"].rolling(5, min_periods=5).mean()
        daily["turnover_ratio_5d"] = np.where(
            daily["turnover_avg_5d"] > 0,
            daily["turnover"] / daily["turnover_avg_5d"],
            np.nan,
        )
        daily["source"] = "local_kline_proxy"
        daily["is_extreme_panic"] = daily.apply(_is_extreme_panic, axis=1)
        daily["trade_date"] = daily["trade_date"].dt.strftime("%Y-%m-%d")

        rows_out = [
            {
                "trade_date": r["trade_date"],
                "advancers": float(r["advancers"]),
                "decliners": float(r["decliners"]),
                "advance_ratio": float(r["advance_ratio"]) if pd.notna(r["advance_ratio"]) else 0,
                "turnover": float(r["turnover"]),
                "turnover_avg_5d": float(r["turnover_avg_5d"]) if pd.notna(r["turnover_avg_5d"]) else 0,
                "turnover_ratio_5d": float(r["turnover_ratio_5d"]) if pd.notna(r["turnover_ratio_5d"]) else 0,
                "source": r["source"],
                "is_extreme_panic": int(r["is_extreme_panic"]),
            }
            for _, r in daily.iterrows()
        ]

        try:
            db_ops.upsert_market_regime_rows(rows_out)
        except Exception as e:
            return {"status": "error", "message": f"写库失败: {e}"}

        msg = f"market_regime {len(rows_out)} 天 ({rows_out[0]['trade_date']} ~ {rows_out[-1]['trade_date']})"
        return {
            "status": "ok",
            "message": msg,
            "rows": len(rows_out),
            "date_min": rows_out[0]["trade_date"],
            "date_max": rows_out[-1]["trade_date"],
        }