import quantopian.algorithm as algo
from quantopian.pipeline import Pipeline
from quantopian.pipeline.data.builtin import USEquityPricing
from quantopian.pipeline.filters import Q500US


def initialize(context):
    #Schedule our rebalance function to run once a month
    schedule_function(rebalance , date_rules.month_end(), time_rules.market_close(minutes=10))
    schedule_function(stops, date_rules.every_day(), time_rules.market_close(minutes=14))

    #This code simply "attaches" our pipeline to this algo. This code never changes
    algo.attach_pipeline(make_pipeline(), 'pipeline')

    #Grab SPY for the Trend Following Filter
    context.spy = sid(8554)

    # Specify percentage based intial protective stop and trailing stop (TS) levels

    context.IPS = 0.8
    context.TS = 0.8

    # Dictionary to hold trailing stop prices {stock1:TS1, etc...}
    context.TS_prices={}


def make_pipeline():

    # Base universe set to the Q500US
    universe = Q500US()

    # This is pipeline code
    # All we are doing is making an object, 'pipe' containing the current 500 most liquid US stocks
    pipe = Pipeline(screen=universe)
    return pipe

def before_trading_start(context, data):

    # Called every day before market open, this grabs the output from our function "make_pipeline"
    context.output = algo.pipeline_output('pipeline')

    # Here we grab in "index" of our pipeline
    # This makes a list of securities "context.security_list" containing 500 stocks
    context.security_list = context.output.index

def rebalance(context, data):
    #This code is for our Trend Following Regime Filter:
    #We get the history for SPY, current price for SPY and calculates the moving average
    spy_history = data.history(context.spy,"close", 200, "1d")
    spy_MA = spy_history.mean()
    spy_current = data.current(context.spy,'price')

    #This code grabs the pricing history for all 500 stocks in our universe
    #This DataFrame, "history", contains 253 rows (days) and 500 columns (one for each stock)
    history = data.history(context.security_list,"close", 253, "1d")

    #This grabs the current momentum for each stock
    #Default is 105 days (5 months) lookback
    #We use .iloc[-1] to get the last value (most recent momentum measure)
    current_momentum = history.pct_change(105).iloc[-1]

    #Here we use "nlargest" to grab the 20 stocks with the best momentum
    top_n_stock = current_momentum.nlargest(20)

    #Loop through all stocks in our portfolio, selling any that arent still in the top 20 based on momentum
    for security in context.portfolio.positions:
        if security not in top_n_stock:
            order_target_percent(security, 0.0)

    #Here we buy new stocks
    #We first check if the current price of SPY is abve its 200 day MA
    #We then loop through our "top_n_stock" list (20 with best momentum) and buy each we 5% of our portfolio
    if spy_current > spy_MA:
        for x in top_n_stock.index:
            if x not in context.portfolio.positions:
                order_target_percent(x , 0.05)

def stops(context, data):

    for stock in context.portfolio.positions:
        current_price = data.current(stock, 'price')
        current_position = context.portfolio.positions[stock].amount
        cost_basis = context.portfolio.positions[stock].cost_basis

    # initial protective stop (IPS)
        IPS_price = context.IPS * cost_basis

    # trailing stop  (TS)
        context.TS_price[stock] = max(context.TS_prices[stock] if stock in context.TS_prices
   else 0,context.TS * current_price)

        # Apply stops
        if (current_position > 0) and ((current_price < context.TS_prices[stock]) or (current_price < IPS_price)):
            order_target_percent(stock, 0)
            del context.TS_prices[stock]
            print('Stopped out of {} at price = {}, cost_basis = {}'.format(stock, round(current_price,4), round(cost_basis,4)))