import pandas as pd, numpy as np, sqlite3, json, sys
DB='/root/cb-allotment/data/fundamental_cache.db'
con=sqlite3.connect(DB)
d=pd.read_parquet('/root/ml_micro_export/daily.parquet')
d=d[d.ts_code.str[-2:].isin(['SH','SZ'])]
P=lambda c: d.pivot(index='trade_date',columns='ts_code',values=c).sort_index()
close=P('close'); opn=P('open'); pre=P('pre_close'); pct=P('pct_chg')/100
del d
dates=close.index
b=pd.read_sql("select trade_date,ts_code,circ_mv,turnover_rate from daily_basic where trade_date>='20150101' and trade_date<'20190101'",con)
b=pd.concat([b,pd.read_parquet('/root/ml_micro_export/basic.parquet',columns=['trade_date','ts_code','circ_mv','turnover_rate'])]).drop_duplicates(['trade_date','ts_code'])
mv=b.pivot(index='trade_date',columns='ts_code',values='circ_mv');mv=mv.reindex(mv.index.union(dates)).ffill(limit=7).reindex(index=dates,columns=close.columns).where(close.notna())
va=pd.read_parquet('/root/ml_micro_export/va.parquet');amt=va.pivot(index='trade_date',columns='ts_code',values='amount').reindex(index=dates,columns=close.columns);del va
to=amt/mv
del b
print('basic range',mv.dropna(how='all').index.min(),file=sys.stderr)
adj=(1+pct.fillna(0)).cumprod()
adjc=adj.where(close.notna())
adjo=(opn/close*adj)
# ST
nc=pd.read_sql("select ts_code,name,start_date from namechange",con)
sb=pd.read_sql("select ts_code,list_date from stock_basic",con).set_index('ts_code').list_date
nc=nc.dropna(subset=['start_date']).drop_duplicates(['ts_code','start_date'],keep='last')
nc['isst']=nc.name.str.contains('ST').astype(float)
stw=nc.pivot(index='start_date',columns='ts_code',values='isst')
st=stw.reindex(stw.index.union(dates)).ffill().reindex(index=dates,columns=close.columns).fillna(0).astype(bool)
listd=pd.Series(sb).reindex(close.columns)
# month ends
dt=pd.to_datetime(dates)
me=[dates[i] for i in range(len(dates)-1) if dt[i].month!=dt[i+1].month]
me=[x for x in me if x>='20151231']
# signals
r=pct
mkt=r.where(mv.notna()).mean(axis=1)
ret20=r.rolling(20,min_periods=15)
# ivol: residual std vs market over 20d
cov=r.rolling(20,min_periods=15).cov(mkt) if False else None
mr=mkt
mvar=mr.rolling(20).var()
beta=(r.mul(mr,axis=0).rolling(20,min_periods=15).mean()-r.rolling(20,min_periods=15).mean().mul(mr.rolling(20).mean(),axis=0)).div(mvar,axis=0)
resid=r.sub(beta.mul(mr,axis=0))
ivol=resid.rolling(20,min_periods=15).std()
tov=to.rolling(20,min_periods=15).mean()
ma21=adjc.rolling(21,min_periods=18).mean(); ma200=adjc.rolling(200,min_periods=180).mean()
mad=ma21/ma200
mom=adjc.shift(21)/adjc.shift(252)-1
rev=adjc/adjc.shift(21)-1
def univ(t):
    i=dates.get_loc(t)
    ok=close.loc[t].notna()&mv.loc[t].notna()&(~st.loc[t])&(close.loc[t]>=2)
    ld=listd.fillna('99999999')
    ok&=(ld<=dates[max(i-250,0)])
    ok&=~close.columns.str.startswith(('8','4','9'))
    return ok
def rk(s): return s.rank(pct=True)
COST=0.002
def bt(sigfun, q=0.1, side='top'):
    rows=[]; prev=set()
    for k in range(len(me)-1):
        t,t2=me[k],me[k+1]
        i=dates.get_loc(t); j=dates.get_loc(t2)
        if j+1>=len(dates): break
        e=dates[i+1]; x=dates[j+1]
        u=univ(t); s=sigfun(t)[u].dropna()
        if len(s)<300:
            continue
        # exclude limit-up at entry open
        lim=np.where(s.index.str.startswith(('30','68')),0.195,0.095)
        gap=(opn.loc[e,s.index]/pre.loc[e,s.index]-1)
        fr=(adjo.loc[x,s.index]/adjo.loc[e,s.index]-1)
        # delisted/suspended -> use last available adj close
        lastc=adjc.loc[e:x,s.index].ffill().iloc[-1]/adjo.loc[e,s.index]-1
        fr=fr.fillna(lastc)
        n=max(int(len(s)*q),20)
        srt=s.sort_values()
        top=srt.index[-n:]; bot=srt.index[:n]
        okb=gap[top]<lim[[s.index.get_loc(c) for c in top]]
        topb=top[okb.values]
        ew=fr.mean()
        lt=fr[topb].mean(); lb=fr[bot].mean()
        cur=set(topb); turn=1-len(cur&prev)/max(len(cur),1); prev=cur
        rows.append(dict(d=t2,top=lt,bot=lb,ew=ew,turn=turn,net=lt-2*COST*turn))
    return pd.DataFrame(rows).set_index('d')
def stats(x,per=12):
    x=x.dropna(); nav=(1+x).cumprod()
    yrs=len(x)/per; cagr=nav.iloc[-1]**(1/yrs)-1
    dd=(nav/nav.cummax()-1).min(); sh=x.mean()/x.std()*np.sqrt(per)
    return dict(cagr=round(float(cagr)*100,1),mdd=round(float(dd)*100,1),sharpe=round(float(sh),2),n=len(x))
def summarize(df):
    out={}
    for nm,sl in [('all',slice(None)),('IS',slice(None,'20201231')),('OOS',slice('20210101',None))]:
        z=df.loc[sl]
        out[nm]=dict(long_net=stats(z.net),ew=stats(z.ew),excess=stats(z.net-z.ew),ls=stats(z.top-z.bot),turn=round(float(z.turn.mean()),2))
    return out
t=me[60];u=univ(t);print('dbg',t,u.sum(),close.loc[t].notna().sum(),mv.loc[t].notna().sum(),st.loc[t].sum(),ivol.loc[t].notna().sum(),mad.loc[t].notna().sum(),file=sys.stderr)
rng=np.random.default_rng(0)
sigs={
 'rip_low': lambda t: -(rk(ivol.loc[t])+rk(tov.loc[t])+rk(-np.log(mv.loc[t])))/3,   # long low speculative
 'rip_low_nosize': lambda t: -(rk(ivol.loc[t])+rk(tov.loc[t]))/2,
 'rip_low_bigsize': lambda t: -(rk(ivol.loc[t])+rk(tov.loc[t])+rk(np.log(mv.loc[t])))/3,
 'ivol_low': lambda t: -ivol.loc[t],
 'mad': lambda t: mad.loc[t],
 'mad_low': lambda t: -mad.loc[t],
 'mom12_1': lambda t: mom.loc[t],
 'placebo': lambda t: pd.Series(rng.random(len(close.columns)),index=close.columns),
}
res={}; series={}
for k,f in sigs.items():
    df=bt(f); res[k]=summarize(df); series[k]=df
    print(k,json.dumps(res[k]['OOS']),file=sys.stderr)
# MAD within smallcap half & MAD residual on momentum corr
cs=[]
for t in me[12:-1]:
    u=univ(t); a=mad.loc[t][u]; m=mom.loc[t][u]; cs.append(a.rank().corr(m.rank()))
res['mad_mom_rankcorr']=round(float(np.nanmean(cs)),2)
# MAD market timing: ew index of universe, hold when mkt MAD>1 at month end else cash
mi=(1+mkt.fillna(0)).cumprod(); mmad=mi.rolling(21).mean()/mi.rolling(200).mean()
dr=mkt.copy(); pos=(mmad>1).shift(2).fillna(False).astype(float)  # signal at t close, in from t+1 close (conservative)
tr=pos.diff().abs().fillna(0)
tim=pos*dr-tr*0.0005
for nm,sl in [('all',slice('20160101',None)),('IS',slice('20160101','20201231')),('OOS',slice('20210101',None))]:
    res.setdefault('mad_timing',{})[nm]=dict(timed=stats(tim.loc[sl],252),bh=stats(dr.loc[sl],252),inmkt=round(float(pos.loc[sl].mean()),2))
json.dump(res,open('/root/papers3_stk.json','w'),indent=1,ensure_ascii=False)
nav=pd.DataFrame({k:(1+v.net).cumprod() for k,v in series.items()}); nav['ew']=(1+series['mad'].ew).cumprod()
nav.to_csv('/root/papers3_nav.csv')
print('done')