# Primary tests, period breakdowns, regime tables, regime-filter & meta-switch OOS tests (PREREG), BH, charts.
import pandas as pd,numpy as np,json,os
from scipy import stats
import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
L='/workspace/research/popular-strategies'; RS=f'{L}/results'; CHD=f'{L}/charts'
NAMES={'P1':'ORB15 breakout','P2':'ORB5 first 5-min candle','P3':'Prior-day high/low breakout','P4':'Narrow-CPR breakout + VWAP','P5':'Supertrend(10,3) flip + VWAP','P6':'EMA 9/21 cross + VWAP','P7':'RSI(14) 30/70 reversal','P8':'Gap fade (gap-fill)'}
C=list(NAMES)
T=pd.read_csv(f'{RS}/trades_options.csv',parse_dates=['date']); T=T[T.status=='ok']
SP=pd.read_csv(f'{RS}/trades_spot.csv',parse_dates=['date']).dropna(subset=['spot_pts'])
RR=pd.read_csv(f'{RS}/trades_random.csv',parse_dates=['date','src_date'])
LAB=pd.read_csv(f'{RS}/regime_labels.csv',parse_dates=['date']).set_index('date')
LBS=[c for c in LAB.columns if c[:2] in ('L1','L2','L3','L4','L5','L6','L7','L8')]
O=T[T.ex=='T50'].copy(); O=O.join(LAB[LBS],on='date'); SP=SP.join(LAB[LBS],on='date'); RR=RR.join(LAB[LBS],on='date')
def lotq(d): return 65 if d>=pd.Timestamp('2026-01-01') else 75
def tp(p): return stats.ttest_1samp(p,0,alternative='greater').pvalue if len(p)>2 else np.nan
def bh(p):
    p=np.asarray(p,float); p=np.where(np.isnan(p),1,p); m=len(p); o=np.argsort(p); q=np.empty(m); q[o]=np.minimum.accumulate((p[o]*m/np.arange(1,m+1))[::-1])[::-1]; return np.minimum(q,1)
rng=np.random.default_rng(7)
def summ(p,dates=None):
    p=np.asarray(p,float); w=p[p>0]; l=p[p<=0]
    out=dict(n=len(p),mean=p.mean() if len(p) else np.nan,median=np.median(p) if len(p) else np.nan,win=(p>0).mean()*100 if len(p) else np.nan,
             payoff=(w.mean()/abs(l.mean())) if len(w) and len(l) and l.mean()!=0 else np.nan,t_p=tp(p))
    if dates is not None and len(p):
        rs=p*np.array([lotq(d) for d in dates]); cum=np.cumsum(rs); out.update(rs_total=rs.sum(),maxdd_rs=(cum-np.maximum.accumulate(np.r_[0,cum])[1:]).min())
    return out
def rand_pct(sub,cand):
    R=RR[RR.cand==cand]; G={k:g.pnl.values for k,g in R.groupby('src_date')}; arrs=[G[d] for d in sub.date if d in G]
    if not arrs: return np.nan,np.nan
    means=np.array([np.mean([a[rng.integers(len(a))] for a in arrs]) for _ in range(2000)]); return (means<sub.pnl.mean()).mean()*100,np.mean(np.concatenate(arrs))
# ---------- 1. primary ----------
rows=[]
for c in C:
    for sp in ('train','test'):
        x=O[(O.cand==c)&(O.split==sp)].sort_values('date'); pc,rm=rand_pct(x,c)
        rows.append(dict(cand=c,name=NAMES[c],split=sp,**summ(x.pnl.values,list(x.date)),random_mean=rm,random_pctile=pc))
P=pd.DataFrame(rows); P.loc[P.split=='test','q_bh']=bh(P[P.split=='test'].t_p.values); P.round(3).to_csv(f'{RS}/primary_results.csv',index=False)
E1=T[T.ex=='E1'].groupby(['cand','split']).pnl.agg(n='size',mean='mean',win=lambda v:(v>0).mean()*100).round(2).reset_index(); E1.to_csv(f'{RS}/secondary_E1.csv',index=False)
spy=SP.assign(year=SP.date.dt.year).groupby(['cand','year']).agg(n=('spot_net','size'),spot_net=('spot_net','mean'),spot_E1=('spot_pts_E1','mean')).round(2).reset_index()
spy.to_csv(f'{RS}/spot_by_year.csv',index=False)
# ---------- 2. period breakdowns ----------
def per(df,col,val):
    d=df.copy(); d['year']=d.date.dt.year.astype(str); d['half']=d.year+'H'+((d.date.dt.month>6)+1).astype(str); d['quarter']=d.date.dt.to_period('Q').astype(str); d['month']=d.date.dt.to_period('M').astype(str)
    out=[]
    for g in ('year','half','quarter','month'):
        z=d.groupby(['cand',g])[val].agg(n='size',mean='mean',total='sum',win=lambda v:(v>0).mean()*100).reset_index().rename(columns={g:'period'}); z['level']=g; out.append(z)
    return pd.concat(out).round(2)
po=per(O,'cand','pnl'); po.to_csv(f'{RS}/options_by_period.csv',index=False)
ps=per(SP,'cand','spot_net'); ps.to_csv(f'{RS}/spot_by_period.csv',index=False)
# share of positive quarters
pq=pd.concat([ps[ps.level=='quarter'].groupby('cand').apply(lambda z:pd.Series(dict(spot_q_pos=f"{(z['mean']>0).sum()}/{len(z)}"))),
              po[po.level=='quarter'].groupby('cand').apply(lambda z:pd.Series(dict(opt_q_pos=f"{(z['mean']>0).sum()}/{len(z)}")))],axis=1)
pq.to_csv(f'{RS}/positive_quarters.csv')
# ---------- 3. regime bucket tables ----------
rb=[]
for c in C:
    for lb in LBS:
        for dsname,df,val,splitfn in (('options',O,'pnl',lambda d:np.where(d<pd.Timestamp('2026-01-01'),'train','test')),('spot',SP,'spot_net',lambda d:np.where(d<pd.Timestamp('2026-01-01'),'2022-25','2026'))):
            x=df[df.cand==c].copy(); x['sp']=splitfn(x.date)
            for (b,s),g in x.dropna(subset=[lb]).groupby([lb,'sp']): rb.append(dict(cand=c,data=dsname,label=lb,bucket=b,split=s,n=len(g),mean=g[val].mean(),win=(g[val]>0).mean()*100))
RB=pd.DataFrame(rb).round(2); RB.to_csv(f'{RS}/regime_buckets.csv',index=False)
NCUTS=int(LAB[LBS].nunique().sum())
# ---------- 4. regime filter selection ----------
def select(x,val,nb,nf):
    best=None; tot0=x[val].sum()
    for lb in LBS:
        z=x.dropna(subset=[lb]); g=z.groupby(lb)[val].agg(['size','mean'])
        good=list(g[(g['size']>=nb)&(g['mean']>0)].index)
        if not good: continue
        f=z[z[lb].isin(good)]
        if len(f)<nf: continue
        if best is None or f[val].sum()>best['train_total']: best=dict(label=lb,buckets=good,train_n=len(f),train_mean=f[val].mean(),train_total=f[val].sum(),train_always_mean=x[val].mean())
    return best
def subset_pct(allp,fp):
    if len(fp)==0 or len(allp)<=len(fp): return np.nan
    m=np.array([rng.choice(allp,len(fp),replace=False).mean() for _ in range(2000)]); return (m<fp.mean()).mean()*100
FAM=[]; CH={}
for c in C:
    x=O[O.cand==c]; tr=x[x.split=='train']; te=x[x.split=='test']; s=select(tr,'pnl',10,20); CH[c]=s
    if s is None: FAM.append(dict(test='opt_regime',cand=c,rule='no filter found on train',oos_n=0,oos_mean=np.nan,always_mean=te.pnl.mean(),p=1.0)); continue
    f=te[te[s['label']].isin(s['buckets'])]
    # random-entry baseline restricted to same buckets (random day in good buckets), test period
    R=RR[(RR.cand==c)&(RR.split=='test')&(RR[s['label']].isin(s['buckets']))]; G={k:g.pnl.values for k,g in R.groupby('src_date')}; arrs=[G[d] for d in f.date if d in G]
    rpc=(np.array([np.mean([a[rng.integers(len(a))] for a in arrs]) for _ in range(2000)])<f.pnl.mean()).mean()*100 if arrs else np.nan
    FAM.append(dict(test='opt_regime',cand=c,rule=f"{s['label']} in {s['buckets']}",train_n=s['train_n'],train_mean=s['train_mean'],train_always=s['train_always_mean'],oos_n=len(f),oos_mean=f.pnl.mean() if len(f) else np.nan,always_mean=te.pnl.mean(),
                    p=tp(f.pnl.values) if len(f)>2 else 1.0,subset_pctile=subset_pct(te.pnl.values,f.pnl.values),random_regime_pctile=rpc,random_regime_mean=np.mean(np.concatenate(arrs)) if arrs else np.nan))
# spot walk-forward
WF=[]
for c in C:
    x=SP[SP.cand==c]; oos=[]; allo=[]
    for Y in (2022,2023,2024,2025):
        tr=x[x.date.dt.year<=Y]; te=x[x.date.dt.year==Y+1]
        if len(te)==0 or len(tr)==0: continue
        s=select(tr,'spot_net',15,30); allo.append(te)
        if s is None: WF.append(dict(cand=c,test_year=Y+1,rule='none (trade nothing)',train_n=len(tr),oos_n=0,oos_mean=np.nan,always_mean=te.spot_net.mean())); continue
        f=te[te[s['label']].isin(s['buckets'])]; oos.append(f)
        WF.append(dict(cand=c,test_year=Y+1,rule=f"{s['label']} in {s['buckets']}",train_n=s['train_n'],train_mean=s['train_mean'],oos_n=len(f),oos_mean=f.spot_net.mean() if len(f) else np.nan,always_mean=te.spot_net.mean()))
    fo=pd.concat(oos) if oos else x.iloc[:0]; ao=pd.concat(allo) if allo else x.iloc[:0]
    FAM.append(dict(test='spot_walkforward',cand=c,rule='walk-forward (see wf table)',oos_n=len(fo),oos_mean=fo.spot_net.mean() if len(fo) else np.nan,always_mean=ao.spot_net.mean(),
                    p=tp(fo.spot_net.values) if len(fo)>2 else 1.0,subset_pctile=subset_pct(ao.spot_net.values,fo.spot_net.values)))
pd.DataFrame(WF).round(2).to_csv(f'{RS}/spot_walkforward.csv',index=False)
# meta switch
def switch(x,val,N=None,weeks=None):
    x=x.sort_values('date'); keep=[]; p=x[val].values; d=x.date.values
    for i in range(len(x)):
        prior=x[x.date<x.date.iloc[i]]
        if weeks: prior=prior[prior.date>=x.date.iloc[i]-pd.Timedelta(weeks=weeks)]; ok=len(prior)>0 and prior[val].sum()>0
        else: ok=len(prior)>=N and prior[val].iloc[-N:].sum()>0
        keep.append(ok)
    return x[np.array(keep)]
MS=[]
for c in C:
    x=O[O.cand==c]; te=x[x.split=='test']
    for nm,kw in (('M20',dict(N=20)),('M40',dict(N=40)),('8w',dict(weeks=8))):
        f=switch(x,'pnl',**kw); f=f[f.split=='test']
        r=dict(data='options_2026',cand=c,rule=nm,oos_n=len(f),oos_mean=f.pnl.mean() if len(f) else np.nan,always_n=len(te),always_mean=te.pnl.mean(),p=tp(f.pnl.values) if len(f)>2 else 1.0,subset_pctile=subset_pct(te.pnl.values,f.pnl.values))
        MS.append(r)
        if nm=='M20': FAM.append(dict(test='opt_M20',**{k:v for k,v in r.items() if k not in('data','rule','always_n')},rule='M20'))
    y=SP[SP.cand==c]; ye=y[y.date.dt.year>=2023]
    for nm,kw in (('M20',dict(N=20)),('M40',dict(N=40)),('8w',dict(weeks=8))):
        f=switch(y,'spot_net',**kw); f=f[f.date.dt.year>=2023]
        r=dict(data='spot_2023_26',cand=c,rule=nm,oos_n=len(f),oos_mean=f.spot_net.mean() if len(f) else np.nan,always_n=len(ye),always_mean=ye.spot_net.mean(),p=tp(f.spot_net.values) if len(f)>2 else 1.0,subset_pctile=subset_pct(ye.spot_net.values,f.spot_net.values))
        MS.append(r)
        if nm=='M20': FAM.append(dict(test='spot_M20',**{k:v for k,v in r.items() if k not in('data','rule','always_n')},rule='M20'))
pd.DataFrame(MS).round(3).to_csv(f'{RS}/meta_switch.csv',index=False)
FM=pd.DataFrame(FAM); FM['q_bh']=bh(FM.p.values)
FM['counts']=(FM.q_bh<0.10)&(FM.oos_mean>FM.always_mean)&(FM.subset_pctile>=90)
FM.round(3).to_csv(f'{RS}/regime_family.csv',index=False)
json.dump({c:CH[c] for c in C},open(f'{RS}/frozen_regime_choices.json','w'),default=str,indent=1)
json.dump(dict(n_labels=len(LBS),n_buckets=NCUTS,labels=LBS,bucket_results_reported=len(RB),family_size=len(FM),
               selection_space_per_strategy_options=f'{len(LBS)} labels x all subsets of positive buckets (rule-based, 1 pick)'),open(f'{RS}/cuts_count.json','w'),indent=1)
pd.set_option('display.width',250); pd.set_option('display.max_colwidth',60)
print(P.round(2).to_string()); print(FM.round(3).to_string()); print(pd.DataFrame(WF).round(2).to_string()); print(pd.DataFrame(MS).round(2).to_string()); print('cuts',NCUTS,len(RB))
# ---------- 5. charts ----------
def grid(fn,title,draw):
    fig,ax=plt.subplots(2,4,figsize=(20,8.5)); 
    for i,c in enumerate(C): a=ax.flat[i]; draw(a,c); a.set_title(f'{c} {NAMES[c]}',fontsize=10)
    fig.suptitle(title); fig.tight_layout(); fig.savefig(f'{CHD}/{fn}',dpi=80); plt.close(fig)
def roll(a,c):
    y=SP[SP.cand==c].sort_values('date'); a.plot(y.date,y.spot_net.rolling(60).mean(),label='spot net pts (2022-26)',color='tab:blue')
    x=O[O.cand==c].sort_values('date'); a.plot(x.date,x.pnl.rolling(60).mean(),label='option net pts (ATM, T50)',color='tab:red')
    a.axhline(0,color='k',lw=.8); a.axvline(pd.Timestamp('2026-01-01'),color='gray',ls='--',lw=.8); a.legend(fontsize=7)
grid('rolling60.png','Rolling 60-trade expectancy (net points per trade). Above 0 = "on" period. Dashed line = start of 2026 test',roll)
def heat(a,c):
    y=SP[SP.cand==c]; m=y.groupby([y.date.dt.year,y.date.dt.month]).spot_net.mean().unstack()
    m=m.reindex(columns=range(1,13)); im=a.imshow(m.values,cmap='RdYlGn',vmin=-15,vmax=15,aspect='auto'); a.set_yticks(range(len(m))); a.set_yticklabels(m.index); a.set_xticks(range(12)); a.set_xticklabels('JFMAMJJASOND')
    for i in range(m.shape[0]):
        for j in range(12):
            v=m.values[i,j]
            if v==v: a.text(j,i,f'{v:.0f}',ha='center',va='center',fontsize=6)
grid('heatmap_spot_month.png','Spot net points per trade by month (2022-2026, trail-50% exit, -1 pt cost)',heat)
def heato(a,c):
    y=O[O.cand==c]; m=y.groupby([y.date.dt.year,y.date.dt.month]).pnl.mean().unstack().reindex(columns=range(1,13))
    a.imshow(m.values,cmap='RdYlGn',vmin=-10,vmax=10,aspect='auto'); a.set_yticks(range(len(m))); a.set_yticklabels(m.index); a.set_xticks(range(12)); a.set_xticklabels('JFMAMJJASOND')
    for i in range(m.shape[0]):
        for j in range(12):
            v=m.values[i,j]
            if v==v: a.text(j,i,f'{v:.0f}',ha='center',va='center',fontsize=6)
grid('heatmap_options_month.png','Option buyer net points per trade by month (ATM weekly, Oct-2024..Oct-2026, REAL costs)',heato)
# regime heatmap: options train vs test means by bucket
for dsname,s1,s2 in (('options','train','test'),('spot','2022-25','2026')):
    R_=RB[RB.data==dsname]; R_=R_.assign(col=R_.label.str.split('_').str[0]+':'+R_.bucket.astype(str))
    fig,ax=plt.subplots(2,1,figsize=(22,11))
    for a,s in zip(ax,(s1,s2)):
        m=R_[R_.split==s].pivot_table(index='cand',columns='col',values='mean'); n=R_[R_.split==s].pivot_table(index='cand',columns='col',values='n')
        lim=10 if dsname=='options' else 15
        a.imshow(m.values,cmap='RdYlGn',vmin=-lim,vmax=lim,aspect='auto'); a.set_yticks(range(len(m))); a.set_yticklabels(m.index); a.set_xticks(range(m.shape[1])); a.set_xticklabels(m.columns,rotation=60,fontsize=8); a.set_title(f'{dsname} {s}: mean net pts by regime bucket (n in brackets)')
        for i in range(m.shape[0]):
            for j in range(m.shape[1]):
                v=m.values[i,j]
                if v==v: a.text(j,i,f'{v:.0f}\n({n.values[i,j]:.0f})',ha='center',va='center',fontsize=7)
    fig.tight_layout(); fig.savefig(f'{CHD}/regime_heat_{dsname}.png',dpi=80); plt.close(fig)
print('charts done')