# -*- coding: utf-8 -*-
"""龙虎榜事件研究: 游资上榜之后, 跟还是反? (东财 akshare + 本地日线)

数据: stock_lhb_detail_em(start,end) 按月批量 2016-2026(上榜日/净买额/买入卖出额/解读),
     缓存 parquet; 收益用本地 daily(全市场复权 pct_chg)自算(解禁研究同款)。
问: ①上榜股 次日/5日/20日 vs 全市场? ②净买入榜 vs 净卖出榜? ③净买额占成交比分档?
   ④机构买入(解读含"机构")与游资有别吗? 剔次日一字板(买不进)。
"""
import sys, os, json, time, sqlite3
sys.path.insert(0, "/root/cb-allotment/scripts")
import numpy as np, pandas as pd
import akshare as ak
from lib.factor_lab import DB

RAW = "/root/cb-allotment/data/factor_lab/lhb_raw.parquet"
if os.path.exists(RAW) and "--refresh" not in sys.argv:
    D = pd.read_parquet(RAW); print(f"缓存 {len(D)} 条", file=sys.stderr)
else:
    parts = []
    rng = pd.date_range("2016-01-01", "2026-07-01", freq="MS")
    for i in range(len(rng) - 1):
        s_, e_ = rng[i].strftime("%Y%m%d"), (rng[i + 1] - pd.Timedelta(days=1)).strftime("%Y%m%d")
        got = None
        for a in range(4):
            try:
                d = ak.stock_lhb_detail_em(start_date=s_, end_date=e_)
                if d is not None: got = d; break
            except Exception:
                time.sleep(2 + a)
        if got is not None and len(got):
            parts.append(got)
            if i % 12 == 0: print(f"  {s_[:6]}: 累计 {sum(len(x) for x in parts)}", file=sys.stderr)
        time.sleep(0.5)
    D = pd.concat(parts, ignore_index=True).drop_duplicates()
    D.to_parquet(RAW); print(f"拉取 {len(D)} 条", file=sys.stderr)

D = D.rename(columns={"代码": "code", "名称": "name", "上榜日": "date", "解读": "note",
                      "龙虎榜净买额": "netbuy", "净买额占总成交比": "netbuy_ratio", "涨跌幅": "chg"})
D["date"] = pd.to_datetime(D["date"])
D["code"] = D["code"].astype(str).str.zfill(6)
D = D[~D["code"].str[0].isin(["4", "8", "9"])]
def suf(c): return c + (".SH" if c[0] == "6" else ".SZ")
D["ts_code"] = D["code"].map(suf)
D["netbuy_ratio"] = pd.to_numeric(D["netbuy_ratio"], errors="coerce")
# 同股同日多榜合并(取净买额合计/占比均值)
D = D.groupby(["ts_code", "date"], as_index=False).agg(
    netbuy=("netbuy", "sum"), netbuy_ratio=("netbuy_ratio", "mean"),
    chg=("chg", "first"), note=("note", "first"))
print(f"事件 {len(D)} (合并同日) {D.date.min().date()}..{D.date.max().date()}", file=sys.stderr)

# 本地日线
conn = sqlite3.connect(DB)
codes = D.ts_code.unique().tolist()
qs = ",".join(f"'{c}'" for c in codes)
dd = pd.read_sql(f"SELECT trade_date,ts_code,close,pre_close,pct_chg,open,high,low FROM daily WHERE ts_code IN ({qs}) AND trade_date>='20151101'", conn)
conn.close()
dd["dt"] = pd.to_datetime(dd.trade_date, format="%Y%m%d")
RET = dd.pivot(index="dt", columns="ts_code", values="pct_chg").sort_index() / 100
OPEN = dd.pivot(index="dt", columns="ts_code", values="open").sort_index()
LOW = dd.pivot(index="dt", columns="ts_code", values="low").sort_index()
HIGH = dd.pivot(index="dt", columns="ts_code", values="high").sort_index()
tdays = RET.index
L = np.log1p(RET.fillna(0)).where(RET.notna())
F1 = RET.shift(-1)
F5 = np.expm1(L.shift(-1).rolling(5).sum().shift(-4))
F20 = np.expm1(L.shift(-1).rolling(20).sum().shift(-19))
# 次日一字(买不进): low==high 次日
YZ_next = (LOW == HIGH).shift(-1)

def get(P, code, dt):
    try:
        v = P.at[dt, code]
        return float(v) if pd.notna(v) else np.nan
    except KeyError:
        return np.nan

rows = []
for _, r in D.iterrows():
    if r.ts_code not in RET.columns: continue
    if r.date not in tdays: continue
    yz = get(YZ_next, r.ts_code, r.date)
    rows.append({"ts": r.ts_code, "dt": r.date, "netbuy": r.netbuy, "nbr": r.netbuy_ratio,
                 "chg": r.chg, "yz_next": bool(yz) if yz == yz else False,
                 "f1": get(F1, r.ts_code, r.date), "f5": get(F5, r.ts_code, r.date), "f20": get(F20, r.ts_code, r.date)})
E = pd.DataFrame(rows)
E_tradable = E[~E.yz_next]   # 剔次日一字
print(f"可交易事件 {len(E_tradable)}/{len(E)}", file=sys.stderr)

def stat(s):
    s = pd.Series(s).dropna()
    if len(s) < 100: return {}
    return {"n": int(len(s)), "mean": round(float(s.mean()) * 100, 2), "median": round(float(s.median()) * 100, 2),
            "win": round(float((s > 0).mean()) * 100, 0)}

TRAIN = pd.Timestamp("2022-01-01")
res = {"n_events": int(len(E)), "n_tradable": int(len(E_tradable)),
       "span": [str(E.dt.min().date()), str(E.dt.max().date())]}
# 基准: 全市场
res["baseline_f5"] = {"mean": round(float(np.nanmean(F5.values)) * 100, 2)}
res["baseline_f20"] = {"mean": round(float(np.nanmean(F20.values)) * 100, 2)}
# 全部上榜
res["all"] = {k: stat(E_tradable[k]) for k in ["f1", "f5", "f20"]}
# 净买 vs 净卖
buy = E_tradable[E_tradable.netbuy > 0]; sell = E_tradable[E_tradable.netbuy <= 0]
res["net_buy"] = {k: stat(buy[k]) for k in ["f1", "f5", "f20"]}
res["net_sell"] = {k: stat(sell[k]) for k in ["f1", "f5", "f20"]}
# 净买占比分5档
E2 = E_tradable[E_tradable.nbr.notna()].copy()
E2["q"] = pd.qcut(E2.nbr.rank(method="first"), 5, labels=["Q1强卖", "Q2", "Q3", "Q4", "Q5强买"])
res["by_netbuy_ratio"] = {str(q): {"nbr_avg": round(float(E2[E2.q == q].nbr.mean()), 1),
                                    "f1": stat(E2[E2.q == q].f1), "f5": stat(E2[E2.q == q].f5), "f20": stat(E2[E2.q == q].f20)}
                          for q in E2["q"].cat.categories}
# 上榜且当日涨(涨停上榜) vs 当日跌(跌停/大跌上榜)
up = E_tradable[E_tradable.chg > 5]; dn = E_tradable[E_tradable.chg < -5]
res["listed_bigup"] = {k: stat(up[k]) for k in ["f1", "f5", "f20"]}
res["listed_bigdn"] = {k: stat(dn[k]) for k in ["f1", "f5", "f20"]}
# train/test (Q5强买 f5)
q5 = E2[E2.q == "Q5强买"]
res["q5_f5_train"] = stat(q5[q5.dt < TRAIN].f5); res["q5_f5_test"] = stat(q5[q5.dt >= TRAIN].f5)

json.dump(res, open("/root/cb-allotment/research/lhb_result.json", "w"), ensure_ascii=False, indent=1)
print("all:", res["all"], file=sys.stderr)
print("net_buy f5/f20:", res["net_buy"].get("f5"), res["net_buy"].get("f20"), file=sys.stderr)
print("by_nbr f20:", {q: v["f20"].get("mean") for q, v in res["by_netbuy_ratio"].items()}, file=sys.stderr)
print("saved", file=sys.stderr)