# -*- coding: utf-8 -*-
"""财报公告月溢价 (Frazzini-Lamont "The Earnings Announcement Premium and Trading Volume") A股版。

美股: 股票在其(可预期的)财报公告月系统性跑赢——关注度/不确定性溢价。
A股测: ①实际公告月(有前视,上限) ②可预期版=去年同月该股有公告→今年同月视为"预期公告月"(无前视)。
口径: income_period 定期报告 ann_date(季/半年/年报), 月频全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 Ctx, TRAIN_END, DB

ctx = Ctx.get()
RET = ctx.mo  # 月收益面板
months = RET.index
conn = sqlite3.connect(DB)
A = pd.read_sql("SELECT ts_code, ann_date FROM income_period WHERE ann_date IS NOT NULL AND ann_date LIKE '2%'", conn)
conn.close()
A["ym"] = A.ann_date.str[:6]
A = A.drop_duplicates(["ts_code", "ym"])
print(f"公告 stock-month {len(A)}", file=sys.stderr)

# 面板 flag: 该股该月有定期报告公告
ym_index = {f"{t.year}{t.month:02d}": t for t in months}
FLAG = pd.DataFrame(False, index=months, columns=RET.columns)
grp = A.groupby("ym")
for ym, g in grp:
    if ym not in ym_index: continue
    t = ym_index[ym]
    cols = [c for c in g.ts_code if c in FLAG.columns]
    FLAG.loc[t, cols] = True
# 可预期版: 去年同月有公告
PRED = FLAG.shift(12).fillna(False).astype(bool)

def pooled(P):
    r = RET.values.ravel(); f = P.values.ravel()
    m = ~np.isnan(r)
    r, f = r[m], f[m].astype(bool)
    return (round(float(np.nanmean(r[f])) * 100, 2), int(f.sum()),
            round(float(np.nanmean(r[~f])) * 100, 2), int((~f).sum()))

def portfolio(P):
    """每月持有 flag=1 的股票(等权) vs 全市场等权, 月度超额序列。"""
    ex = []
    for t in months:
        row = RET.loc[t]; fl = P.loc[t]
        u = row.dropna()
        sel = u[fl.reindex(u.index).fillna(False).astype(bool)]
        if len(sel) < 100 or len(u) < 500: continue
        ex.append((t, float(sel.mean() - u.mean())))
    s = pd.Series(dict(ex))
    tr = s[s.index < pd.Timestamp(TRAIN_END)]; te = s[s.index >= pd.Timestamp(TRAIN_END)]
    def ann(x): return round(float(np.nanmean(x) * 12 * 100), 2)
    t_ = float(s.mean() / s.std() * np.sqrt(len(s)))
    return {"excess_ann": ann(s), "t": round(t_, 1), "n_months": int(len(s)),
            "train": ann(tr), "test": ann(te)}

res = {}
a_in, n_in, a_out, n_out = pooled(FLAG)
res["actual"] = {"ann_month_ret": a_in, "n_in": n_in, "other_ret": a_out, "n_out": n_out,
                 "portfolio": portfolio(FLAG)}
a_in, n_in, a_out, n_out = pooled(PRED)
res["predictable"] = {"ann_month_ret": a_in, "n_in": n_in, "other_ret": a_out, "n_out": n_out,
                      "portfolio": portfolio(PRED)}
json.dump(res, open("/root/cb-allotment/research/ann_month_result.json", "w"), ensure_ascii=False, indent=1)
for k in ["actual", "predictable"]:
    r = res[k]
    print(f"{k}: 公告月 {r['ann_month_ret']}% vs 其他月 {r['other_ret']}% | 组合超额 {r['portfolio']}", file=sys.stderr)
print("saved", file=sys.stderr)