import pandas as pd
import matplotlib.pyplot as plt

def engineered_momentum(end, start, sp500, world_small, treasuries):
    assets = [sp500, world_small, treasuries]
    df = pd.DataFrame(columns=['Ratio','Short','Long']) 
    data = get_pricing(assets, start_date=start, end_date=end)['close_price'].dropna()
    
    #Get best asset class within each subcategory
    df = pd.DataFrame(columns=['Weight','Score','1Year','6Mon','3Mon','1Mon']) 
    
    #Calculate Momentum Ratios
    for stock in assets:
        his = data[stock][-252:]
        df.loc[stock, '1Mon'] = his[-1] / his[-21] - 1
        df.loc[stock, '3Mon'] = his[-1] / his[-63] - 1
        df.loc[stock, '6Mon'] = his[-1] / his[-126] - 1
        df.loc[stock, '1Year'] = his[-1] / his[0] - 1

    #Check Term Trend is Positive
    df = df.astype(float)
    df['Score'] = df['1Mon'] + df['3Mon'] + df['6Mon']
    
    
    '''
    Determine Asset
    '''    
    df['Weight'] = 0.0
    df.loc[sp500, 'Weight'] = df.loc[sp500,'Score']
    df.loc[world_small, 'Weight'] = df.loc[world_small,'Score']
    df.loc[df['Weight'] < 0, 'Weight'] = 0.0
    df.loc[~df.index.isin(df['Weight'].nlargest(1).index.tolist()),'Weight'] = 0.0
    if len(df[df.Weight > 0]) == 0:
        df.loc[treasuries, 'Weight'] = 1.0
    df.loc[df.Weight > 0, 'Weight'] = 1.0
    
    return df, data

"""
Define end date either by now, or offsetting etc.
"""
end = pd.Timestamp.utcnow() #pd.to_datetime('2016-11-01') # #2017-04-01
print end
start = end - 300 * pd.tseries.offsets.BDay()

#Run strategy
df, data = engineered_momentum(end,
                               start,
                               symbols('SPY'),
                         symbols('VSS'),
                         symbols('VGLT'))

#Plot and display results
returns = data.dropna().pct_change().dropna()+1
fig, ax = plt.subplots()
returns[-126:].cumprod().plot(ax = ax)
df