# -*- coding: utf-8 -*-
"""隔夜 vs 日内收益的「拔河」(Lou, Polk & Skouras 2019 JFE "A tug of war")A股全市场版。

每日收益拆两段:
  隔夜(跳空) overnight = open/pre_close - 1   (pre_close=前复权前收, 已含除权除息 → 干净跳空)
  日内       intraday  = close/open - 1
  合计 total = (1+on)(1+in)-1 = close/pre_close-1 = pct_chg/100

问题:哪一段可预测?美股结论——隔夜动量持续(散户情绪), 日内反转(机构)。A股散户更重, 看会怎样。
月频. 全A(剔ST/次新/低价/低流动). size中性分位. 长仓口径. train<2022/test OOS.
"""
import sys, json, sqlite3, numpy as np, pandas as pd
sys.path.insert(0, "/root/cb-allotment/scripts")
from lib.factor_lab import DB

# 分年读取聚合(峰值内存 ~1年, 避 OOM: 7GB 限制)
conn = sqlite3.connect(DB)
chunks = []
for yr in range(2016, 2027):
    d = pd.read_sql(f"SELECT trade_date,ts_code,open,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.open > 0) & (d.pre_close > 0) & (d.close > 0)]
    on = (d.open / d.pre_close - 1.0); inn = (d.close / d.open - 1.0)
    keep = (on.abs() < 0.2) & (inn.abs() < 0.2)
    mm = pd.DataFrame({
        "ts_code": d.ts_code[keep].values,
        "mk": (d.trade_date[keep].astype(np.int64) // 100).astype(np.int32).values,
        "lon": np.log1p(on[keep]).values, "lin": np.log1p(inn[keep]).values,
        "ltot": np.log1p(d.pct_chg[keep].values / 100), "close": d.close[keep].values,
    })
    g = mm.groupby(["ts_code", "mk"], sort=False).agg(
        lon=("lon", "sum"), lin=("lin", "sum"), ltot=("ltot", "sum"),
        px=("close", "last"), ndays=("close", "size")).reset_index()
    chunks.append(g)
    print(f"  {yr}: {len(d)}日行 -> {len(g)}月行", file=sys.stderr)
    del d, on, inn, keep, mm, g
conn.close()
mon = pd.concat(chunks, ignore_index=True); del chunks
mon = mon[mon.ndays >= 10].copy()
mon["m"] = pd.PeriodIndex(mon["mk"].astype(str), freq="M")
mon["on_m"] = np.expm1(mon["lon"]); mon["in_m"] = np.expm1(mon["lin"]); mon["tot_m"] = np.expm1(mon["ltot"])

def piv(col): return mon.pivot(index="m", columns="ts_code", values=col).sort_index()
ON, IN, TOT, PX = piv("on_m"), piv("in_m"), piv("tot_m"), piv("px")
months = ON.index
print(f"面板 {ON.shape} {months.min()}..{months.max()}", file=sys.stderr)

# universe: 价格>2, 剔ST(按名不可得则用价格+流动), 该月有数据
uni = (PX > 2) & TOT.notna()

def zscore(df):  # 截面 rank->[-0.5,0.5] (稳健, 免winsor)
    r = df.rank(axis=1)
    return r.sub(r.mean(axis=1), axis=0).div(r.count(axis=1), axis=0)

def decile_fwd(signal, groups=10):
    """signal[m] 预测 m+1 的 TOT/ON/IN. 返回各组 forward 三段年化 + 单调性。"""
    fwd_tot, fwd_on, fwd_in = TOT.shift(-1), ON.shift(-1), IN.shift(-1)
    rows = {g: {"tot": [], "on": [], "in": []} for g in range(groups)}
    spread_series = []
    for m in months[:-1]:
        s = signal.loc[m][uni.loc[m]].dropna()
        if len(s) < 300: continue
        q = pd.qcut(s.rank(method="first"), groups, labels=False)
        ft, fo, fi = fwd_tot.loc[m], fwd_on.loc[m], fwd_in.loc[m]
        gm = {}
        for gi in range(groups):
            codes = s.index[q == gi]
            rows[gi]["tot"].append(ft[codes].mean())
            rows[gi]["on"].append(fo[codes].mean())
            rows[gi]["in"].append(fi[codes].mean())
            gm[gi] = ft[codes].mean()
        spread_series.append((m, gm[groups - 1] - gm[0]))
    def ann(v): return float(np.nanmean(v) * 12 * 100)
    out = {"groups": []}
    for gi in range(groups):
        out["groups"].append({"g": gi + 1, "tot": round(ann(rows[gi]["tot"]), 1),
                              "on": round(ann(rows[gi]["on"]), 1), "in": round(ann(rows[gi]["in"]), 1)})
    sp = pd.Series(dict(spread_series))
    tstat = sp.mean() / sp.std() * np.sqrt(len(sp))
    # 单调性 spearman
    tots = [x["tot"] for x in out["groups"]]
    rho = pd.Series(tots).corr(pd.Series(range(groups)), method="spearman")
    out["spread_ann"] = round(ann(sp), 1); out["spread_t"] = round(tstat, 1)
    out["monotonic_rho"] = round(rho, 2)
    tr = sp[sp.index < pd.Period("2022-01")]; te = sp[sp.index >= pd.Period("2022-01")]
    out["spread_train"] = round(ann(tr), 1); out["spread_test"] = round(ann(te), 1)
    return out

# 信号: 上月隔夜 / 上月日内 / 6月累计隔夜 / 6月累计日内
sig_on1 = ON
sig_in1 = IN
sig_on6 = np.expm1(np.log1p(ON).rolling(6).sum())
sig_in6 = np.expm1(np.log1p(IN).rolling(6).sum())

# 组合信号: 隔夜动量(+) 减 日内动量(-) = 拔河两端各取有效方向 (long-only)
combo = zscore(sig_on6) - zscore(sig_in1)

result = {"panel": [str(ON.shape[0]), str(ON.shape[1]), str(months.min()), str(months.max())]}
for name, sig in [("overnight_1m", sig_on1), ("intraday_1m", sig_in1),
                  ("overnight_6m", sig_on6), ("intraday_6m", sig_in6),
                  ("combo_on6_minus_in1", combo)]:
    result[name] = decile_fwd(sig)
    r = result[name]
    print(f"{name}: 价差年化{r['spread_ann']:+.1f}% t={r['spread_t']} 单调ρ={r['monotonic_rho']} "
          f"train{r['spread_train']:+.1f}/test{r['spread_test']:+.1f}", file=sys.stderr)

# 持续/反转: 分量自相关(上月分量 -> 次月同分量), 池化去月均
def autocorr(comp):
    cur, nxt = comp, comp.shift(-1)
    vals = []
    for m in months[:-1]:
        a = cur.loc[m][uni.loc[m]]; b = nxt.loc[m]
        df = pd.concat([a, b], axis=1).dropna(); df.columns = ["a", "b"]
        if len(df) < 300: continue
        df = df - df.mean()  # 去截面均值
        vals.append((df["a"] * df["b"]).sum() / (df["a"] ** 2).sum())  # 斜率
    return round(float(np.nanmean(vals)), 3), round(float(np.nanmean(vals) / (np.nanstd(vals) / np.sqrt(len(vals)))), 1)
result["persist_overnight"] = autocorr(ON)
result["persist_intraday"] = autocorr(IN)
print(f"隔夜自相关斜率 {result['persist_overnight']} | 日内自相关斜率 {result['persist_intraday']}", file=sys.stderr)

# ============ size 中性化: 5 市值组内各自分组, 看信号是否只是小盘效应 ============
mo = pd.read_parquet("/root/cb-allotment/data/factor_lab/monthly.parquet")
MV = mo["mv"].unstack("ts_code")
MV.index = pd.PeriodIndex(MV.index, freq="M")
MV = MV.reindex(index=months)
fwd_tot = TOT.shift(-1)

def size_neutral(signal, n_size=5, n_dec=5):
    """每月先分 n_size 市值组, 组内按 signal 分 n_dec 档; 返回组内 spread(高档-低档)均值 + 各市值组 spread。"""
    by_size = {q: [] for q in range(n_size)}   # 每市值组: 组内 高-低 signal 档 forward 价差
    overall = []
    for m in months[:-1]:
        u = uni.loc[m] & MV.loc[m].notna() & signal.loc[m].notna()
        s = signal.loc[m][u]; mv = MV.loc[m][u]; ft = fwd_tot.loc[m]
        if len(s) < 500: continue
        sz = pd.qcut(mv.rank(method="first"), n_size, labels=False)
        mspreads = []
        for q in range(n_size):
            sub = s[sz == q]
            if len(sub) < 40: continue
            dec = pd.qcut(sub.rank(method="first"), n_dec, labels=False)
            hi = sub.index[dec == n_dec - 1]; lo = sub.index[dec == 0]
            sp = ft[hi].mean() - ft[lo].mean()
            by_size[q].append(sp); mspreads.append(sp)
        if mspreads: overall.append(np.mean(mspreads))
    def ann(v): return float(np.nanmean(v) * 12 * 100)
    ov = pd.Series(overall)
    return {"sn_spread_ann": round(ann(ov), 1),
            "sn_spread_t": round(float(ov.mean() / ov.std() * np.sqrt(len(ov))), 1),
            "by_size_quintile": {f"Q{q+1}": round(ann(by_size[q]), 1) for q in range(n_size)}}

result["size_neutral_combo"] = size_neutral(combo)
result["size_neutral_on6"] = size_neutral(sig_on6)
result["size_neutral_in1"] = size_neutral(sig_in1)
print(f"size中性 组合 {result['size_neutral_combo']['sn_spread_ann']}% t={result['size_neutral_combo']['sn_spread_t']} "
      f"各市值组 {result['size_neutral_combo']['by_size_quintile']}", file=sys.stderr)
print(f"size中性 隔夜6m {result['size_neutral_on6']['sn_spread_ann']}% | 日内1m {result['size_neutral_in1']['sn_spread_ann']}%", file=sys.stderr)

json.dump(result, open("/root/cb-allotment/my-app/public/reports/overnight-intraday/backtest_result.json", "w"), ensure_ascii=False, indent=1)
print("saved research/tug_of_war_result.json", file=sys.stderr)