market_sync/app/tasks/task_sector_features.py

"""同步任务：行业聚合特征（衍生计算）。

输入：kline_stock（日 K 线）+ industry（股票→行业映射）
输出：
  - sector_indices         ：行业日线（trade_date × sector_name × close × sector_amplitude）
  - sector_features_daily  ：行业特征（+ sector_ret + ema10/20/200 + score）

算法（参考 dashboard/api/services/sync/sector.py 重写，持久层改为 PG）：
  1) 每只股票日涨跌幅 pct_chg = (close/prev_close - 1) * 100
  2) 按行业聚合 sector_ret = 该行业所有股票当日 pct_chg 的均值
  3) 行业指数 close = 100 × ∏(1 + sector_ret/100) （基点 100，复合收益）
  4) 行业振幅 sector_amplitude = 该行业所有股票当日 (high-low)/close 的均值
  5) EMA10/20/200  = close 的指数移动平均（adjust=False）
  6) score: 0/1/2 — close>ema200 + 1，ema10>ema20 + 1

无外部 API 调用，纯本地计算；增量逻辑以"全量重算"实现（数据量小 ~75K 行）。
"""
from __future__ import annotations

import time
from datetime import datetime, timedelta
from typing import Any

import pandas as pd
from sqlalchemy import create_engine

from app.core.config import settings
from app.core.db import ops as db_ops
from app.core.sync.base import SyncTask
from app.core.utils.logging import get_logger

logger = get_logger("sync.sector_features")


# 起算点：DB 里 kline_stock 最早一天 往前 5 年，或固定一个默认起点
# （不用写死 2021-04-24，让数据自己决定；下面自动算）
DEFAULT_START_YEARS_BACK = 5

# 输出基点（行业指数 close 从多少开始）
INDEX_BASE = 100.0


class SyncSectorFeatures(SyncTask):
    dataset_id = "sector_features"

    def _run(
        self,
        *,
        trigger_source: str = "manual",
        codes: list[str] | None = None,
        max_workers: int = 1,  # 本任务纯本地计算，单线程足够
        **kwargs,
    ) -> dict[str, Any]:
        t0 = time.time()
        logger.info("[sector] 启动行业聚合特征计算")

        # ── 1) 拉 kline_stock（日线）只取需要的列 ──
        # 走 SQLAlchemy engine + PG URL（pd.read_sql 一次性读 1.1kw 行 → pandas DataFrame）
        engine = create_engine(settings.pg_sqlalchemy_url())
        df_kline = pd.read_sql(
            'SELECT stock_code, trade_date, open, high, low, "close", volume '
            'FROM market_data.kline_stock ORDER BY stock_code, trade_date',
            engine,
        )
        if df_kline.empty:
            return {"status": "error", "message": "kline_stock 为空"}
        df_kline["stock_code"] = df_kline["stock_code"].astype(str).str.zfill(6)
        df_kline["trade_date"] = pd.to_datetime(df_kline["trade_date"], errors="coerce")
        logger.info(f"[sector] kline_stock: {len(df_kline):,} 行, {df_kline['stock_code'].nunique()} 只, "
                    f"{df_kline['trade_date'].min().date()} ~ {df_kline['trade_date'].max().date()}")

        # ── 2) 拉 industry（股票→行业映射）──
        df_ind = pd.read_sql(
            "SELECT code, industry_name FROM market_data.industry WHERE industry_name IS NOT NULL",
            engine,
        )
        if df_ind.empty:
            return {"status": "error", "message": "industry 表为空"}
        df_ind["code"] = df_ind["code"].astype(str).str.zfill(6)
        df_ind = df_ind.drop_duplicates(subset=["code"], keep="first")
        logger.info(f"[sector] industry 映射: {len(df_ind):,} 行, {df_ind['industry_name'].nunique()} 个行业")

        # ── 3) 算每只股票日涨跌幅 + 振幅 ──
        df = df_kline.merge(df_ind, left_on="stock_code", right_on="code", how="inner")
        df = df.sort_values(["stock_code", "trade_date"])
        df["prev_close"] = df.groupby("stock_code")["close"].shift(1)
        df["stock_amplitude"] = (df["high"] - df["low"]) / df["close"]
        df = df.dropna(subset=["prev_close", "industry_name"])
        df["pct_chg"] = (df["close"] / df["prev_close"] - 1.0) * 100.0
        logger.info(f"[sector] 合并后: {len(df):,} 行, 行业 {df['industry_name'].nunique()} 个")

        # ── 4) 按行业 + 日期聚合 ──
        sector_daily = (
            df.groupby(["industry_name", "trade_date"], as_index=False).agg(
                sector_ret=("pct_chg", "mean"),
                sector_amplitude=("stock_amplitude", "mean"),
            )
            .rename(columns={"industry_name": "sector_name"})
            .sort_values(["sector_name", "trade_date"])
            .reset_index(drop=True)
        )
        logger.info(f"[sector] 行业日聚合: {len(sector_daily):,} 行")

        # ── 5) 算行业指数 close（基点 100，复合）──
        sector_index = self._build_index(sector_daily)

        # ── 6) EMA + score ──
        sector_index = self._calc_ema(sector_index)
        sector_index["score"] = sector_index.apply(
            lambda r: self._calc_score(r["close"], r["ema10"], r["ema20"], r["ema200"]),
            axis=1,
        )

        # 整理列名
        out_cols = ["trade_date", "sector_name", "sector_ret", "sector_amplitude",
                    "close", "ema10", "ema20", "ema200", "score"]
        sector_index = sector_index[out_cols]
        sector_index["trade_date"] = sector_index["trade_date"].dt.strftime("%Y-%m-%d")
        # NaN → None（让 MySQL 接受 NULL）
        sector_index = sector_index.where(pd.notnull(sector_index), None)

        # ── 7) 写入两张表 ──
        # 7a. sector_indices (4 列)
        si_rows = sector_index[["trade_date", "sector_name", "close", "sector_amplitude"]].to_dict("records")
        db_ops.replace_all_sector_indices(si_rows)
        # 7b. sector_features_daily (9 列)
        sf_rows = sector_index.to_dict("records")
        db_ops.replace_all_sector_features(sf_rows)

        elapsed = round(time.time() - t0, 1)
        msg = (
            f"行业聚合 {sector_index['sector_name'].nunique()} 个行业 "
            f"× {sector_index['trade_date'].nunique()} 天, "
            f"共 {len(sector_index):,} 行, {elapsed}s"
        )
        logger.info(f"[sector] {msg}")
        return {
            "status": "ok",
            "message": msg,
            "sectors": int(sector_index["sector_name"].nunique()),
            "days": int(sector_index["trade_date"].nunique()),
            "rows": int(len(sector_index)),
            "elapsed_sec": elapsed,
        }

    @staticmethod
    def _build_index(sector_daily: pd.DataFrame) -> pd.DataFrame:
        """行业指数 close = INDEX_BASE × ∏(1 + sector_ret/100)"""
        out = []
        for sector_name, group in sector_daily.groupby("sector_name", sort=False):
            g = group.sort_values("trade_date").copy()
            close_vals = []
            cur_base = INDEX_BASE
            for ret in g["sector_ret"].to_numpy(dtype=float):
                if pd.isna(ret):
                    close_vals.append(cur_base)
                else:
                    cur_base = cur_base * (1 + float(ret) / 100.0)
                    close_vals.append(cur_base)
            g["close"] = close_vals
            out.append(g)
        return pd.concat(out, ignore_index=True)

    @staticmethod
    def _calc_ema(sector_index: pd.DataFrame) -> pd.DataFrame:
        """每个行业分别算 EMA10/20/200"""
        out = []
        for sector_name, group in sector_index.groupby("sector_name", sort=False):
            g = group.sort_values("trade_date").copy()
            g["ema10"] = g["close"].ewm(span=10, adjust=False).mean()
            g["ema20"] = g["close"].ewm(span=20, adjust=False).mean()
            g["ema200"] = g["close"].ewm(span=200, adjust=False).mean()
            out.append(g)
        return pd.concat(out, ignore_index=True)

    @staticmethod
    def _calc_score(close, ema10, ema20, ema200) -> int:
        score = 0
        if pd.notna(close) and pd.notna(ema200) and close > ema200:
            score += 1
        if pd.notna(ema10) and pd.notna(ema20) and ema10 > ema20:
            score += 1
        return int(score)