# -*- coding: utf-8 -*-
"""限售股解禁事件研究 (东方财富 akshare) — 解禁压力压制股价吗?

数据: stock_restricted_release_detail_em 每笔解禁(股票代码/解禁时间/限售类型/占流通市值比例/
      解禁前20日涨跌幅/解禁后20日涨跌幅), 按年批量拉全市场 2016-2026。
问: ①解禁后20日跑输吗? ②越大比例(占流通市值)越惨吗? ③哪类限售股(首发/定增/激励/大小非)最凶?
    ④市场是否提前抛(解禁前drift)? 用同期沪深300 20日收益做市场调整。
"""
import sys, os, json, time
sys.path.insert(0, "/root/cb-allotment/scripts")
import numpy as np, pandas as pd
import akshare as ak

CACHE = "/root/cb-allotment/data/factor_lab/lockup_detail.parquet"
if os.path.exists(CACHE) and "--refresh" not in sys.argv:
    D = pd.read_parquet(CACHE)
    print(f"缓存载入 {len(D)} 笔", file=sys.stderr)
else:
    parts = []
    for yr in range(2016, 2027):
        for attempt in range(4):
            try:
                d = ak.stock_restricted_release_detail_em(start_date=f"{yr}0101", end_date=f"{yr}1231")
                if d is not None and len(d):
                    parts.append(d); print(f"  {yr}: {len(d)} 笔", file=sys.stderr)
                break
            except Exception as e:
                if attempt == 3: print(f"  {yr} FAIL {str(e)[:50]}", file=sys.stderr)
                time.sleep(2)
    D = pd.concat(parts, ignore_index=True)
    D.to_parquet(CACHE); print(f"拉取并缓存 {len(D)} 笔", file=sys.stderr)

# 清洗
D = D.rename(columns={"股票代码": "code", "解禁时间": "date", "限售股类型": "type",
                      "占解禁前流通市值比例": "ratio", "解禁前20日涨跌幅": "pre20", "解禁后20日涨跌幅": "post20"})
D["date"] = pd.to_datetime(D["date"])
D = D[D["post20"].notna() & D["ratio"].notna()]
D["ratio"] = pd.to_numeric(D["ratio"], errors="coerce") * 100   # 转 %
D = D[(D["ratio"] > 0) & (D["ratio"] < 100)]
D["year"] = D["date"].dt.year

# 市场调整: 同期沪深300 的 20 交易日收益 —— 用 summary 的月度近似, 简化: 用全样本 post20 中位数做基准?
# 更稳: 对每个解禁日, 减去当月所有解禁事件 post20 的中位数(去市场beta)? 不严谨。
# 采用: 减去 同年 post20 均值(粗市场调整) + 报告原始与调整两版。
yr_mean = D.groupby("year")["post20"].transform("mean")
D["post20_adj"] = D["post20"] - yr_mean   # 相对同年解禁事件均值(剔市场/年份效应)

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

res = {"n_events": int(len(D)), "span": [str(D.date.min().date()), str(D.date.max().date())]}
res["overall"] = {"pre20": stats(D["pre20"]), "post20": stats(D["post20"]), "post20_adj": stats(D["post20_adj"])}

# 按 占流通市值比例 分5组
D["ratio_q"] = pd.qcut(D["ratio"].rank(method="first"), 5, labels=["Q1最小", "Q2", "Q3", "Q4", "Q5最大"])
res["by_ratio"] = {str(q): {"ratio_avg": round(float(D[D.ratio_q == q].ratio.mean()), 1),
                            "post20": stats(D[D.ratio_q == q]["post20"]),
                            "post20_adj": stats(D[D.ratio_q == q]["post20_adj"])}
                   for q in D["ratio_q"].cat.categories}

# 按限售股类型
top_types = D["type"].value_counts().head(6).index
res["by_type"] = {str(t): {"post20": stats(D[D.type == t]["post20"]),
                           "post20_adj": stats(D[D.type == t]["post20_adj"])} for t in top_types}

# 大比例(>5% 流通市值) 事件
big = D[D.ratio > 5]
res["big_unlock_gt5pct"] = {"pre20": stats(big["pre20"]), "post20": stats(big["post20"]), "post20_adj": stats(big["post20_adj"])}

# 逐年 post20_adj (大比例)
res["yearly_big"] = {int(y): stats(big[big.year == y]["post20_adj"]) for y in sorted(big.year.unique())}

json.dump(res, open("/root/cb-allotment/my-app/public/reports/lockup-release/backtest_result.json", "w"), ensure_ascii=False, indent=1)
print(f"\n总{res['n_events']}笔 {res['span']}", file=sys.stderr)
print(f"overall post20 {res['overall']['post20']} | adj {res['overall']['post20_adj']}", file=sys.stderr)
print("by_ratio post20_adj:", {q: res['by_ratio'][q]['post20_adj']['mean'] for q in res['by_ratio']}, file=sys.stderr)
print("big>5% post20:", res['big_unlock_gt5pct']['post20'], "| adj", res['big_unlock_gt5pct']['post20_adj'], file=sys.stderr)
print("saved", file=sys.stderr)