# -*- coding: utf-8 -*-
"""涨停板研究: 打板到底赚不赚钱? (全A 本地日线, 零抓取)

事件: 涨停(收盘=最高 且 pct_chg 达板)/一字板(low=high)/换手板(非一字)/炸板(盘中触板收盘未板)。
板幅按代码+日期: 688→20%; 300→2020-08-24 前 10% 后 20%; 4/8/9(北交)剔除; 其余 10%。
ST 5% 档无历史名单, 统一不识别(口径注明)。
连板数 = 连续涨停天数。
问: ①首板次日/后5日/20日? ②连板越高越危险? ③一字 vs 换手板? ④炸板多惨? ⑤打板者(收盘价买)能赚吗?
基准: 全市场当日均值。train<2022/test。
"""
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

conn = sqlite3.connect(DB)
# 分年读, 只留必要列, float32
chunks = []
for yr in range(2015, 2027):
    d = pd.read_sql(f"SELECT trade_date,ts_code,open,high,low,close,pre_close,pct_chg FROM daily "
                    f"WHERE trade_date>='{yr}0101' AND trade_date<='{yr}1231'", conn)
    if not len(d): continue
    d = d[(d.pre_close > 0) & (d.close > 0) & ~d.ts_code.str.endswith(".BJ")]
    d = d[~d.ts_code.str[0].isin(["4", "8", "9"])]
    for c in ["open", "high", "low", "close", "pre_close", "pct_chg"]:
        d[c] = d[c].astype("float32")
    chunks.append(d)
    print(f"  {yr}: {len(d)}", file=sys.stderr)
conn.close()
D = pd.concat(chunks, ignore_index=True); del chunks
D["dt"] = pd.to_datetime(D.trade_date, format="%Y%m%d")

# 板幅
pfx3 = D.ts_code.str[:3]
lim = np.where(pfx3 == "688", 0.20,
      np.where((pfx3 == "300") & (D.trade_date >= "20200824"), 0.20, 0.10)).astype("float32")
r = D.pct_chg.values / 100
# 涨停: 收盘=最高 且 涨幅≥板幅-0.5pp (交易所取整误差)
is_lu = (D.close.values >= D.high.values - 1e-6) & (r >= lim - 0.005)
# 一字: low=high(全天一价) 且涨停
is_yz = is_lu & (D.low.values >= D.high.values - 1e-6)
# 炸板: 盘中触板(high 涨幅≥板幅-0.5pp) 但收盘未涨停
touched = (D.high.values / D.pre_close.values - 1) >= (lim - 0.005)
is_zb = touched & (~is_lu)
D["lu"] = is_lu; D["yz"] = is_yz; D["zb"] = is_zb
print(f"总样本 {len(D)} 涨停 {is_lu.sum()} 一字 {is_yz.sum()} 炸板 {is_zb.sum()}", file=sys.stderr)

# 面板化 (float32/bool)
RET = D.pivot(index="dt", columns="ts_code", values="pct_chg").sort_index() / 100
LU = D.pivot(index="dt", columns="ts_code", values="lu").sort_index().fillna(False).astype(bool)
YZ = D.pivot(index="dt", columns="ts_code", values="yz").sort_index().fillna(False).astype(bool)
ZB = D.pivot(index="dt", columns="ts_code", values="zb").sort_index().fillna(False).astype(bool)
OPEN = D.pivot(index="dt", columns="ts_code", values="open").sort_index()
CLOSE = D.pivot(index="dt", columns="ts_code", values="close").sort_index()
del D

# 连板数: 迭代行
lb = np.zeros(LU.shape, dtype=np.int16)
luv = LU.values
for i in range(1, LU.shape[0]):
    lb[i] = np.where(luv[i], lb[i - 1] + 1, 0)
lb[0] = luv[0].astype(np.int16)
LB = pd.DataFrame(lb, index=LU.index, columns=LU.columns)

# forward 收益(log 累加)
L = np.log1p(RET.fillna(0)).where(RET.notna())
def fwd_cum(n):
    # t+1..t+n 累计
    return np.expm1(L.shift(-1).rolling(n, min_periods=n).sum().shift(-(n - 1)))
F1 = RET.shift(-1)
F5 = fwd_cum(5)
F20 = fwd_cum(20)
# 次日开盘溢价(隔夜): open_{t+1}/close_t - 1
GAP1 = OPEN.shift(-1) / CLOSE - 1
# 次日也一字涨停(买不进) 概率用 YZ.shift(-1)
YZ_next = YZ.shift(-1).fillna(False)

TRAIN = pd.Timestamp("2022-01-01")
def stat_mask(mask, F):
    v = F.values[mask.values]
    v = v[~np.isnan(v)]
    if len(v) < 50: return {}
    return {"n": int(len(v)), "mean": round(float(np.mean(v)) * 100, 2),
            "median": round(float(np.median(v)) * 100, 2), "win": round(float((v > 0).mean()) * 100, 0)}
def full_stats(mask):
    out = {"gap1": stat_mask(mask, GAP1), "f1": stat_mask(mask, F1), "f5": stat_mask(mask, F5), "f20": stat_mask(mask, F20)}
    # train/test on f5
    mtr = mask.copy(); mtr.loc[mtr.index >= TRAIN] = False
    mte = mask.copy(); mte.loc[mte.index < TRAIN] = False
    out["f5_train"] = stat_mask(mtr, F5); out["f5_test"] = stat_mask(mte, F5)
    return out

res = {}
res["baseline_all"] = {"f1": stat_mask(RET.notna(), F1), "f5": stat_mask(RET.notna(), F5), "f20": stat_mask(RET.notna(), F20)}
# 按连板数(换手板=非一字, 可成交)
HB = LU & ~YZ   # 换手板(收盘能买到)
for k, nm in [(1, "首板"), (2, "二板"), (3, "三板")]:
    res[f"{nm}_换手板"] = full_stats(HB & (LB == k))
res["高位板4+_换手板"] = full_stats(HB & (LB >= 4))
res["一字板_全部"] = full_stats(YZ)
res["一字板_次日仍一字概率"] = round(float(YZ_next.values[YZ.values].mean()) * 100, 1)
res["炸板"] = full_stats(ZB)
res["涨停总数"] = int(LU.values.sum()); res["一字数"] = int(YZ.values.sum()); res["炸板数"] = int(ZB.values.sum())

json.dump(res, open("/root/cb-allotment/my-app/public/reports/limit-up/backtest_result.json", "w"), ensure_ascii=False, indent=1)
for k in ["baseline_all", "首板_换手板", "二板_换手板", "三板_换手板", "高位板4+_换手板", "一字板_全部", "炸板"]:
    v = res[k]
    print(f"{k}: gap1 {v.get('gap1',{}).get('mean','-')} | f1 {v['f1'].get('mean','-')} | f5 {v['f5'].get('mean','-')}(win{v['f5'].get('win','-')}) | f20 {v['f20'].get('mean','-')}", file=sys.stderr)
print(f"一字次日仍一字: {res['一字板_次日仍一字概率']}%", file=sys.stderr)
print("saved", file=sys.stderr)