akg-factor-bridge/common.py

"""共用：股票码规范化、覆盖池、交易日历、幂等写因子表、注册 factor_metadata（自适应列）。"""
import json

import pandas as pd

import config
import db


def to_prefix(ts_code: str) -> str:
    """600000.SH -> SH600000；已是前缀式(SZ002625)或纯代码则原样返回。"""
    s = str(ts_code).strip()
    if "." in s:
        num, exch = s.split(".", 1)
        return f"{exch.upper()}{num}"
    return s


def load_universe() -> set[str]:
    """KG 覆盖池 ts_code 集合（600000.SH 形态），来自基座视图 v_factor_universe。

    ⚠️ 这是**当前态**：industry_pools 只存最新快照、每周一 refresh_pools 自动长大。
    用它过滤历史数据会引入成员性前视——见 universe_filter 与 config.SUBFACTOR_UNIVERSE。
    """
    df = db.read_pg("SELECT ts_code FROM v_factor_universe")
    return set(df["ts_code"].astype(str).str.strip())


def universe_filter(df: pd.DataFrame, col: str = "stock_code",
                    as_prefix: bool = True) -> pd.DataFrame:
    """按 config.SUBFACTOR_UNIVERSE 决定子因子是否受覆盖池限制。

    'pool'（默认，保持原行为）：只留池内成员。
    'market'（评审 §4 建议）：全市场出行，池只在 akg_score 侧当过滤器。

    为什么建议 market：industry_pools 只存最新态且每周一长大，用它过滤历史子因子，
    值是点时正确的、**成员性却是前视的**（用今天的池去筛历史，而今天的池里有大量
    主题是后来的研报才长出来的——正是 §2.2 批评传导时用的同一条论证）；
    同时截面统计量会在每周一发生结构性跳变，IC 序列带周期性伪影。
    子因子表是「可独立观察的仪表」，过滤该发生在消费端。
    """
    if config.SUBFACTOR_UNIVERSE == "pool":
        uni = load_universe()
        if as_prefix:
            uni = {to_prefix(x) for x in uni}
        return df[df[col].astype(str).str.strip().isin(uni)]
    return df


def trading_days(start: str, end: str) -> list:
    """目标区间交易日历。

    ⚠️ 曾用热度表 stock_fund_heat_scores —— 但它最早只有 2025-03-26，
    导致 `build akg_event --mode history --start 2024-01-01`（README 里的示例命令）
    拿到空日历 → 静默返回空 → 只打印「无数据（跳过）」，不报错。
    改用行情表 gp_day_data（实测 5584 天，覆盖全历史）。
    """
    df = db.read_mysql(
        "price",
        "SELECT DISTINCT `timestamp` AS trade_date FROM gp_day_data "
        "WHERE `timestamp` BETWEEN %s AND %s ORDER BY `timestamp`",
        (start, end))
    if df.empty:
        print(f"  ⚠️ 交易日历为空（gp_day_data 在 {start}~{end} 无数据）"
              f"——上层会跳过该区间，请核对区间与行情表覆盖")
    return list(pd.to_datetime(df["trade_date"]))


_CREATE_FACTOR_TABLE = """
CREATE TABLE IF NOT EXISTS {t} (
  trade_date DATE NOT NULL,
  stock_code VARCHAR(15) NOT NULL,
  factor_value DOUBLE,
  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
  PRIMARY KEY (trade_date, stock_code),
  KEY idx_stock (stock_code)
) ENGINE=InnoDB
"""


def ensure_table(table: str) -> None:
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(_CREATE_FACTOR_TABLE.format(t=table))
        conn.commit()


def write_factor(table: str, df: pd.DataFrame, mode: str = "daily") -> None:
    """df 列 [trade_date, stock_code, factor_value] → 幂等写因子表。
    daily/history 都是「删涉及日期区间 → 批插」。stock_code 统一转前缀式。

    分块提交（评审）：热度全史约 54 万行，原来单事务 executemany 走 ShardingSphere
    代理有风险。改成 DELETE 一个事务 + INSERT 按 config.WRITE_CHUNK_ROWS 分块提交。
    代价是中途失败会留下部分区间——但整段写入按日期区间幂等，重跑即修复。
    """
    if df is None or df.empty:
        print(f"  {table}: 无数据（跳过）")
        return
    df = df.dropna(subset=["trade_date", "stock_code", "factor_value"]).copy()
    df["stock_code"] = df["stock_code"].map(to_prefix)
    df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
    df = df.drop_duplicates(["trade_date", "stock_code"], keep="last")
    if df.empty:
        print(f"  {table}: 清洗后无数据")
        return
    dmin, dmax = df["trade_date"].min(), df["trade_date"].max()
    ensure_table(table)
    rows = list(df[["trade_date", "stock_code", "factor_value"]]
                .itertuples(index=False, name=None))
    chunk = max(1, config.WRITE_CHUNK_ROWS)
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(f"DELETE FROM {table} WHERE trade_date BETWEEN %s AND %s",
                        (dmin, dmax))
        conn.commit()
        for i in range(0, len(rows), chunk):
            with conn.cursor() as cur:
                cur.executemany(
                    f"INSERT INTO {table} (trade_date, stock_code, factor_value) "
                    f"VALUES (%s,%s,%s)", rows[i:i + chunk])
            conn.commit()
            if len(rows) > chunk:
                print(f"    ...{min(i + chunk, len(rows))}/{len(rows)}")
    print(f"  {table}: 写入 {len(rows)} 行, 日期 {dmin}~{dmax}")


def register(factor_code, display_name, table, category, desc,
             author="akg-factor-bridge") -> None:
    """注册 factor_metadata。自适应实际存在的列（避免猜死 schema）——
    列名以线上 factor_metadata 为准，本函数只写它有的列。"""
    values = {
        "factor_code": factor_code, "display_name": display_name,
        "target_ds_name": "ds_a", "target_table_name": table,
        "author": author, "description": desc,
        "category": json.dumps(category, ensure_ascii=False),
        "frequency": "daily", "status": "active", "factor_type": "single",
    }
    with db.factor_conn() as conn:
        with conn.cursor() as cur:
            # ShardingSphere 代理不支持 information_schema 内省 → 用真实查询取列名
            cur.execute("SELECT * FROM factor_metadata LIMIT 0")
            cols = [d[0] for d in (cur.description or [])]
            if not cols:   # 极端兜底：代理连列元数据都不给时，用平台已知 schema
                cols = ["factor_code", "display_name", "target_ds_name", "target_table_name",
                        "start_date", "end_date", "author", "description", "category",
                        "frequency", "status", "factor_type"]
            use = [k for k in values if k in cols]
            if "factor_code" not in use:
                raise RuntimeError(f"factor_metadata 取不到 factor_code 列；实际列: {cols}")
            ph = ",".join(["%s"] * len(use))
            upd = ",".join(f"{k}=VALUES({k})" for k in use if k != "factor_code")
            cur.execute(
                f"INSERT INTO factor_metadata ({','.join(use)}) VALUES ({ph}) "
                f"ON DUPLICATE KEY UPDATE {upd}", [values[k] for k in use])
        conn.commit()
    print(f"  注册 {factor_code} -> {table}（写入列: {sorted(use)}）")
    if "factor_type" not in cols:
        print("  ⚠️ factor_metadata 无 factor_type 列——平台 /mining/factors/all 按 "
              "factor_type IN('single','multiple') 过滤，缺列可能查不到本因子，请核实。")