#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""migrate_constituent.py — 单元3: 合并成份股 -> constituent_unified_staging (全集型)。

设计 (spec §14, 治幸存者偏差选股池):
- baostock (bs_index_constituent 988 时点) -> 聚合成全集 (hs300/zz500/sz50 -> 000300/000905/000016),
  in_current=最后时点成份, was_removed=历史入选过但已踢出
- 深证/国证 _union.parquet (399001/399006/399005/399330) -> 直入 (akshare cni)
- 中证1000/2000 _snapshot.parquet (000852/932000) -> 当前 (akshare csindex)
- 新浪 _sina.parquet -> 丢弃 (baostock 300/500/50 已权威)
- code 统一 6 位无前缀 (sh.600000 / 600519.SH -> 600519)

输出 constituent_unified_staging(index_code, code, code_name, source, in_current, was_removed)。
验证 OK 后手动 rename: staging->constituent_unified, bs_index_constituent->_old。
"""
import glob
import os
import re
import sqlite3

import pandas as pd

DB = r"C:\sanguo_vnpy_v2\data\quant_trading.db"
HIST = r"C:\sanguo_vnpy_v2\data\index_const_hist"
STAGING = "constituent_unified_staging"
BS_MAP = {"hs300": "000300", "zz500": "000905", "sz50": "000016"}


def norm_code(code):
    s = re.sub(r"^(sh|sz|SH|SZ)\.?", "", str(code))
    s = re.sub(r"\.(SH|SZ|sh|sz)$", "", s)
    return s.zfill(6) if s.isdigit() and len(s) <= 6 else s


def main():
    c = sqlite3.connect(DB, timeout=60)
    c.execute("PRAGMA busy_timeout = 60000")

    # 1. baostock -> 全集
    df_bs = pd.read_sql(
        "SELECT updateDate, index_code, code, code_name FROM bs_index_constituent", c)
    df_bs["index_code"] = df_bs["index_code"].map(BS_MAP)
    df_bs["code"] = df_bs["code"].apply(norm_code)
    last_sets = {}
    for idx, grp in df_bs.groupby("index_code"):
        last_d = grp["updateDate"].max()
        last_sets[idx] = set(grp[grp["updateDate"] == last_d]["code"])
    pool = (df_bs.groupby(["index_code", "code"])["code_name"]
            .first().reset_index())
    pool["in_current"] = pool.apply(
        lambda r: r["code"] in last_sets.get(r["index_code"], set()), axis=1)
    pool["was_removed"] = ~pool["in_current"]
    pool["source"] = "baostock"
    print(f"[baostock] pool rows={len(pool)} (300/500/50 全集)")

    # 2. 深证 union
    deep = []
    for f in sorted(glob.glob(os.path.join(HIST, "*_union.parquet"))):
        d = pd.read_parquet(f)[["index_code", "code", "code_name",
                                "in_current", "was_removed"]]
        d["source"] = "akshare_cni"
        deep.append(d)
    df_deep = pd.concat(deep, ignore_index=True) if deep else pd.DataFrame(
        columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
    df_deep["code"] = df_deep["code"].apply(norm_code)
    print(f"[深证 union] rows={len(df_deep)}")

    # 3. 中证 snapshot
    snap = []
    for f in [os.path.join(HIST, "000852_snapshot.parquet"),
              os.path.join(HIST, "932000_snapshot.parquet")]:
        if os.path.exists(f):
            d = pd.read_parquet(f)[["index_code", "code", "code_name"]]
            d["in_current"] = True
            d["was_removed"] = False
            d["source"] = "akshare_csindex"
            snap.append(d)
    df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame(
        columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"])
    df_snap["code"] = df_snap["code"].apply(norm_code)
    print(f"[中证 snapshot] rows={len(df_snap)}")

    # 合并 + 去重 (同 index+code+source)
    all_df = pd.concat([pool, df_deep, df_snap], ignore_index=True)
    all_df = all_df.drop_duplicates(["index_code", "code", "source"])
    print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, "
          f"{all_df['index_code'].nunique()} indices")

    print("\n各指数分布:")
    print(all_df.groupby("index_code").agg(
        n=("code", "count"), src=("source", "first"),
        in_cur=("in_current", "sum"), removed=("was_removed", "sum")))

    # 写 staging
    c.execute(f"DROP TABLE IF EXISTS {STAGING}")
    c.execute(f"""CREATE TABLE {STAGING} (
        index_code TEXT, code TEXT, code_name TEXT, source TEXT,
        in_current INTEGER, was_removed INTEGER)""")
    work = all_df[["index_code", "code", "code_name", "source",
                   "in_current", "was_removed"]].copy()
    work["in_current"] = work["in_current"].astype(int)
    work["was_removed"] = work["was_removed"].astype(int)
    c.executemany(f"INSERT INTO {STAGING} VALUES (?,?,?,?,?,?)",
                  work.itertuples(index=False, name=None))
    c.commit()
    n = c.execute(f"SELECT COUNT(*) FROM {STAGING}").fetchone()[0]
    # 抽样验证
    print(f"\n[staging] {STAGING}: {n} rows")
    print("sample 300:", c.execute(
        "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
        f"{STAGING} WHERE index_code='000300'").fetchone())
    print("sample 399001:", c.execute(
        "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM "
        f"{STAGING} WHERE index_code='399001'").fetchone())
    print("sample 000852:", c.execute(
        "SELECT COUNT(*) FROM " f"{STAGING} WHERE index_code='000852'").fetchone())
    c.close()
    print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)")


if __name__ == "__main__":
    main()