from quantopian.algorithm import attach_pipeline, pipeline_output
from quantopian.pipeline import Pipeline
from quantopian.pipeline.filters import Q500US
from quantopian.pipeline.factors import RSI

import pandas as pd
import numpy as np
import scipy.stats as stats


def make_pipeline(context):
    pipe = Pipeline()
    universe = Q500US()
    rsi = RSI(mask=universe, window_length = 15)
    pipe.add(rsi, "rsi")
    pipe.set_screen(universe)
    return pipe
   
def initialize(context):
        
    context.max_long_sec         = 10
    context.max_short_sec        = 10
    
    attach_pipeline(make_pipeline(context), 'factors')   
    
    schedule_function(rebalance,       date_rules.every_day(), time_rules.market_open())
    schedule_function(my_record_close, date_rules.every_day(), time_rules.market_close())
    
def my_record_close(context, data):
    record(leverage_close=context.account.leverage)    
    record(positions=len(context.portfolio.positions))

# Compute final rank and assign long and short baskets.
def before_trading_start(context, data):
    #
    # get and clean pipeline results
    #
    results = pipeline_output('factors')
    results = results.replace([np.inf, -np.inf], np.nan)
    results = results.dropna()

    print 'Basket of stocks %d' % len(results)

    context.longs  = results["rsi"].order(ascending=False).tail(context.max_long_sec)
    context.shorts = results["rsi"].order(ascending=False).head(context.max_short_sec)

    # equal weight
    context.longs[:]  = 1.0/len(context.longs.index)
    context.shorts[:] = 1.0/len(context.shorts.index)
    context.longs  /= 2
    context.shorts /= 2
        
    print 'longs  weighted (length %d, sum %f):\n' % (len(context.longs.index), context.longs.sum())
    print 'shorts weighted (length %d, sum %f):\n' % (len(context.shorts.index), context.shorts.sum())
       
def rebalance(context, data):
        
    for security in context.shorts.index:
        if get_open_orders(security):
            continue
        if data.can_trade(security):
            order_target_percent(security, -context.shorts[security])
            
    for security in context.longs.index:
        if get_open_orders(security):
            continue
        if data.can_trade(security):
            order_target_percent(security, context.longs[security])
            
    for security in context.portfolio.positions:
        if get_open_orders(security):
            continue
        if data.can_trade(security) and security not in (context.longs.index | context.shorts.index):
            order_target_percent(security, 0) #, style=LimitOrder(data.current(security, 'price')))