import numpy as np  
import pandas as pd

from datetime import timedelta

class QualityMomentumModel(QCAlgorithm):
   
    def Initialize(self):
        self.SetStartDate(2020, 2, 10)  # Set Start Date
        self.SetEndDate(2020, 2, 11)
        self.SetCash(1000)  # Set Strategy Cash
        self.SPY = self.AddEquity("SPY", Resolution.Minute) #add SPY to use for trends
        #list of bond etfs for when markets down.
        self.AddEquity("TLT").Symbol
        self.AddEquity("IEF").Symbol
        self.BONDS = ['TLT', 'IEF']
        # Add bonds 
        # Set target number of securities to hold and top ROE qty to filter  
        self.TARGET_SECURITIES = 5
        self.TOP_ROE_QTY = 50 #First sort by ROE
        self.UniverseSettings.Resolution = Resolution.Minute #update the universe every minute
        #adding a universe of stocks
        # self.AddUniverse(self.FineSelectionFunction)
        #determine how many symbols to select in the coarse filter
        self.num_coarse = 1000
        self.num_fine = 50
        
        #trend following filter
        self.TF_LOOKBACK = 200  
        self.TF_CURRENT_LOOKBACK = 20
        
        #determining momentum
        self.MOMENTUM_LOOKBACK_DAYS = 126 #how many days to lookback  
        self.MOMENTUM_SKIP_DAYS = 10  #how many days to skip
        overall_lookback = (self.MOMENTUM_LOOKBACK_DAYS + self.MOMENTUM_SKIP_DAYS)
        
        # Initialize any other variables before starting trading  
        # setting the weights for each type of securitie
        self.stock_weights = pd.Series()  
        self.bond_weights = pd.Series()
 
   
        # #schedule function for selecting stocks and weights
        # self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), self.select_stocks_set_weights)
        # self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 30), self.select_stocks_set_weights)
        # #schedule function for making trades
        # self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 30), self.trade)
        # self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 30), self.trade)

    
        # #50 Day moving average of SPY
        # self.spy_ma_fast = self.SMA("SPY", 50)
        # #200 Day moving average of SPY
        # self.spy_ma_slow = self.SMA("SPY", 200)
    
        # self.trend_up = self.spy_ma_fast >= self.spy_ma_slow
    
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
    
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data    
        '''
        
    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
    
        selectedcoarse = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 10],
            key=lambda x: x.DollarVolume, reverse=True)

        return [x.Symbol for x in selectedcoarse[:self.num_coarse]]


    def FineSelectionFunction(self, fine):
        filtered_fine = [x for x in fine if x.OperationRatios.ROIC.SixMonths
                                        and x.ValuationRatios.CashReturn
                                        and x.ValuationRatios.FCFYield
                                        and x.OperationRatios.LongTermDebtEquityRatio.NineMonths]
            
        value = [x for x in filtered_fine if x.ValuationRatios.CashReturn
                                         and x.ValuationRatios.FCFYield]
    
        quality = [x for x in value if x.OperationRatios.ROIC
                                    and x.OperationRatios.LongTermDebtEquityRatio]
    
        returns_overall = sorted(filtered_fine,
            key = lambda f: f.OperationRatios.RevenueGrowth.overall_lookback)
            
            
        returns_recent = sorted(filtered_fine,
            key = lambda f: f.OperationRatios.RevenueGrowth.MOMENTUM_SKIP_DAYS)
            
        momentum = sorted(filtered_fine,
            key = lambda f: f.returns_overall +
                            f.returns_recent)
                        
 
        top_quality = sorted(filtered_fine,
            key = lambda f: f.OperationRatios.ROE.OneMonth)
            
        return [x.Symbol for x in top_quality[:self.TOP_ROE_QTY]] #search for top 50 equities with the highest ROE
        
        # this won't get called
        return [x.Symbol for x in momentum[:self.TARGET_SECURITIES]] #search for specified number of securities with the highest momentum