"""
Trading Strategy using Fundamental Data
1. Look at stocks in the QTradableStocksUS.
2. Go long in the top 50 stocks by pe_ratio.
3. Go short in the bottom 50 stocks by pe_ratio.
4. Rebalance once a month at market open.
"""

from quantopian.algorithm import attach_pipeline, pipeline_output
from quantopian.pipeline import Pipeline
from quantopian.pipeline.data import Fundamentals
from quantopian.pipeline.filters import QTradableStocksUS
from quantopian.pipeline.factors import BusinessDaysSincePreviousEvent


import pandas as pd


def initialize(context):
    # Rebalance monthly on the first day of the month at market open
    schedule_function(rebalance,
                      date_rule=date_rules.month_start(days_offset=1),
                      time_rule=time_rules.market_open())

    attach_pipeline(make_pipeline(), 'fundamentals_pipeline')

def make_pipeline():

    growth_grade = Fundamentals.growth_grade.latest
    
   
    universe = QTradableStocksUS()
    Aplus_grade = growth_grade.startswith('A')
    Anormal_grade = growth_grade.startswith('A') # This seems to be the same thing as Aplus_grade?
    F_grade = growth_grade.startswith('F')
   

    graded  = (F_grade | Aplus_grade | Anormal_grade )
    A_grade = (Anormal_grade | Aplus_grade)
    tradable = universe & graded
    
    pipe = Pipeline(
              columns={
                'Growth Grade': growth_grade,
                'long': A_grade,
                'short': F_grade 
              }, screen= tradable
          )

    return pipe

    
"""
Runs our fundamentals pipeline before the marke opens every day.
"""
def before_trading_start(context, data): 

    context.raw_pipe_output = pipeline_output('fundamentals_pipeline')
    context.long_pipe_output = context.raw_pipe_output[context.raw_pipe_output['long']].head(100)
    context.short_pipe_output = context.raw_pipe_output[context.raw_pipe_output['short']].head(100)
    context.pipe_output = context.long_pipe_output.append(context.short_pipe_output)
    
    print context.pipe_output

    # The high p/e stocks that we want to long.
    context.longs = context.pipe_output[context.pipe_output['long']].index

    # The low p/e stocks that we want to short.
    context.shorts = context.pipe_output[context.pipe_output['short']].index
    
    record(leverage = context.account.leverage)
    
    longs = shorts = 0
    for position in context.portfolio.positions.itervalues():
        if position.amount > 0:
            longs += 1
        if position.amount < 0:
            shorts += 1
    record(long_count=longs, short_count=shorts)

    
    

def rebalance(context, data):
    
    my_positions = context.portfolio.positions
    
    # If we have at least one long and at least one short from our pipeline. Note that
    # we should only expect to have 0 of both if we start our first backtest mid-month
    # since our pipeline is scheduled to run at the start of the month.
    if (len(context.longs) > 0) and (len(context.shorts) > 0):

        # Equally weight all of our long positions and all of our short positions.
        long_weight = .75/len(context.longs)
        short_weight = -.25/len(context.shorts)
        
        # Get our target names for our long and short baskets. We can display these
        # later.
        target_long_symbols = [s.symbol for s in context.longs]
        target_short_symbols = [s.symbol for s in context.shorts]

        log.info("Opening long positions each worth %.2f of our portfolio in: %s" \
                 % (long_weight, ','.join(target_long_symbols)))
        
        log.info("Opening long positions each worth %.2f of our portfolio in: %s" \
                 % (short_weight, ','.join(target_short_symbols)))
        
        # Open long positions in our high p/e stocks.
        for security in context.longs:
            if data.can_trade(security):
                if security not in my_positions:
                    order_target_percent(security, long_weight)
            else:
                log.info("Didn't open long position in %s" % security)

        # Open short positions in our low p/e stocks.
        for security in context.shorts:
            if data.can_trade(security):
                if security not in my_positions:
                    order_target_percent(security, short_weight)
            else:
                log.info("Didn't open short position in %s" % security)
                  

    closed_positions = []
    
    # Close our previous positions that are no longer in our pipeline.
    for security in my_positions:
        if security not in context.longs and security not in context.shorts \
        and data.can_trade(security):
            order_target_percent(security, 0)
            closed_positions.append(security)
    
    log.info("Closing our positions in %s." % ','.join([s.symbol for s in closed_positions]))