# -*- coding: utf-8 -*-
"""龙虎榜事件研究(向量化版, 内存安全): 游资上榜之后, 跟还是反?

读 lhb_raw.parquet(175k行, 2016-2026) + 本地daily(分年chunk, float32),
事件→forward收益用 numpy 花式索引(不 iterrows)。剔次日一字板。
"""
import sys, json, sqlite3
sys.path.insert(0, "/root/cb-allotment/scripts")
import numpy as np, pandas as pd
from lib.factor_lab import DB

D = pd.read_parquet("/root/cb-allotment/data/factor_lab/lhb_raw.parquet")
D = D.rename(columns={"代码": "code", "上榜日": "date", "解读": "note", "龙虎榜净买额": "netbuy",
                      "净买额占总成交比": "nbr", "涨跌幅": "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"])]
D["ts_code"] = D["code"] + np.where(D["code"].str[0] == "6", ".SH", ".SZ")
D["nbr"] = pd.to_numeric(D["nbr"], errors="coerce")
D["chg"] = pd.to_numeric(D["chg"], errors="coerce")
E = D.groupby(["ts_code", "date"], as_index=False).agg(
    netbuy=("netbuy", "sum"), nbr=("nbr", "mean"), chg=("chg", "first"))
del D
print(f"事件 {len(E)}", file=sys.stderr)

codes = sorted(E.ts_code.unique())
conn = sqlite3.connect(DB)
chunks = []
for yr in range(2015, 2027):
    d = pd.read_sql(f"SELECT trade_date,ts_code,pct_chg,high,low FROM daily "
                    f"WHERE trade_date>='{yr}0101' AND trade_date<='{yr}1231'", conn)
    d = d[d.ts_code.isin(codes)]
    d["pct_chg"] = d.pct_chg.astype("float32")
    d["yz"] = (d.low >= d.high - 1e-9)
    chunks.append(d[["trade_date", "ts_code", "pct_chg", "yz"]])
conn.close()
dd = pd.concat(chunks, ignore_index=True); del chunks
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
YZ = dd.pivot(index="dt", columns="ts_code", values="yz").sort_index().fillna(False).astype(bool)
del dd
tdays = RET.index
L = np.log1p(RET.fillna(0)).where(RET.notna()).astype("float32")
F1 = RET.shift(-1)
F5 = np.expm1(L.shift(-1).rolling(5).sum().shift(-4)).astype("float32")
F20 = np.expm1(L.shift(-1).rolling(20).sum().shift(-19)).astype("float32")
YZn = YZ.shift(-1).fillna(False)
print(f"面板 {RET.shape}", file=sys.stderr)

# 向量化事件对齐
col_map = {c: i for i, c in enumerate(RET.columns)}
E = E[E.ts_code.isin(col_map)]
ri = tdays.searchsorted(E.date.values)
ok = (ri < len(tdays) - 1)
E = E[ok]; ri = ri[ok]
# 上榜日若非交易日(极少) searchsorted 落到下一交易日, 统一以 ri 为事件日(上榜=当日盘后知晓, forward从ri+1起)
ci = np.array([col_map[c] for c in E.ts_code])
def take(P): return P.values[ri, ci]
E = E.assign(f1=take(F1), f5=take(F5), f20=take(F20), yz_next=take(YZn))
ET = E[~E.yz_next.astype(bool)]
print(f"可交易 {len(ET)}/{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(ET)), "span": [str(E.date.min().date()), str(E.date.max().date())]}
res["baseline"] = {"f1": round(float(np.nanmean(F1.values)) * 100, 2), "f5": round(float(np.nanmean(F5.values)) * 100, 2),
                   "f20": round(float(np.nanmean(F20.values)) * 100, 2)}
res["all"] = {k: stat(ET[k]) for k in ["f1", "f5", "f20"]}
buy = ET[ET.netbuy > 0]; sell = ET[ET.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"]}
E2 = ET[ET.nbr.notna()].copy()
E2["q"] = pd.qcut(E2.nbr.rank(method="first"), 5, labels=["Q1强卖", "Q2", "Q3", "Q4", "Q5强买"])
res["by_nbr"] = {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}
up = ET[ET.chg > 5]; dn = ET[ET.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"]}
q5 = E2[E2.q == "Q5强买"]
res["q5_f5_train"] = stat(q5[q5.date < TRAIN].f5); res["q5_f5_test"] = stat(q5[q5.date >= TRAIN].f5)
q1 = E2[E2.q == "Q1强卖"]
res["q1_f5_train"] = stat(q1[q1.date < TRAIN].f5); res["q1_f5_test"] = stat(q1[q1.date >= TRAIN].f5)

json.dump(res, open("/root/cb-allotment/my-app/public/reports/dragon-tiger/backtest_result.json", "w"), ensure_ascii=False, indent=1)
print("baseline:", res["baseline"], file=sys.stderr)
print("all:", {k: (v.get("mean"), v.get("win")) for k, v in res["all"].items()}, file=sys.stderr)
print("net_buy f1/f5/f20:", [res["net_buy"][k].get("mean") for k in ["f1", "f5", "f20"]], file=sys.stderr)
print("net_sell f1/f5/f20:", [res["net_sell"][k].get("mean") for k in ["f1", "f5", "f20"]], file=sys.stderr)
print("by_nbr f20:", {q: v["f20"].get("mean") for q, v in res["by_nbr"].items()}, file=sys.stderr)
print("bigup f5/f20:", res["listed_bigup"]["f5"].get("mean"), res["listed_bigup"]["f20"].get("mean"),
      "| bigdn:", res["listed_bigdn"]["f5"].get("mean"), res["listed_bigdn"]["f20"].get("mean"), file=sys.stderr)
print("Q5 tr/te:", res["q5_f5_train"].get("mean"), res["q5_f5_test"].get("mean"),
      "| Q1 tr/te:", res["q1_f5_train"].get("mean"), res["q1_f5_test"].get("mean"), file=sys.stderr)
print("saved", file=sys.stderr)