# -*- coding: utf-8 -*-
"""业绩预告事件研究: 报喜追得上吗, 暴雷逃得掉吗? (本地 forecast 缓存 + 全市场日线)

数据: fundamental_cache.db forecast 表(47896 条, 2019-2026, ann_date/type/p_change_min/max)。
事件窗(相对公告日 t=0, 公告多在盘后→次日反应):
  pre = [-20,-1] 公告前(提前泄露/抢跑?)  jump = [1,2] 公告后即时反应
  drift = [3,20] 追/逃窗口 —— 看到公告再行动, 还有肉/还来得及吗?
分组: 报喜(预增/扭亏/略增/续盈) vs 暴雷(预减/首亏/续亏/略减); 预增再按幅度(p_change中点)分档。
剔事件次日一字板(买不进/卖不出的近似)。市场调整: 减同窗全A均值。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)
F = pd.read_sql("SELECT ann_date, ts_code, end_date, type, p_change_min, p_change_max FROM forecast", conn)
F["ann"] = pd.to_datetime(F.ann_date, format="%Y%m%d")
F = F[~F.ts_code.str[0].isin(["4", "8", "9"])]
F = F.drop_duplicates(subset=["ts_code", "ann_date", "end_date"])   # 去重坑
GOOD = {"预增", "扭亏", "略增", "续盈"}; BAD = {"预减", "首亏", "续亏", "略减"}
F["dir"] = np.where(F.type.isin(GOOD), "good", np.where(F.type.isin(BAD), "bad", "na"))
F = F[F["dir"] != "na"]
F["pmid"] = (pd.to_numeric(F.p_change_min, errors="coerce") + pd.to_numeric(F.p_change_max, errors="coerce")) / 2
print(f"事件 {len(F)} good {(F['dir']=='good').sum()} bad {(F['dir']=='bad').sum()}", file=sys.stderr)

codes = F.ts_code.unique().tolist()
qs = ",".join(f"'{c}'" for c in codes)
dd = pd.read_sql(f"SELECT trade_date,ts_code,pct_chg,high,low FROM daily WHERE ts_code IN ({qs}) AND trade_date>='20181001'", 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
YZ = (dd.pivot(index="dt", columns="ts_code", values="low") == dd.pivot(index="dt", columns="ts_code", values="high")).sort_index()
tdays = RET.index
L = np.log1p(RET.fillna(0)).where(RET.notna())
MKT = RET.mean(axis=1)   # 全A等权日收益
LM = np.log1p(MKT)

def win_ret(code, pos, a, b):
    """[a,b] 交易日窗口累计(相对事件日 pos), 超界 nan。"""
    i, j = pos + a, pos + b
    if i < 0 or j >= len(tdays): return np.nan, np.nan
    seg = L[code].iloc[i:j + 1]
    if seg.isna().all(): return np.nan, np.nan
    r = float(np.expm1(np.nansum(seg)))
    m = float(np.expm1(LM.iloc[i:j + 1].sum()))
    return r, r - m

rows = []
Lcols = set(L.columns)
tdset = set(tdays)
for _, e in F.iterrows():
    if e.ts_code not in Lcols: continue
    pos = tdays.searchsorted(e.ann)
    if pos >= len(tdays) - 3: continue
    # 反应日对齐: 交易日盘后公告→次日反应(pos+1); 非交易日公告→searchsorted已落到下一交易日(pos)
    rx = pos + 1 if e.ann in tdset else pos
    pre, pre_a = win_ret(e.ts_code, rx, -21, -2)      # 公告前20日(不含反应日前一日贴近泄露窗)
    jump, jump_a = win_ret(e.ts_code, rx, 0, 1)        # 反应日+次日
    drift, drift_a = win_ret(e.ts_code, rx, 2, 19)     # 之后18日 = 看到公告还追/逃得及吗
    try: yz = bool(YZ[e.ts_code].iloc[rx])
    except Exception: yz = False
    rows.append({"dir": e["dir"], "type": e.type, "pmid": e.pmid, "ann": e.ann, "yz": yz,
                 "pre": pre_a, "jump": jump_a, "drift": drift_a, "drift_raw": drift})
E = pd.DataFrame(rows)
print(f"有效事件 {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))}
for d in ["good", "bad"]:
    sub = E[(E["dir"] == d) & (~E.yz)]   # 剔次日一字
    res[d] = {w: stat(sub[w]) for w in ["pre", "jump", "drift"]}
    res[d]["drift_train"] = stat(sub[sub.ann < TRAIN].drift); res[d]["drift_test"] = stat(sub[sub.ann >= TRAIN].drift)
# 分type
res["by_type_drift"] = {t: stat(E[(E.type == t) & (~E.yz)].drift) for t in list(GOOD | BAD)}
# 预增按幅度
yz_ok = E[(E.type == "预增") & (~E.yz) & E.pmid.notna()].copy()
if len(yz_ok) > 500:
    yz_ok["q"] = pd.qcut(yz_ok.pmid.rank(method="first"), 3, labels=["小增", "中增", "大增"])
    res["yuzeng_by_mag"] = {str(q): {"pmid_avg": round(float(yz_ok[yz_ok.q == q].pmid.mean()), 0),
                                     "jump": stat(yz_ok[yz_ok.q == q].jump), "drift": stat(yz_ok[yz_ok.q == q].drift)}
                            for q in yz_ok["q"].cat.categories}
res["yz_next_prob"] = {"good": round(float(E[E["dir"]=="good"].yz.mean())*100,1), "bad": round(float(E[E["dir"]=="bad"].yz.mean())*100,1)}

json.dump(res, open("/root/cb-allotment/my-app/public/reports/earnings-forecast/backtest_result.json", "w"), ensure_ascii=False, indent=1, default=str)
for d in ["good", "bad"]:
    r = res[d]
    print(f"{d}: pre {r['pre'].get('mean')} | jump {r['jump'].get('mean')} | drift {r['drift'].get('mean')}(win{r['drift'].get('win')}) tr{r['drift_train'].get('mean')}/te{r['drift_test'].get('mean')}", file=sys.stderr)
print("by_type drift:", {k: v.get("mean") for k, v in res["by_type_drift"].items()}, file=sys.stderr)
if "yuzeng_by_mag" in res:
    print("预增分档 drift:", {k: v["drift"].get("mean") for k, v in res["yuzeng_by_mag"].items()}, file=sys.stderr)
print("saved", file=sys.stderr)