""" Shareholder Yield Algo.
This is for learning the QuantConnect Algo Framework. I don't expect it to be useful outside of that.

The goal of this algo is:
- create a universe of liquid stocks with the best share holder yield (dividends + buybacks) relative to the market.
- buy an equal weighted portfolio of these stocks.
- rebalance monthly based on share holder yield.

Classes:
- ShareHolderYieldAlgo() is the main class that configures and instantiates the QuantConnect algo framework.
- LiquidTotalYieldUniverseSelectionModel() is a class which creates a liquid universe of stocks ranking by
  dollar volume, then filters the list down for the top X yielders based on share holder yield.
- LongShareHolderYieldAlphaModel() is an alpha model class which emits buy and sell signals based on changes
  to the universe selection.
"""
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class ShareHolderYieldAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(LiquidTotalYieldUniverseSelectionModel())
        self.AddAlpha(LongShareHolderYieldAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())

class LiquidTotalYieldUniverseSelectionModel(FundamentalUniverseSelectionModel):
    
    def __init__(self):
        super().__init__(True, None, None)
        self.lastMonth = -1
        # Use this to determined how many stocks we take off the dollar volume sorted list.
        self.dollarVolumeListSize = 200
        # Number of securities to select and return from universe.
        self.selectionSize = 10
        
    def SelectCoarse(self, algorithm, coarse):
        if self.lastMonth == algorithm.Time.month:
            return Universe.Unchanged
            
        self.lastMonth = algorithm.Time.month

        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData],
            key=lambda x: x.DollarVolume, reverse=True)

        return [x.Symbol for x in sortedByDollarVolume[:self.dollarVolumeListSize]]

    def SelectFine(self, algorithm, fine):
        sortedByYields = sorted(fine, key=lambda f: f.ValuationRatios.TotalYield, reverse=True)
        return [f.Symbol for f in sortedByYields][:self.selectionSize]
        
class LongShareHolderYieldAlphaModel(AlphaModel):

    def __init__(self):
        self.lastMonth = -1
        # Average number of trading days per month
        self.signalTimeDelta = timedelta(days = 21)
        self.symbols = []

    def Update(self, algorithm, data):
        
        insights = []
        
        # Skip if we have already run this month.
        if self.lastMonth == algorithm.Time.month:
            return insights
            
        self.lastMonth = algorithm.Time.month 
  
        for symbol in self.symbols:
            insights.append(Insight.Price(symbol, self.signalTimeDelta, InsightDirection.Up))
        
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.symbols:
                self.symbols.append(security.Symbol)
                
        for security in changes.RemovedSecurities:
            self.symbols.remove(security.Symbol)