# -*- coding: utf-8 -*-
"""因子动量 A股版 (Gupta-Kelly 2019 "Factor Momentum Everywhere" / Ehsani-Linnainmaa 2022 JF)。

问: 因子本身能追涨吗——上月/近12月赢的因子, 下月继续赢?
用自家 factor_lab 32 个因子, 每因子提月度序列两条(与 decile_test 完全同口径):
  spread = D10−D1(因子多空价差, 研究口径) ; topex = D10 − 全市场(长仓超额, 落地口径)
测: ①因子价差的1月自相关(池化) ②TS因子动量策略: 持有过去K月为正的因子(等权), K=1/3/12,
    vs 静态等权全因子 ③长仓版: topex 同样测。train<2022/test。
直接回答: 果仁的 IC 择时开关(用近期因子IC选因子)该不该用、什么窗口。
"""
import sys, json, warnings
warnings.filterwarnings("ignore")
sys.path.insert(0, "/root/cb-allotment/scripts")
import numpy as np, pandas as pd
from lib.factor_lab import Ctx, TRAIN_END
from lib.factors import REGISTRY
from lib.factors2 import REGISTRY2

ctx = Ctx.get()
ALL = {**REGISTRY, **REGISTRY2}
fwd = ctx.mo.shift(-1)

def spread_series(sig, start="2017-05-31", end=None, annual_may=False, winsor=0.02, min_universe=300, **_):
    months = ctx.months(start, end)
    sp, tex = {}, {}
    held = None
    for t in months:
        if annual_may:
            if t.month == 5 and t in sig.index: held = sig.loc[t]
            if held is None: continue
            s = held.dropna()
        else:
            if t not in sig.index: continue
            s = sig.loc[t].dropna()
        s = s.replace([np.inf, -np.inf], np.nan).dropna()
        if winsor: s = s.clip(s.quantile(winsor), s.quantile(1 - winsor))
        u = [c for c in s.index if c in fwd.columns and pd.notna(fwd.loc[t, c])]
        s = s[u]
        if len(s) < min_universe: continue
        qb = pd.qcut(s.rank(method="first"), 10, labels=range(1, 11))
        row = fwd.loc[t]
        g10 = float(row[list(s.index[qb == 10])].mean()); g1 = float(row[list(s.index[qb == 1])].mean())
        sp[t] = g10 - g1
        tex[t] = g10 - float(row[u].mean())
    return pd.Series(sp), pd.Series(tex)

SP, TX = {}, {}
for slug, (fn, kw) in ALL.items():
    try:
        sig = fn(ctx)
        kw2 = {k: v for k, v in kw.items() if k in ("start", "end", "annual_may", "winsor", "min_universe")}
        kw2.setdefault("start", "2017-05-31")
        sp, tx = spread_series(sig, **kw2)
        if len(sp) >= 36:
            SP[slug] = sp; TX[slug] = tx
            print(f"  {slug}: {len(sp)}月 spread年化{sp.mean()*12*100:+.1f}%", file=sys.stderr)
        else:
            print(f"  {slug}: 序列太短({len(sp)}) 跳过", file=sys.stderr)
    except Exception as e:
        print(f"  {slug} FAIL {str(e)[:60]}", file=sys.stderr)
S = pd.DataFrame(SP).sort_index()   # 月 x 因子 (D10-D1)
T = pd.DataFrame(TX).sort_index()   # 月 x 因子 (D10-市场)
print(f"因子面板 {S.shape}", file=sys.stderr)

TRAIN = pd.Timestamp(TRAIN_END)
def ann(x): return float(np.nanmean(x) * 12 * 100)
def sharpe(x):
    x = pd.Series(x).dropna()
    return float(x.mean() / x.std() * np.sqrt(12)) if x.std() > 0 else np.nan

# ① 池化自相关: 因子上月收益 -> 本月(截面去均值)
def pooled_ac(P):
    a = P.shift(1).values.ravel(); b = P.values.ravel()
    m = ~(np.isnan(a) | np.isnan(b))
    a, b = a[m], b[m]
    a = a - a.mean(); b = b - b.mean()
    slope = (a * b).sum() / (a * a).sum()
    r = np.corrcoef(a, b)[0, 1]
    t = r * np.sqrt(len(a) - 2) / np.sqrt(1 - r * r)
    return round(slope, 3), round(r, 3), round(t, 1), int(len(a))

# ② TS 因子动量: 过去K月累计>0 的因子等权持有(spread口径: 正持有/负翻转两版; topex口径: 只持正)
def ts_fm(P, K, flip=False):
    rets = []
    for i in range(K, len(P) - 0):
        past = P.iloc[i - K:i].sum()
        cur = P.iloc[i]
        pos = past[past > 0].index.intersection(cur.dropna().index)
        neg = past[past <= 0].index.intersection(cur.dropna().index)
        if flip:
            vals = list(cur[pos].values) + list(-cur[neg].values)
        else:
            vals = list(cur[pos].values)
        rets.append((P.index[i], np.mean(vals) if vals else 0.0))
    return pd.Series(dict(rets))

def strat_stats(s):
    tr = s[s.index < TRAIN]; te = s[s.index >= TRAIN]
    return {"ann": round(ann(s), 1), "sharpe": round(sharpe(s), 2),
            "train_ann": round(ann(tr), 1), "test_ann": round(ann(te), 1),
            "train_sh": round(sharpe(tr), 2), "test_sh": round(sharpe(te), 2), "n": int(len(s))}

res = {"n_factors": int(S.shape[1]), "months": [str(S.index.min().date()), str(S.index.max().date())]}
res["pooled_autocorr_spread"] = dict(zip(["slope", "corr", "t", "n"], pooled_ac(S)))
res["pooled_autocorr_topex"] = dict(zip(["slope", "corr", "t", "n"], pooled_ac(T)))

res["static_ew_spread"] = strat_stats(S.mean(axis=1))
res["static_ew_topex"] = strat_stats(T.mean(axis=1))
for K in (1, 3, 12):
    res[f"fm{K}_spread_long"] = strat_stats(ts_fm(S, K))
    res[f"fm{K}_spread_flip"] = strat_stats(ts_fm(S, K, flip=True))
    res[f"fm{K}_topex_long"] = strat_stats(ts_fm(T, K))
# 因子平均持有个数(K=12 spread)
past12 = S.rolling(12).sum()
res["avg_n_positive_12m"] = round(float((past12 > 0).sum(axis=1).mean()), 1)

json.dump(res, open("/root/cb-allotment/my-app/public/reports/factor-momentum/backtest_result.json", "w"), ensure_ascii=False, indent=1)
print("池化自相关 spread:", res["pooled_autocorr_spread"], file=sys.stderr)
print("static EW spread:", res["static_ew_spread"], file=sys.stderr)
for K in (1, 3, 12):
    print(f"FM{K} spread_long:", res[f"fm{K}_spread_long"], file=sys.stderr)
    print(f"FM{K} topex_long:", res[f"fm{K}_topex_long"], file=sys.stderr)
print("saved", file=sys.stderr)