import numpy as np, pandas as pd, itertools, json
from load import load
P=load()
ETF="SPY QQQ IWM DIA XLK XLF XLE XLV XLY XLP XLI XLB XLU XLRE XLC GLD TLT".split()
C=P['C']; cols=C.columns
def arr(k,u): return P[k][u].reindex(C.index).values
def ma(x,n): return pd.DataFrame(x).rolling(n,min_periods=n).mean().values
def signals(c,v):
    m={n:ma(c,n) for n in (5,10,15,20,30)}
    sd=pd.DataFrame(c).rolling(20,min_periods=20).std().values
    up=m[20]+2*sd; lo=m[20]-2*sd
    prev=lambda x: np.vstack([np.zeros((1,x.shape[1]),bool),x[:-1]])
    cross=(m[5]>m[20])&~prev(m[5]>m[20])&(c>m[30])
    stack=(m[5]>m[10])&(m[10]>m[15])&(m[15]>m[20])&(m[20]>m[30]); stack_new=stack&~prev(stack)
    bbu=c>up; bb_break=bbu&~prev(bbu)
    bbl=c<lo; bb_rev=bbl&~prev(bbl)
    bb_stack=bb_break&stack
    dv=pd.DataFrame(c*v).rolling(20,min_periods=20).mean().values
    return dict(MAcross=cross,MAstack=stack_new,BBbreak=bb_break,BBrevert=bb_rev,BBbreak_stack=bb_stack),dv
def sim(univ,sig,dv,o,h,l,c,tp,sl,maxh,K=5,slip=0.0005,fixed_fee=0.0,cap0=5000.0,elig=None,tp_first=False):
    T,N=c.shape; eq=np.full(T,np.nan); cash=cap0; pos={}  # j:(shares,entry,tpx,slx,days)
    trades=[]
    rank=np.nan_to_num(dv,nan=-1)
    for t in range(1,T):
        # entries at open t from signals at t-1
        nopen=K-len(pos)
        if nopen>0:
            s=sig[t-1].copy()
            if elig is not None: s&=elig[t-1]
            s&=~np.isnan(o[t])
            cand=[j for j in np.where(s)[0] if j not in pos]
            cand.sort(key=lambda j:-rank[t-1,j])
            mtm=cash+sum(sh*c[t-1,j] for j,(sh,e,a,b,d) in pos.items())
            for j in cand[:nopen]:
                alloc=min(mtm/K,cash)
                if alloc<50: break
                px=o[t,j]*(1+slip); sh=(alloc-fixed_fee)/px
                cash-=alloc; pos[j]=(sh,px,px/(1+slip)*(1+tp),px/(1+slip)*(1-sl),0)
        # exits during day t
        for j in list(pos):
            sh,e,tpx,slx,d=pos[j]
            if np.isnan(c[t,j]):  # missing bar: hold
                continue
            d+=1; ex=None
            if d>1 and o[t,j]<=slx: ex=o[t,j]
            elif d>1 and o[t,j]>=tpx: ex=o[t,j]
            elif tp_first and h[t,j]>=tpx: ex=tpx
            elif l[t,j]<=slx: ex=slx
            elif h[t,j]>=tpx: ex=tpx
            elif d>=maxh: ex=c[t,j]
            if ex is not None:
                fill=ex*(1-slip); cash+=sh*fill-fixed_fee; trades.append((t,j,fill/e-1,d)); del pos[j]
            else: pos[j]=(sh,e,tpx,slx,d)
        eq[t]=cash+sum(sh*(c[t,j] if not np.isnan(c[t,j]) else e) for j,(sh,e,a,b,d) in pos.items())
    eq[0]=cap0
    return eq,trades
def metrics(eq,idx,start=None,end=None):
    s=pd.Series(eq,idx).dropna()
    if start: s=s[s.index>=start]
    if end: s=s[s.index<=end]
    r=s.pct_change().dropna(); yrs=len(r)/252
    cagr=(s.iloc[-1]/s.iloc[0])**(1/yrs)-1 if yrs>0 else np.nan
    sh=r.mean()/r.std()*np.sqrt(252) if r.std()>0 else 0
    mdd=(s/s.cummax()-1).min()
    return dict(ret=s.iloc[-1]/s.iloc[0]-1,cagr=cagr,vol=r.std()*np.sqrt(252),sharpe=sh,mdd=mdd)