# -*- coding: utf-8 -*-
"""聚焦基本面超预期: 行业等权top-K动量 × 行业内超预期选股。
用 daily.pct_chg(已复权口径, 对齐部署基线) 修正持仓收益。
超预期两个用法: (幅度)按surprise z排序取topN; (符号)filter surprise>0等权。
月频0.2%/边, 剔ST/剔北交所, train<2022/test≥2022。"""
import sys, json, datetime, sqlite3
import numpy as np, pandas as pd
from collections import defaultdict
DB="/root/cb-allotment/data/fundamental_cache.db"; c=sqlite3.connect(DB); c.execute("PRAGMA busy_timeout=120000")
sb=pd.read_sql("SELECT ts_code,name FROM stock_basic",c).set_index('ts_code')
sw=pd.read_sql("SELECT l1_name,ts_code FROM sw_member_full",c).drop_duplicates('ts_code'); io=dict(zip(sw.ts_code,sw.l1_name))
# 用 pct_chg 修正收益
px=pd.read_sql("SELECT trade_date,ts_code,pct_chg FROM daily WHERE trade_date>='20141001'",c); px=px[~px.ts_code.str[0].isin(['4','8','9'])]
ret=px.pivot(index='trade_date',columns='ts_code',values='pct_chg').sort_index()/100.0
dates=list(ret.index); codes=list(ret.columns)
st_mask=pd.Series(sb['name'].reindex(codes).fillna('').str.contains('ST').values,index=codes).values
indarr=np.array([io.get(x,'') for x in codes]); inds=sorted(set(io.values()))
sw2=sw[['ts_code','l1_name']].copy(); sw2['is_st']=sw2['ts_code'].map(sb['name']).fillna('').str.contains('ST').astype(int); sw2.to_sql('tmp_map',c,if_exists='replace',index=False)
_on=pd.read_sql("SELECT d.trade_date td,mp.l1_name ind,AVG(d.open*1.0/d.pre_close-1) onr FROM daily d JOIN tmp_map mp ON d.ts_code=mp.ts_code WHERE d.trade_date>='20141001' AND d.pre_close>0 AND mp.is_st=0 AND ABS(d.open*1.0/d.pre_close-1)<0.5 GROUP BY d.trade_date,mp.l1_name",c)
MOM=((1+_on.pivot(index='td',columns='ind',values='onr').reindex(index=dates,columns=inds).fillna(0)).cumprod()); MOM=MOM/MOM.shift(180)-1
# 超预期
rep=pd.read_sql("SELECT ts_code,report_date,quarter,np FROM report_rc WHERE np IS NOT NULL AND quarter LIKE '%Q4'",c); rep['fy']=rep['quarter'].str[:4]; rep_g={k:v[['report_date','np']].values for k,v in rep.groupby(['ts_code','fy'])}
fc=pd.read_sql("SELECT ts_code,ann_date,end_date,net_profit_min,net_profit_max FROM forecast WHERE end_date LIKE '%1231'",c); fc['val']=fc[['net_profit_min','net_profit_max']].mean(axis=1)
ex=pd.read_sql("SELECT ts_code,ann_date,end_date,n_income/10000.0 AS val FROM express WHERE end_date LIKE '%1231' AND n_income IS NOT NULL",c)
inc=pd.read_sql("SELECT ts_code,ann_date,end_date,n_income_attr_p/10000.0 AS val FROM income_period WHERE end_date LIKE '%1231' AND n_income_attr_p IS NOT NULL",c)
candf=pd.concat([fc[['ts_code','ann_date','end_date','val']],ex,inc]).dropna(subset=['ann_date','val']).sort_values('ann_date').drop_duplicates(['ts_code','end_date'],keep='first')
c.close()
SUR=defaultdict(list)
for r in candf.itertuples():
    rows=rep_g.get((r.ts_code,r.end_date[:4]))
    if rows is None: continue
    ann=datetime.datetime.strptime(r.ann_date,"%Y%m%d"); lb=(ann-datetime.timedelta(days=270)).strftime("%Y%m%d")
    nps=[float(v) for rd,v in rows if lb<=rd<r.ann_date]
    if len(nps)<3: continue
    cons=float(np.median(nps))
    if abs(cons)<100: continue
    SUR[r.ts_code].append((r.ann_date,(r.val-cons)/abs(cons)))
for k in SUR: SUR[k].sort()
def latest(code,d,win=450):
    ev=SUR.get(code)
    if not ev: return None
    lb=(datetime.datetime.strptime(d,"%Y%m%d")-datetime.timedelta(days=win)).strftime("%Y%m%d"); best=None
    for ann,v in ev:
        if lb<=ann<=d: best=v
    return best
print("就绪",flush=True)
retv=np.where(st_mask,np.nan,ret.values); posd={x:i for i,x in enumerate(dates)}
memidx={ind:np.where(indarr==ind)[0] for ind in inds}
months=pd.Series(dates).groupby(pd.Series(dates).str[:6]).first().values; ms=[m for m in months if m>='20170101']
def run(K,mode,perN=8):
    # mode: 'all' 基线全成分 | 'sign' 超预期>0过滤 | 'mag' 按幅度排序取perN | 'magpool' 全池按幅度取总topN(=perN*K)
    nav=1.0; ser=[]; cur=None; rset=set(ms)
    for di in range(posd[ms[0]],len(dates)):
        d=dates[di]; r=0.0
        if cur is not None:
            tot=0.0
            for w,idx in cur:
                v=retv[di,idx]; v=v[~np.isnan(v)]
                if len(v): tot+=w*float(v.mean())
            r=tot
        if d in rset:
            sc=MOM.iloc[di-1].dropna()
            if len(sc)>=10:
                topinds=list(sc.sort_values(ascending=False).head(K).index)
                if mode=='magpool':
                    pool=[]
                    for ind in topinds:
                        for j in memidx.get(ind,[]):
                            if st_mask[j] or np.isnan(retv[di-1,j]): continue
                            a=latest(codes[j],d)
                            if a is not None: pool.append((j,a))
                    pool.sort(key=lambda x:-x[1]); sel=[j for j,_ in pool[:perN*K]]
                    cur=[(1.0/max(len(sel),1),np.array(sel,dtype=int))] if sel else cur
                    if sel: r-=0.004
                else:
                    legs=[]
                    for ind in topinds:
                        mem=[j for j in memidx.get(ind,[]) if not st_mask[j] and not np.isnan(retv[di-1,j])]
                        if mode=='all': sel=mem
                        elif mode=='sign': sel=[j for j in mem if (latest(codes[j],d) or -9)>0]
                        elif mode=='mag':
                            sc2=[(j,latest(codes[j],d)) for j in mem]; sc2=[(j,a) for j,a in sc2 if a is not None]
                            sc2.sort(key=lambda x:-x[1]); sel=[j for j,_ in sc2[:perN]] if len(sc2)>=perN else mem
                        if sel: legs.append(np.array(sel,dtype=int))
                    if legs:
                        w=1.0/len(legs); cur=[(w,leg) for leg in legs]; r-=0.004
        nav*=1+r; ser.append((d,nav))
    return ser
def stats(ser,lo=None,hi=None):
    s=[(d,v) for d,v in ser if (lo is None or d>=lo) and (hi is None or d<hi)]
    if len(s)<50: return None
    nv=np.array([v for _,v in s]); nv=nv/nv[0]; yrs=len(nv)/244; rr=np.diff(nv)/nv[:-1]
    return {"cagr":round((nv[-1]**(1/yrs)-1)*100,1),"sharpe":round(float(np.mean(rr)/np.std(rr)*np.sqrt(244)),2),"mdd":round(float((nv/np.maximum.accumulate(nv)-1).min())*100,1)}
out={}
for K in (5,3,2):
    for mode,lbl in (('all','全成分基线'),('sign','超预期>0'),('mag','幅度topN/行业')):
        ser=run(K,mode); key=f"K{K}_{mode}"
        out[key]={"label":f"top{K}·{lbl}","full":stats(ser),"train":stats(ser,hi='20220101'),"test":stats(ser,lo='20220101')}
        print(f"{key} {lbl}: full {out[key]['full']} test {out[key]['test']}",flush=True)
# 浓缩版: 全池按幅度取总top-N
for K in (5,3):
    for N in (15,10):
        ser=run(K,'magpool',perN=N//K if K else 3); key=f"K{K}_magpool{N}"
        out[key]={"label":f"top{K}·超预期幅度全池top{N}","full":stats(ser),"train":stats(ser,hi='20220101'),"test":stats(ser,lo='20220101')}
        print(f"{key}: full {out[key]['full']} test {out[key]['test']}",flush=True)
json.dump(out,open("research/industry_surprise_focus.json","w"),ensure_ascii=False,indent=1)
print("DONE",flush=True)

# ========== 追加: 给 K5 超预期>0 赢家叠加匹配择时(宽基MA, 非微盘) ==========
import sqlite3 as _s3
_c=_s3.connect(DB)
def _idx(code):
    try:
        s=pd.read_sql(f"SELECT trade_date,close FROM index_daily WHERE ts_code='{code}'",_c).set_index('trade_date')['close']
        return s.reindex(dates).ffill()
    except: return None
_c.close()
# 候选宽基: 中证全指000985 / 沪深300 000300 / 中证500 000905
def run_timed(K,mode,idxcode,ma,confirm=1):
    sig=_idx(idxcode)
    if sig is None: return None
    ma_s=sig.rolling(ma).mean(); on_raw=(sig>ma_s).values
    # 连续confirm日确认
    on=on_raw.copy()
    if confirm>1:
        for i in range(len(on)):
            on[i]=all(on_raw[max(0,i-confirm+1):i+1])
    nav=1.0; ser=[]; cur=None; rset=set(ms)
    for di in range(posd[ms[0]],len(dates)):
        d=dates[di]; r=0.0
        risk_on=on[di-1] if di>0 else True
        if cur is not None and risk_on:
            tot=0.0
            for w,idx in cur:
                v=retv[di,idx]; v=v[~np.isnan(v)]
                if len(v): tot+=w*float(v.mean())
            r=tot
        if d in rset:
            sc=MOM.iloc[di-1].dropna()
            if len(sc)>=10:
                topinds=list(sc.sort_values(ascending=False).head(K).index); legs=[]
                for ind in topinds:
                    mem=[j for j in memidx.get(ind,[]) if not st_mask[j] and not np.isnan(retv[di-1,j])]
                    sel=[j for j in mem if (latest(codes[j],d) or -9)>0] if mode=='sign' else mem
                    if sel: legs.append(np.array(sel,dtype=int))
                if legs:
                    newcur=[(1.0/len(legs),leg) for leg in legs]
                    if risk_on: r-=0.004
                    cur=newcur
        nav*=1+r; ser.append((d,nav))
    return ser
print("=== 择时叠加 ===",flush=True)
tout={}
for idxc,nm in (('000985.CSI','中证全指'),('000300.SH','沪深300'),('000905.SH','中证500')):
    for ma in (20,60):
        ser=run_timed(5,'sign',idxc,ma,confirm=2)
        if ser is None: print(f"{nm} MA{ma}: 无指数",flush=True); continue
        key=f"timed_{idxc}_{ma}"
        tout[key]={"label":f"K5超预期>0 + {nm}MA{ma}择时","full":stats(ser),"train":stats(ser,hi='20220101'),"test":stats(ser,lo='20220101')}
        print(f"{nm} MA{ma}: full {tout[key]['full']} test {tout[key]['test']}",flush=True)
json.dump(tout,open("research/industry_surprise_timed.json","w"),ensure_ascii=False,indent=1)
print("TIMED_DONE",flush=True)