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.data import morningstar 
from quantopian.pipeline.factors import AverageDollarVolume, SimpleMovingAverage
from quantopian.pipeline.filters.morningstar import IsPrimaryShare


import numpy as np
from collections import defaultdict

#
# define custom classes
#
class simple_momentum(CustomFactor):    
   inputs = [USEquityPricing.close]   
   window_length = 1  
     
   def compute(self, today, assets, out, close):  
     out[:] = close[-1]/close[0]
        

        
        
class market_cap(CustomFactor):    
   inputs = [USEquityPricing.close, morningstar.valuation.shares_outstanding]   
   window_length = 1  
     
   def compute(self, today, assets, out, close, shares):      
     out[:] = close[-1] * shares[-1]          
        
        
class get_fcf_per_share(CustomFactor):    
   inputs = [morningstar.valuation_ratios.cash_return]   
   window_length = 1  
     
   def compute(self, today, assets, out, cash_return):   
     out[:] = cash_return

        
class get_last_close(CustomFactor):    
   inputs = [USEquityPricing.close]   
   window_length = 1  
     
   def compute(self, today, assets, out, close):  
     out[:] = close[-1]
        
def initialize(context):  

#: Declares which stocks we currently held and how many days we've held them dict[stock:days_held]
      
    
    schedule_function(cancel_open_orders, date_rules.every_day(),
                      time_rules.market_close())


    
#
# schedule methods
#
    schedule_function(func=periodic_rebalance,
                      date_rule=date_rules.month_start(days_offset=5),
                      time_rule=time_rules.market_open(), half_days=True)  
#
# set portfolis parameters
#
    set_do_not_order_list(security_lists.leveraged_etf_list)
    context.acc_leverage = 1.0
    
#
# Set commission model to be used
#
    set_commission(commission.PerShare(cost=0.005, min_trade_cost=1.00))
#
#
# 
    
#
# Establish pipeline
#
    pipe = Pipeline()  
    attach_pipeline(pipe, 'ranked_stocks')  
#
# Define the five momentum factors used in ranking stocks
#

# Week
    week = simple_momentum(window_length=5)
    pipe.add(week, 'week')      
# Month    
    month = simple_momentum(window_length=20)
    pipe.add(month, 'month')  
# Quarter    
    quarter = simple_momentum(window_length=60)
    pipe.add(quarter, 'quarter')  
# Six Months   
    six_months = simple_momentum(window_length=125)  
    pipe.add(six_months, 'six_months')  
# 52 Weeks     
    one_year = simple_momentum(window_length=252) 
    pipe.add(one_year, 'one_year') 
    
    
  
#
# Define other factors that may be used in stock screening
#
    free_cash_flow = get_fcf_per_share()
    pipe.add(free_cash_flow, 'fcf')
    mkt_cap = market_cap()
    pipe.add(mkt_cap, 'mkt_cap')
    factor8 = get_last_close()
    pipe.add(factor8, 'factor_8')
    

    
    ytd_filter = six_months > one_year   # only consider stocks with six month performance greater than one year
    mkt_cap_filter = mkt_cap > 2000000000   # only consider stocks mkt cap above 2b
    free_cash_flow_filter = free_cash_flow > 0 # only consider stocks with positive free cash flow
    close_price_filter = factor8 > 5.00  # only consider stocks that close above this value

    common_stock = morningstar.share_class_reference.security_type.latest.eq('ST00000001')
    not_lp_name = ~morningstar.company_reference.standard_name.latest.matches('.* L[\\. ]?P\.?$')
    not_lp_balance_sheet = morningstar.balance_sheet.limited_partnership.latest.isnull()
    have_data = morningstar.valuation.market_cap.latest.notnull()
    not_otc = ~morningstar.share_class_reference.exchange_id.latest.startswith('OTC')
    not_wi = ~morningstar.share_class_reference.symbol.latest.endswith('.WI')
    not_depository = ~morningstar.share_class_reference.is_depositary_receipt.latest
    primary_share = IsPrimaryShare()    
    
    
    
#
# Establish screen used to establish candidate stock list
#
# Establish screen to filter unwanted 
   
    total_filter = (common_stock
                    & not_lp_name
                    & not_lp_balance_sheet
                    & have_data
                    & not_otc
                    & not_wi
                    & not_depository
                    & primary_share
                    & ytd_filter 
                    & mkt_cap_filter 
                    & free_cash_flow_filter
                    & close_price_filter
                   )
                              

    pipe.set_screen(total_filter)
    
    
def before_trading_start(context,data):

    
    context.holdings_bottom_large_cap = 17
  
    context.holdings_bottom_mid_cap = 8

    context.holdings_bottom_small_cap = 6
    
    context.output = pipeline_output('ranked_stocks')
    
  

    
    ranked_large_cap_stocks = context.output[context.output.mkt_cap > 25000000000]
    ranked_mid_cap_stocks = context.output[context.output.mkt_cap > 10000000000] 
    ranked_small_cap_stocks = context.output[context.output.mkt_cap > 5000000000]
    
#
# Establishing ranking based on different market caps Large, Mid and Small. Sorting by six months return.
#

#
# Large Cap
#

 
    
    context.stock_factors_bottom_large_cap = ranked_large_cap_stocks.sort(['six_months'], ascending=True).iloc[:context.holdings_bottom_large_cap]
    
    
#
# Mid cap
#
        

    
    
    context.stock_factors_bottom_mid_cap = ranked_mid_cap_stocks.sort(['six_months'], ascending=True).iloc[:context.holdings_bottom_mid_cap]
 
    
#
# Small Cap
#
                                                                                                           


    
    context.stock_factors_bottom_small_cap = ranked_small_cap_stocks.sort(['six_months'], ascending=True).iloc[:context.holdings_bottom_small_cap]
                                                                                                      
                                                                                                           
                                                                                                          
                                                                                                          

    context.stock_list_bottom_large_cap_set = set(context.stock_factors_bottom_large_cap.index)
    context.stock_list_large_cap_set = context.stock_list_bottom_large_cap_set
    

    context.stock_list_bottom_mid_cap_set = set(context.stock_factors_bottom_mid_cap.index)
    context.stock_list_mid_cap_set = context.stock_list_bottom_mid_cap_set
    
    
    context.stock_list_bottom_small_cap_set = set(context.stock_factors_bottom_small_cap.index)
    context.stock_list_small_cap_set = context.stock_list_bottom_small_cap_set
    
  
    context.stock_list_large_mid_set = context.stock_list_large_cap_set.union(context.stock_list_mid_cap_set)
    context.stock_list_set = context.stock_list_small_cap_set.union(context.stock_list_large_mid_set)
    

    
    
    
def cancel_open_orders(context, data):
    for stock in get_open_orders():
        for order in get_open_orders(stock):
            cancel_order(order)
        
  

    
def periodic_rebalance(context,data):


# Get any stocks we ordered last time that still have open orders.
    open_orders = get_open_orders()
    
    
    for stock in context.portfolio.positions:
        costBasis = context.portfolio.positions[stock].cost_basis
        if data.can_trade(stock):
            if 'price' < costBasis:
                order_target(stock, 0)
    
    
            
    long_weight = context.acc_leverage / len(context.stock_list_set)
    for stock in context.stock_list_set:
        if stock not in open_orders:
            if data.can_trade(stock):
                order_target_percent(stock, long_weight)

    
    
    n10 = len(context.stock_list_set)
    record(leverage=context.account.leverage, 
           positions= len(context.portfolio.positions),
    candidates=n10)