from quantopian.algorithm import attach_pipeline, pipeline_output
from quantopian.pipeline import Pipeline
from quantopian.pipeline import CustomFactor
from quantopian.pipeline.data.builtin import USEquityPricing
from quantopian.pipeline.factors import SimpleMovingAverage
from quantopian.pipeline.data import morningstar
import numpy as np

def offset_sma_factor(days, **kwargs):
    kwargs['window_length'] += days
    class OffsetSMA(CustomFactor):
        def compute(self, today, assets, out, values):
            out[:] = np.nanmean(values[:-days], axis=0)
    return OffsetSMA(**kwargs)

def initialize(context):
    context.security = symbol('SPY')
    pipe = Pipeline()
    pipe = attach_pipeline(pipe, name='co_pipeline')
    sma_50 = SimpleMovingAverage(inputs=[USEquityPricing.close], window_length=50)
    sma_200 = SimpleMovingAverage(inputs=[USEquityPricing.close], window_length=200)
    sma_50_1day_ago = offset_sma_factor(1, inputs=[USEquityPricing.close], window_length=50)
    sma_200_1day_ago = offset_sma_factor(1, inputs=[USEquityPricing.close], window_length=200)

    current_co = (sma_50 > sma_200)
    past_not_co = (sma_50_1day_ago < sma_200_1day_ago);

    pipe.add(sma_50, 'sma_50')
    pipe.add(sma_200, 'sma_200')
    pipe.add(sma_50_1day_ago, 'sma_50_1day_ago')
    pipe.add(sma_200_1day_ago, 'sma_200_1day_ago')
    pipe.set_screen(current_co & past_not_co)

def before_trading_start(context, data):
    # get the output of your pipeline
    context.output = pipeline_output('co_pipeline').fillna(0)

    context.etf_list = [sid(24), sid(5061)]

    list = context.output[context.output.index.isin(context.etf_list)]
    context.stock_list = list.sort(ascending=True).iloc[:2]
    context.assets = [context.stock_list]

def handle_data(context, data):

    co_output = context.assets
    log.info(co_output)

    average_price = data[context.security].mavg(5)
    current_price = data[context.security].price

    cash = context.portfolio.cash

    if current_price > 1.01*average_price and cash > current_price:

        # Need to calculate how many shares we can buy
        number_of_shares = int(cash/current_price)

        # Place the buy order (positive means buy, negative means sell)
        order(context.security, +number_of_shares)
        log.info("Buying %s" % (context.security.symbol))


    elif current_price < average_price:

        # Sell all of our shares by setting the target position to zero
        order_target(context.security, 0)
        log.info("Selling %s" % (context.security.symbol))

    # You can use the record() method to track any custom signal.
    # The record graph tracks up to five different variables.
    # Here we record the Apple stock price.
    record(stock_price=data[context.security].price)