# Taken-vs-skipped feature contrast, depth-2 tree + logistic with leave-one-day-out CV (fit May-Jul), filter extraction, holdout lift (Aug-Oct). PREREG TRADER-DISCRETION.
import pandas as pd,numpy as np,json
from scipy import stats
from sklearn.tree import DecisionTreeClassifier,export_text
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.impute import SimpleImputer
from sklearn.metrics import roc_auc_score
RS='/workspace/research/popular-strategies/results'
CD=pd.read_csv(f'{RS}/disc_candidates.csv',parse_dates=['date'])
FE=['tod','prem_dd','prem_chg15','prem_chg5','spot_pos','spot_mv15','dist_round50','dist_round100','vix','gap','oi_chg15','pcr_chg15','basis','basis_chg15','prior_trades','expiry_day','dte']
def bh(p):
    p=np.where(np.isnan(np.asarray(p,float)),1,np.asarray(p,float)); 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)
rows=[]
for f in FE:
    a=CD[CD.taken==1][f].dropna(); b=CD[CD.taken==0][f].dropna(); sd=np.sqrt((a.var()*(len(a)-1)+b.var()*(len(b)-1))/(len(a)+len(b)-2))
    rows.append(dict(feature=f,taken_mean=a.mean(),skipped_mean=b.mean(),smd=(a.mean()-b.mean())/sd if sd>0 else np.nan,mw_p=stats.mannwhitneyu(a,b).pvalue if len(a)>2 and len(b)>2 else np.nan,n_taken=len(a),n_skipped=len(b)))
FC=pd.DataFrame(rows); FC['q_bh']=bh(FC.mw_p.values); FC=FC.sort_values('mw_p'); FC.round(3).to_csv(f'{RS}/disc_feature_contrast.csv',index=False); print(FC.round(3).to_string())
fit=CD[CD.date<pd.Timestamp('2026-08-01')].reset_index(drop=True); ho=CD[CD.date>=pd.Timestamp('2026-08-01')].reset_index(drop=True)
def mk(kind): return make_pipeline(SimpleImputer(strategy='median'),DecisionTreeClassifier(max_depth=2,min_samples_leaf=10,random_state=0)) if kind=='tree' else make_pipeline(SimpleImputer(strategy='median'),StandardScaler(),LogisticRegression(C=1.0,max_iter=2000))
res={}
for kind in ('tree','logit'):
    pr=np.full(len(fit),np.nan)
    for d in fit.date.unique():
        tr=fit.date!=d
        if fit[tr].taken.nunique()<2: continue
        mdl=mk(kind).fit(fit.loc[tr,FE],fit.loc[tr,'taken']); pr[~tr]=mdl.predict_proba(fit.loc[~tr,FE])[:,1]
    ok=~np.isnan(pr); mdl=mk(kind).fit(fit[FE],fit.taken)
    res[kind]=dict(lodo_auc=roc_auc_score(fit.taken[ok],pr[ok]),holdout_auc=roc_auc_score(ho.taken,mdl.predict_proba(ho[FE])[:,1]))
    if kind=='tree': TREE=mdl; res['tree_text']=export_text(mdl[-1],feature_names=FE)
    else: res['logit_coef']=dict(zip(FE,np.round(mdl[-1].coef_[0],3)))
print(json.dumps(res,indent=1,default=float))
# filter extraction: leaf with highest take rate among leaves with >=10 takes
t=TREE[-1].tree_; Xf=TREE[0].transform(fit[FE]); leaf=TREE[-1].apply(Xf); base=fit.taken.mean(); best=None
for lf in np.unique(leaf):
    s=fit.taken[leaf==lf]
    if s.sum()>=10 and (best is None or s.mean()>best[1]): best=(lf,s.mean(),int(s.sum()),len(s))
def path(lf):
    # reconstruct conditions to leaf
    conds=[]; 
    def rec(node,acc):
        if node==lf: conds.extend(acc); return True
        if t.children_left[node]==-1: return False
        f=FE[t.feature[node]]; th=t.threshold[node]
        return rec(t.children_left[node],acc+[(f,'<=',th)]) or rec(t.children_right[node],acc+[(f,'>',th)])
    rec(0,[]); return conds
rule=None
if best and best[1]>=1.2*base: rule=path(best[0])
med=dict(zip(FE,TREE[0].statistics_))
def apply(df):
    if rule is None: return np.zeros(len(df),bool)
    m=np.ones(len(df),bool)
    for f,op,th in rule:
        x=df[f].fillna(med[f]).values; m&=(x<=th) if op=='<=' else (x>th)
    return m
out=dict(base_take_rate_fit=base,rule=[(f,op,float(th)) for f,op,th in rule] if rule else None,leaf_take_rate=best[1] if best else None,leaf_takes=best[2] if best else None,leaf_n=best[3] if best else None)
if rule:
    h=apply(ho); lift=ho.taken[h].mean()/ho.taken.mean() if h.sum() else np.nan
    rng=np.random.default_rng(1); days=ho.date.unique(); bs=[]
    for _ in range(2000):
        dd=rng.choice(days,len(days)); z=pd.concat([ho[ho.date==x].assign(_h=h[(ho.date==x).values]) for x in dd])
        if z._h.sum() and z.taken.mean()>0: bs.append(z.taken[z._h].mean()/z.taken.mean())
    out.update(holdout_n=int(h.sum()),holdout_takes=int(ho.taken[h].sum()),holdout_base=ho.taken.mean(),holdout_lift=lift,lift_ci90=[float(np.percentile(bs,5)),float(np.percentile(bs,95))])
    out['uses_prior_trades']=any(f=='prior_trades' for f,_,_ in rule)
out.update(res); json.dump(out,open(f'{RS}/disc_model.json','w'),indent=1,default=float); print(json.dumps({k:v for k,v in out.items() if k!='tree_text'},indent=1,default=float))
# outcomes: taken vs skipped candidates
print(CD.groupby('taken')[['xold','xt50']].agg(['size','mean']).round(2)); CD.groupby(['taken']).agg(n=('xold','size'),xold=('xold','mean'),xt50=('xt50','mean'),xold_win=('xold',lambda s:(s>0).mean()*100)).round(2).reset_index().to_csv(f'{RS}/disc_taken_vs_skipped_outcomes.csv',index=False)
json.dump(dict(med=med),open(f'{RS}/disc_impute_medians.json','w'),default=float)