from dateutil.relativedelta import relativedelta

class PuppiesOfTheDow(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 1, 1)    # Set Start Date       
        self.SetCash(100000)           # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        self.coarseUniverseSize = 500
        self.portfolioSize = 5
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.selectUniverse = True
        self.monthCounter = 0
        self.lastRebalancing = self.Time - relativedelta(years=1)
        
        
    def CoarseSelectionFunction(self, coarse):
        if self.lastRebalancing.year != self.Time.year:
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
            self.filteredCoarse = [ x.Symbol for x in filtered[:self.coarseUniverseSize]]
            return self.filteredCoarse
        else: 
            return self.filteredCoarse
        
    def FineSelectionFunction(self, fine):
        if self.lastRebalancing.year != self.Time.year:
            dogsOfTheDow = sorted(fine, key = lambda x: x.ValuationRatios.TotalYield, reverse=True)
            puppiesOfTheDow = sorted(dogsOfTheDow, key = lambda x: x.ValuationRatios.BuyBackYield)[:self.portfolioSize]
            self.filteredFine = [x.Symbol for x in puppiesOfTheDow]
            return self.filteredFine
        else:
            return self.filteredFine

    def OnData(self, data):
        pass

    def OnSecuritiesChanged(self, changes):
        for shortStock in changes.RemovedSecurities: self.Liquidate(shortStock.Symbol)
        for longStock in changes.AddedSecurities: self.SetHoldings(longStock.Symbol,1/self.portfolioSize)