import pandas as pd
import numpy as np

def getDailyVol(close,
                span0=100):
    """SNIPPET 3.1 DAILY VOLATILITY ESTIMATES
    Daily vol reindexed to close
    """ 
    df0=close.index.searchsorted(close.index-pd.Timedelta(days=1))
    df0=df0[df0>0]   
    df0=(pd.Series(close.index[df0-1], 
                   index=close.index[close.shape[0]-df0.shape[0]:]))   
    try:
        df0=close.loc[df0.index]/close.loc[df0.values].values-1 # daily rets
    except Exception as e:
        print(f'error: {e}\nplease confirm no duplicate indices')
    df0=df0.ewm(span=span0).std().rename('dailyVol')
    return df0

def applyPtSlOnT1(close,
                  events,
                  ptSl,
                  molecule):
    """SNIPPET 3.2 TRIPLE-BARRIER LABELING METHOD
    Apply stop loss/profit taking, if it takes place before t1 (end of event)
    """
    events_=events.loc[molecule]
    out=events_[['t1']].copy(deep=True)
    if ptSl[0]>0: pt=ptSl[0]*events_['trgt']
    else: pt=pd.Series(index=events.index) # NaNs
    if ptSl[1]>0: sl=-ptSl[1]*events_['trgt']
    else: sl=pd.Series(index=events.index) # NaNs
    for loc,t1 in events_['t1'].fillna(close.index[-1]).iteritems():
        df0=close[loc:t1] # path prices
        df0=(df0/close[loc]-1)*events_.at[loc,'side'] # path returns
        out.loc[loc,'sl']=df0[df0<sl[loc]].index.min() # earliest stop loss
        out.loc[loc,'pt']=df0[df0>pt[loc]].index.min() # earliest profit taking
    return out

def getEvents(close,
              tEvents,
              ptSl,
              trgt,
              minRet,
              numThreads,
              t1=False,
              side=None):
    """SNIPPET 3.3 GETTING THE TIME OF FIRST TOUCH
    SNIPPET 3.6 EXPANDING getEvents TO INCORPORATEMETA-LABELING
    """
    #1) get target
    trgt=trgt.loc[tEvents]
    trgt=trgt[trgt>minRet] # minRet
    #2) get t1 (max holding period)
    if t1 is False:t1=pd.Series(pd.NaT, index=tEvents)
    #3) form events object, apply stop loss on t1
    if side is None:side_,ptSl_=pd.Series(1.,index=trgt.index), [ptSl[0],ptSl[0]]
    else: side_,ptSl_=side.loc[trgt.index],ptSl[:2]
    events=(pd.concat({'t1':t1,'trgt':trgt,'side':side_}, axis=1)
            .dropna(subset=['trgt']))
    df0=mpPandasObj(func=applyPtSlOnT1,pdObj=('molecule',events.index),
                    numThreads=numThreads,close=close,events=events,
                    ptSl=ptSl_)
    events['t1']=df0.dropna(how='all').min(axis=1) # pd.min ignores nan
    if side is None:events=events.drop('side',axis=1)
    return events

def addVerticalBarrier(tEvents,
                       close,
                       numDays=1):
    """SNIPPET 3.4 ADDING A VERTICAL BARRIER
    """
    t1=close.index.searchsorted(tEvents+pd.Timedelta(days=numDays))
    t1=t1[t1<close.shape[0]]
    t1=(pd.Series(close.index[t1],index=tEvents[:t1.shape[0]]))
    return t1

def getBinsOld(events,
               close):
    """SNIPPET 3.5 LABELING FOR SIDE AND SIZE
    """
    #1) prices aligned with events
    events_=events.dropna(subset=['t1'])
    px=events_.index.union(events_['t1'].values).drop_duplicates()
    px=close.reindex(px,method='bfill')
    #2) create out object
    out=pd.DataFrame(index=events_.index)
    out['ret']=px.loc[events_['t1'].values].values/px.loc[events_.index]-1
    out['bin']=np.sign(out['ret'])
    # where out index and t1 (vertical barrier) intersect label 0
    try:
        locs = out.query('index in @t1').index
        out.loc[locs, 'bin'] = 0
    except:
        pass
    return out

def getBins(events,
            close):
    """SNIPPET 3.7 EXPANDING getBins TO INCORPORATE META-LABELING
    Compute event's outcome (including side information, if provided).
    events is a DataFrame where:
    -events.index is event's starttime
    -events['t1'] is event's endtime
    -events['trgt'] is event's target
    -events['side'] (optional) implies the algo's position side
    Case 1: ('side' not in events): bin in (-1,1) <-label by price action
    Case 2: ('side' in events): bin in (0,1) <-label by pnl (meta-labeling)
    """
    #1) prices aligned with events
    events_=events.dropna(subset=['t1'])
    px=events_.index.union(events_['t1'].values).drop_duplicates()
    px=close.reindex(px,method='bfill')
    #2) create out object
    out=pd.DataFrame(index=events_.index)
    out['ret']=px.loc[events_['t1'].values].values/px.loc[events_.index]-1
    if 'side' in events_:out['ret']*=events_['side'] # meta-labeling
    out['bin']=np.sign(out['ret'])
    if 'side' in events_:out.loc[out['ret']<=0,'bin']=0 # meta-labeling
    return out

def dropLabels(events,
               minPct=.05):
    """SNIPPET 3.8 DROPPING UNDER-POPULATED LABELS
    Apply weights, drop labels with insufficient examples
    """
    while True:
        df0=events['bin'].value_counts(normalize=True)
        if df0.min()>minPct or df0.shape[0]<3:break
        print('dropped label: ', df0.argmin(),df0.min())
        events=events[events['bin']!=df0.argmin()]
    return events


import pandas as pd
import numpy as np

def getDailyVol(close,
                span0=100):
    """SNIPPET 3.1 DAILY VOLATILITY ESTIMATES
    Daily vol reindexed to close
    """ 
    df0=close.index.searchsorted(close.index-pd.Timedelta(days=1))
    df0=df0[df0>0]   
    df0=(pd.Series(close.index[df0-1], 
                   index=close.index[close.shape[0]-df0.shape[0]:]))   
    try:
        df0=close.loc[df0.index]/close.loc[df0.values].values-1 # daily rets
    except Exception as e:
        print(f'error: {e}\nplease confirm no duplicate indices')
    df0=df0.ewm(span=span0).std().rename('dailyVol')
    return df0

def applyPtSlOnT1(close,
                  events,
                  ptSl,
                  molecule):
    """SNIPPET 3.2 TRIPLE-BARRIER LABELING METHOD
    Apply stop loss/profit taking, if it takes place before t1 (end of event)
    """
    events_=events.loc[molecule]
    out=events_[['t1']].copy(deep=True)
    if ptSl[0]>0: pt=ptSl[0]*events_['trgt']
    else: pt=pd.Series(index=events.index) # NaNs
    if ptSl[1]>0: sl=-ptSl[1]*events_['trgt']
    else: sl=pd.Series(index=events.index) # NaNs
    for loc,t1 in events_['t1'].fillna(close.index[-1]).iteritems():
        df0=close[loc:t1] # path prices
        df0=(df0/close[loc]-1)*events_.at[loc,'side'] # path returns
        out.loc[loc,'sl']=df0[df0<sl[loc]].index.min() # earliest stop loss
        out.loc[loc,'pt']=df0[df0>pt[loc]].index.min() # earliest profit taking
    return out

def getEvents(close,
              tEvents,
              ptSl,
              trgt,
              minRet,
              numThreads,
              t1=False,
              side=None):
    """SNIPPET 3.3 GETTING THE TIME OF FIRST TOUCH
    SNIPPET 3.6 EXPANDING getEvents TO INCORPORATEMETA-LABELING
    """
    #1) get target
    trgt=trgt.loc[tEvents]
    trgt=trgt[trgt>minRet] # minRet
    #2) get t1 (max holding period)
    if t1 is False:t1=pd.Series(pd.NaT, index=tEvents)
    #3) form events object, apply stop loss on t1
    if side is None:side_,ptSl_=pd.Series(1.,index=trgt.index), [ptSl[0],ptSl[0]]
    else: side_,ptSl_=side.loc[trgt.index],ptSl[:2]
    events=(pd.concat({'t1':t1,'trgt':trgt,'side':side_}, axis=1)
            .dropna(subset=['trgt']))
    df0=mpPandasObj(func=applyPtSlOnT1,pdObj=('molecule',events.index),
                    numThreads=numThreads,close=close,events=events,
                    ptSl=ptSl_)
    events['t1']=df0.dropna(how='all').min(axis=1) # pd.min ignores nan
    if side is None:events=events.drop('side',axis=1)
    return events

def addVerticalBarrier(tEvents,
                       close,
                       numDays=1):
    """SNIPPET 3.4 ADDING A VERTICAL BARRIER
    """
    t1=close.index.searchsorted(tEvents+pd.Timedelta(days=numDays))
    t1=t1[t1<close.shape[0]]
    t1=(pd.Series(close.index[t1],index=tEvents[:t1.shape[0]]))
    return t1

def getBinsOld(events,
               close):
    """SNIPPET 3.5 LABELING FOR SIDE AND SIZE
    """
    #1) prices aligned with events
    events_=events.dropna(subset=['t1'])
    px=events_.index.union(events_['t1'].values).drop_duplicates()
    px=close.reindex(px,method='bfill')
    #2) create out object
    out=pd.DataFrame(index=events_.index)
    out['ret']=px.loc[events_['t1'].values].values/px.loc[events_.index]-1
    out['bin']=np.sign(out['ret'])
    # where out index and t1 (vertical barrier) intersect label 0
    try:
        locs = out.query('index in @t1').index
        out.loc[locs, 'bin'] = 0
    except:
        pass
    return out

def getBins(events,
            close):
    """SNIPPET 3.7 EXPANDING getBins TO INCORPORATE META-LABELING
    Compute event's outcome (including side information, if provided).
    events is a DataFrame where:
    -events.index is event's starttime
    -events['t1'] is event's endtime
    -events['trgt'] is event's target
    -events['side'] (optional) implies the algo's position side
    Case 1: ('side' not in events): bin in (-1,1) <-label by price action
    Case 2: ('side' in events): bin in (0,1) <-label by pnl (meta-labeling)
    """
    #1) prices aligned with events
    events_=events.dropna(subset=['t1'])
    px=events_.index.union(events_['t1'].values).drop_duplicates()
    px=close.reindex(px,method='bfill')
    #2) create out object
    out=pd.DataFrame(index=events_.index)
    out['ret']=px.loc[events_['t1'].values].values/px.loc[events_.index]-1
    if 'side' in events_:out['ret']*=events_['side'] # meta-labeling
    out['bin']=np.sign(out['ret'])
    if 'side' in events_:out.loc[out['ret']<=0,'bin']=0 # meta-labeling
    return out

def dropLabels(events,
               minPct=.05):
    """SNIPPET 3.8 DROPPING UNDER-POPULATED LABELS
    Apply weights, drop labels with insufficient examples
    """
    while True:
        df0=events['bin'].value_counts(normalize=True)
        if df0.min()>minPct or df0.shape[0]<3:break
        print('dropped label: ', df0.argmin(),df0.min())
        events=events[events['bin']!=df0.argmin()]
    return events


SNIPPET 3.1 DAILY VOLATILITY ESTIMATES     Daily vol reindexed to close # daily rets SNIPPET 3.2 TRIPLE-BARRIER LABELING METHOD     Apply stop loss/profit taking, if it takes place before t1 (end of event) # NaNs # NaNs # path prices # path returns # earliest stop loss # earliest profit taking SNIPPET 3.3 GETTING THE TIME OF FIRST TOUCH     SNIPPET 3.6 EXPANDING getEvents TO INCORPORATEMETA-LABELING #1) get target # minRet #2) get t1 (max holding period) #3) form events object, apply stop loss on t1 # pd.min ignores nan SNIPPET 3.4 ADDING A VERTICAL BARRIER SNIPPET 3.5 LABELING FOR SIDE AND SIZE #1) prices aligned with events #2) create out object # where out index and t1 (vertical barrier) intersect label 0 SNIPPET 3.7 EXPANDING getBins TO INCORPORATE META-LABELING     Compute event's outcome (including side information, if provided).     events is a DataFrame where:     -events.index is event's starttime     -events['t1'] is event's endtime     -events['trgt'] is event's target     -events['side'] (optional) implies the algo's position side     Case 1: ('side' not in events): bin in (-1,1) <-label by price action     Case 2: ('side' in events): bin in (0,1) <-label by pnl (meta-labeling) #1) prices aligned with events #2) create out object # meta-labeling # meta-labeling SNIPPET 3.8 DROPPING UNDER-POPULATED LABELS     Apply weights, drop labels with insufficient examples