# ATM CE buy P&L for signals, controls and random entries (real 1m option bars; no filling)
import pandas as pd,numpy as np,glob,os,sys,json,pyarrow.parquet as pq
sys.path.insert(0,'scripts'); from up import get,q
rng=np.random.default_rng(7)
OD='/workspace/nifty-options-s3/data/options/1m/by_expiry/'
EXPS=sorted(pd.Timestamp(os.path.basename(f)[:10]) for f in glob.glob(OD+'*.parquet'))
OCT=[pd.Timestamp(x) for x in ['2026-10-06','2026-10-13','2026-10-19','2026-10-27']]
ALLE=sorted(set(EXPS)|set(OCT))
E=pd.read_csv('results/p2_events.csv'); E['date']=pd.to_datetime(E.date)
P=pd.read_parquet('data/minute_panel_v.parquet'); P=P[P.hm<='15:14']
first_opt=pd.Timestamp('2024-09-27')  # first day covered by local expiries (2024-10-03 weekly)
# random pool: days with fvr & option coverage, minutes 09:20-14:14
pool=P[(P.date>=pd.Timestamp('2024-10-21'))&P.fvr.notna()&P.close.notna()&(P.hm>='09:20')&(P.hm<='14:14')]
R=pool.sample(4000,random_state=7)[['date','hm','close']].rename(columns={'close':'entry'}); R['kind']='random'
EV=pd.concat([E[(E.date>=first_opt)&E.kind.isin(['signal','control_quiet','price_bar','at_low_any'])],R],ignore_index=True)
def nexp(d,skip0=False):
    c=[e for e in ALLE if e>=d and not (skip0 and e==d)]; return c[0] if c else None
EV['exp0']=EV.date.map(nexp); EV['exp1']=EV.date.map(lambda d:nexp(d,True)); EV['K']=(EV.entry/50).round()*50
need={}
for r in EV.itertuples():
    for e in (r.exp0,r.exp1):
        if e is not None: need.setdefault(e,set()).add((r.date,float(r.K)))
CE={}
cons={}
for e,s in sorted(need.items()):
    ks=sorted({k for _,k in s}); ds={d for d,_ in s}
    if e in EXPS:
        t=pq.read_table(OD+e.strftime('%Y-%m-%d')+'.parquet',columns=['timestamp','option_type','strike','open','high','low','close'],
            filters=[('option_type','=','CE'),('strike','in',ks)]).to_pandas()
        t['date']=t.timestamp.dt.normalize(); t=t[t.date.isin(ds)]; t['hm']=t.timestamp.dt.strftime('%H:%M')
        for (d,k),g in t.groupby(['date','strike']): CE[(d,e,k)]=g.drop_duplicates('hm').set_index('hm').sort_index()
    elif e in OCT:   # Oct-2026: Upstox (expired API for 10-06/10-13/10-19; live for 10-27)
        if e not in cons:
            if e<pd.Timestamp('2026-10-25'):
                c,j=get('/v2/expired-instruments/option/contract',{'instrument_key':'NSE_INDEX|Nifty 50','expiry_date':e.strftime('%Y-%m-%d')})
                cons[e]={float(x['strike_price']):x['instrument_key'] for x in j.get('data',[]) if x['instrument_type']=='CE'}
            else:
                import gzip; ins=json.load(gzip.open('../data/NSE.json.gz'))
                cons[e]={}
                for x in ins:
                    if x.get('segment')=='NSE_FO' and x.get('name')=='NIFTY' and x.get('instrument_type')=='CE':
                        ee=pd.to_datetime(x['expiry'],unit='ms') if isinstance(x['expiry'],(int,float)) else pd.Timestamp(str(x['expiry'])[:10])
                        if ee.normalize()==e: cons[e][float(x['strike_price'])]=x['instrument_key']
        for d,k in s:
            key=cons[e].get(k)
            if key is None: continue
            ds_=d.strftime('%Y-%m-%d')
            path=(f"/v2/expired-instruments/historical-candle/{q(key,safe='')}/1minute/{ds_}/{ds_}" if e<pd.Timestamp('2026-10-25') else f"/v2/historical-candle/{q(key,safe='')}/1minute/{ds_}/{ds_}")
            c,j=get(path)
            if c!=200 or not j['data']['candles']: continue
            g=pd.DataFrame(j['data']['candles'],columns=['ts','open','high','low','close','volume','oi']); g['hm']=g.ts.str[11:16]
            CE[(d,e,k)]=g.drop_duplicates('hm').set_index('hm').sort_index()
print('option series loaded',len(CE),flush=True)
SLIP=3.0
def lot(t): return 25 if t<pd.Timestamp('2024-11-20') else (75 if t<pd.Timestamp('2026-01-01') else 65)
def hmadd(h,m): return (pd.Timestamp('2000-01-01 '+h)+pd.Timedelta(minutes=m)).strftime('%H:%M')
def charges(t,buy,sells):  # buy price, list of (qty_frac, price); per-unit points for 1 lot
    Q=lot(t); n_orders=1+len(sells); brk=20.0*n_orders
    sv=sum(f*p for f,p in sells)*Q; bv=buy*Q
    stt=(0.001 if t>=pd.Timestamp('2024-10-01') else 0.000625)*sv
    txn=0.0003503*(bv+sv); sebi=1e-6*(bv+sv); stamp=3e-5*bv; gst=0.18*(brk+txn+sebi)
    return (brk+stt+txn+sebi+stamp+gst)/Q
def trade(d,e,k,hm,style):
    s=CE.get((d,e,k))
    if s is None: return None
    h1=hmadd(hm,1)
    if h1 not in s.index: return None          # no bar -> no trade (no filling)
    fill=s.open[h1]+SLIP
    end=min(hmadd(hm,61),'15:25'); path=s.loc[h1:end]
    if len(path)<5: return None
    rem=1.0; sells=[]; stop=fill-20; peak=fill; t1=t2=False
    for i,(x,b) in enumerate(path.iterrows()):
        if style=='time': continue
        if b.low<=stop:                          # stop first (conservative)
            px=min(b.open,stop) if i>0 else stop
            sells.append((rem,px-SLIP)); rem=0; break
        if style=='scale':
            if not t1 and b.high>=fill+20: sells.append((0.4,fill+20-SLIP)); rem-=0.4; t1=True; stop=max(stop,fill)
            if not t2 and b.high>=fill+40: sells.append((0.3,fill+40-SLIP)); rem-=0.3; t2=True
            peak=max(peak,b.high)
            if t2: stop=max(stop,peak-20,fill+20)
    if rem>1e-9: sells.append((rem,path.close.iloc[-1]-SLIP))
    gross=sum(f*p for f,p in sells)-fill
    return dict(fill=fill,gross=gross,cost=charges(d,fill,sells),pnl=gross-charges(d,fill,sells),exit_last=path.index[-1] if rem>0 else None)
out=[]
for r in EV.itertuples():
    for lab,e in (('nearest',r.exp0),('skip0DTE',r.exp1)):
        if e is None: continue
        for style in ('scale','sl_only'):
            z=trade(r.date,e,float(r.K),r.hm,style)
            if z: out.append(dict(idx=r.Index,date=r.date,hm=r.hm,kind=r.kind,X=r.X,V=r.V,win=r.win,oi=r.oi,expiry_rule=lab,dte=(e-r.date).days,style=style,**z))
T=pd.DataFrame(out); T.to_csv('results/p2_option_trades.csv',index=False)
print(T.groupby(['kind','expiry_rule','style']).pnl.agg(['count','mean','median']).round(2).to_string())