import json,gzip,pandas as pd,numpy as np
from sklearn.metrics import roc_auc_score
pd.set_option('display.width',250)
R=json.load(gzip.open('/tmp/run/rt.json.gz','rt'))
S={k.lower():v for k,v in R['supply'].items()}; K={k.lower():v for k,v in R['kline'].items()}
CR={t['address'].lower():t['created_ts'] for t in R['tokens']}
F=pd.read_pickle('f30.pkl')[['tok','grp']].set_index('tok').grp
P=pd.read_pickle('paths.pkl').set_index('tok')
rows=[]
for a,k in K.items():
    s=S.get(a)
    if not s or a not in F.index: continue
    H=[c for c in k if c[5]>=1000]
    if len(H)<3: continue
    if k[0][0]>(CR.get(a) or 0)+3600*2: continue      # kline window starts after launch: start not visible
    T0=next((c[0] for c in H if c[5]>=5000),None)
    if T0 is None: continue
    tt=[(c[0]-T0)/3600 for c in H]; cl=[c[4]*s for c in H]; v=[c[5] for c in H]
    r=dict(tok=a,grp=F[a],start=H[0][1]*s)
    for h in (1,2,4,8,12,24,48,72):
        idx=[i for i,x in enumerate(tt) if 0<=x<h]
        if not idx: continue
        r['mc_%dh'%h]=cl[idx[-1]]; r['peak_%dh'%h]=max(cl[i] for i in idx); r['vol_%dh'%h]=sum(v[i] for i in idx)
        r['dd_%dh'%h]=cl[idx[-1]]/r['peak_%dh'%h]; r['volmc_%dh'%h]=r['vol_%dh'%h]/max(1,r['peak_%dh'%h])
        r['green_%dh'%h]=np.mean([H[i][4]>=H[i][1] for i in idx])
    rows.append(r)
x=pd.DataFrame(rows); print(x.grp.value_counts())
x.to_pickle('hstage.pkl')
med=x.groupby('grp').median(numeric_only=True).T
print(med[[c for c in ['R20_small','RT_5to20M','RT_1to5M','random_launch'] if c in med.columns]].round(2).to_string())
A=x[x.grp=='R20_small']; B=x[x.grp=='RT_1to5M']
for c in [c for c in x.columns if any(c.startswith(p) for p in ('mc_','dd_','volmc_','green_'))]:
    a=A[c].dropna(); b=B[c].dropna()
    if len(a)>5 and len(b)>5: print(c,'AUC R20 vs 1-5M',round(roc_auc_score([1]*len(a)+[0]*len(b),pd.concat([a,b])),2))