import pandas as pd, numpy as np, sqlite3, json
con=sqlite3.connect('/root/cb-allotment/data/fundamental_cache.db')
sh=pd.read_sql("select * from etf_share",con).pivot(index='trade_date',columns='ts_code',values='fd_share').sort_index()
px=pd.read_sql("select * from etf_close",con).pivot(index='trade_date',columns='ts_code',values='close').sort_index()
broad=['510050.SH','510300.SH','510310.SH','510500.SH','512050.SH','512100.SH','588000.SH','159915.SH','159915.SZ']
sector=[c for c in px.columns if c not in broad]
px=px.loc['20190102':]; sh=sh.reindex(px.index).ffill(limit=3)
r=px.pct_change()
# guard splits: daily price drop >35% with share jump -> treat return as 0
spl=(r<-0.35)
r=r.mask(spl,0)
f20=sh/sh.shift(20)-1
f20=f20.mask((f20.abs()>3))
pctl=f20.rolling(250,min_periods=120).apply(lambda a:(a[:-1]<a[-1]).mean() if np.isfinite(a[-1]) else np.nan,raw=True)
dt=pd.to_datetime(px.index)
wk=[px.index[i] for i in range(len(px)-1) if dt[i].isocalendar()[1]!=dt[i+1].isocalendar()[1]]
C=0.0005
def run(pick):
    rows=[];prev=set()
    for a,b in zip(wk[:-1],wk[1:]):
        ia=px.index.get_loc(a); ib=px.index.get_loc(b)
        e=px.index[ia+1]; x=px.index[min(ib+1,len(px)-1)]
        sel=pick(a)
        seg=(1+r.loc[e:x].iloc[1:]).prod()-1
        ret=seg[list(sel)].mean() if len(sel) else 0.0
        cur=set(sel); turn=len(cur^prev)/max(len(cur|prev),1) if (cur or prev) else 0; prev=cur
        ew=seg[sector].mean()
        rows.append(dict(d=b,r=ret-2*C*turn,ew=ew,n=len(sel)))
    return pd.DataFrame(rows).set_index('d')
def st(x):
    x=x.dropna();nav=(1+x).cumprod();y=len(x)/52
    return dict(cagr=round((nav.iloc[-1]**(1/y)-1)*100,1),mdd=round((nav/nav.cummax()-1).min()*100,1),sharpe=round(x.mean()/x.std()*np.sqrt(52),2))
rng=np.random.default_rng(1)
picks={
 'extreme_out_p10':lambda a:[c for c in sector if pctl.loc[a,c]<0.10],
 'extreme_in_p90':lambda a:[c for c in sector if pctl.loc[a,c]>0.90],
 'low3_flow':lambda a:list(f20.loc[a,sector].dropna().nsmallest(3).index),
 'high3_flow':lambda a:list(f20.loc[a,sector].dropna().nlargest(3).index),
 'placebo3':lambda a:list(rng.choice(f20.loc[a,sector].dropna().index,3,replace=False)) if f20.loc[a,sector].notna().sum()>=3 else [],
}
res={};nav={}
for k,p in picks.items():
    df=run(p);nav[k]=(1+df.r).cumprod()
    res[k]={nm:dict(strat=st(df.r.loc[sl]),ew=st(df.ew.loc[sl]),excess=st((df.r-df.ew).loc[sl]),avg_n=round(df.n.loc[sl].mean(),1)) for nm,sl in [('all',slice(None)),('IS',slice(None,'20221231')),('OOS',slice('20230101',None))]}
nav['ew_sector']=(1+df.ew).cumprod()
res['sectors']=sector
json.dump(res,open('/root/papers3_etf.json','w'),indent=1,ensure_ascii=False)
pd.DataFrame(nav).to_csv('/root/papers3_etfnav.csv')
print(json.dumps({k:v['OOS'] for k,v in res.items() if k!='sectors'},ensure_ascii=False))
