# -*- coding: utf-8 -*-
"""股权质押风险 (议程⑤): 高质押比例的股票跑输吗? 2018爆仓 regime vs 现在。

数据: 东财质押比例周度快照(akshare stock_gpzy_pledge_ratio_em, date=周五)。
拉每季度末最近周五 2018Q1-2026Q2 (~34快照, 缓存parquet); forward 3月收益本地daily自算。
分组: 零质押 / 有质押分5档; 逐期横截面 → 池化; 分 regime(2018-19 vs 2020+) + train/test。
"""
import sys, os, json, time
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
import sqlite3

CACHE = "/root/cb-allotment/data/factor_lab/pledge_snapshots.parquet"
# 季度末最近的周五
qends = pd.date_range("2018-03-31", "2026-06-30", freq="QE")
fridays = [(q - pd.Timedelta(days=(q.weekday() - 4) % 7)).strftime("%Y%m%d") for q in qends]

if os.path.exists(CACHE) and "--refresh" not in sys.argv:
    P = pd.read_parquet(CACHE); print(f"缓存 {len(P)} 行 {P.snap.nunique()} 快照", file=sys.stderr)
else:
    parts = []
    for dt in fridays:
        got = None
        for a in range(3):
            try:
                d = ak.stock_gpzy_pledge_ratio_em(date=dt)
                if d is not None and len(d) > 100: got = d; break
            except Exception:
                time.sleep(3)
        if got is None:
            # 往前退一周再试
            dt2 = (pd.Timestamp(dt) - pd.Timedelta(days=7)).strftime("%Y%m%d")
            try:
                d = ak.stock_gpzy_pledge_ratio_em(date=dt2)
                if d is not None and len(d) > 100: got = d; dt = dt2
            except Exception: pass
        if got is not None:
            got = got.rename(columns={"股票代码": "code", "质押比例": "ratio"})
            got["code"] = got["code"].astype(str).str.zfill(6)
            got["ratio"] = pd.to_numeric(got["ratio"], errors="coerce")
            got = got[["code", "ratio"]].dropna()
            got["snap"] = dt
            parts.append(got)
            print(f"  {dt}: {len(got)}", file=sys.stderr)
        time.sleep(1)
    P = pd.concat(parts, ignore_index=True)
    P.to_parquet(CACHE); print(f"拉取 {len(P)} 行 {P.snap.nunique()} 快照", file=sys.stderr)

P = P[~P.code.str[0].isin(["4", "8", "9"])]
P["ts_code"] = P.code + np.where(P.code.str[0] == "6", ".SH", ".SZ")

conn = sqlite3.connect(DB)
codes = sorted(P.ts_code.unique())
chunks = []
for yr in range(2018, 2027):
    d = pd.read_sql(f"SELECT trade_date,ts_code,pct_chg FROM daily 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")
    chunks.append(d)
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
del dd
tdays = RET.index
L = np.log1p(RET.fillna(0)).where(RET.notna()).astype("float32")
F63 = np.expm1(L.shift(-1).rolling(63).sum().shift(-62))  # 未来一季(~63交易日)
col_map = {c: i for i, c in enumerate(RET.columns)}
print(f"面板 {RET.shape}", file=sys.stderr)

rows = []
for snap, g in P.groupby("snap"):
    t = pd.Timestamp(snap)
    pos = tdays.searchsorted(t)
    if pos >= len(tdays) - 63: continue
    g = g[g.ts_code.isin(col_map)]
    ci = np.array([col_map[c] for c in g.ts_code])
    fwd = F63.values[pos, ci]
    mkt = float(np.nanmean(F63.values[pos]))   # 全市场同窗
    for r, f in zip(g.ratio.values, fwd):
        if np.isnan(f): continue
        rows.append({"snap": snap, "ratio": r, "fwd": f, "ex": f - mkt})
E = pd.DataFrame(rows)
E["year"] = E.snap.str[:4].astype(int)
print(f"样本 {len(E)}", file=sys.stderr)

def stat(s):
    s = pd.Series(s).dropna()
    if len(s) < 200: 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)}

# 分档: 0质押不在名单里(EM只列有质押的) → 档内5分位
E["q"] = pd.qcut(E.ratio.rank(method="first"), 5, labels=["Q1最低", "Q2", "Q3", "Q4", "Q5最高"])
res = {"n_obs": int(len(E)), "n_snaps": int(E.snap.nunique()), "span": [E.snap.min(), E.snap.max()]}
res["by_q"] = {str(q): {"ratio_avg": round(float(E[E.q == q].ratio.mean()), 1), **stat(E[E.q == q].ex)}
               for q in E["q"].cat.categories}
res["regime_2018_19_q5"] = stat(E[(E.year <= 2019) & (E.q == "Q5最高")].ex)
res["regime_2020p_q5"] = stat(E[(E.year >= 2020) & (E.q == "Q5最高")].ex)
res["regime_2018_19_q1"] = stat(E[(E.year <= 2019) & (E.q == "Q1最低")].ex)
res["regime_2020p_q1"] = stat(E[(E.year >= 2020) & (E.q == "Q1最低")].ex)
# 极高质押(>50%)
res["gt50"] = stat(E[E.ratio > 50].ex); res["gt50_2018_19"] = stat(E[(E.ratio > 50) & (E.year <= 2019)].ex)
res["gt50_2020p"] = stat(E[(E.ratio > 50) & (E.year >= 2020)].ex)

json.dump(res, open("/root/cb-allotment/my-app/public/reports/pledge-risk/backtest_result.json", "w"), ensure_ascii=False, indent=1)
print("by_q ex(mean):", {q: v.get("mean") for q, v in res["by_q"].items()}, file=sys.stderr)
print("Q5 2018-19:", res["regime_2018_19_q5"], file=sys.stderr)
print("Q5 2020+:", res["regime_2020p_q5"], file=sys.stderr)
print(">50%质押 全期/18-19/20+:", res["gt50"].get("mean"), res["gt50_2018_19"].get("mean"), res["gt50_2020p"].get("mean"), file=sys.stderr)
print("saved", file=sys.stderr)