#https://quantpedia.com/Screener/Details/36
from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
import scipy as sp

class NetPayoutYield(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2000, 6, 1)   
        self.SetEndDate(datetime.now())          
        self.SetCash(1000000)            
        self.UniverseSettings.Resolution = Resolution.Daily
        self.sorted_by_npy = None
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.AddEquity("SPY", Resolution.Daily)
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), self.rebalance)
        # Count the number of months that have passed since the algorithm starts
        self.months = -1
        self.yearly_rebalance = True
    def CoarseSelectionFunction(self, coarse):
        if self.yearly_rebalance:
            # drop stocks which have no fundamental data or have low price
            self.filtered_coarse = [x.Symbol for x in coarse if (x.HasFundamentalData)]
            return self.filtered_coarse
        else: 
            return []       

    def FineSelectionFunction(self, fine):
        if self.yearly_rebalance:
            # Filter stocks with nonzero Total Assets
            fine = [x for x in fine if (x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths != 0)] 
            # Calculate the market cap and add the "MakretCap" property to fine universe object
            for i in fine:
                i.MarketCap = float(i.EarningReports.BasicAverageShares.ThreeMonths * (i.EarningReports.BasicEPS.TwelveMonths*i.ValuationRatios.PERatio))
            fine = [x for x in fine if (x.MarketCap != 0)] 
            # sorted stocks in the top market-cap list by Net Payout Yield
            top_npy = sorted(fine, key = lambda x: ((x.ValuationRatios.TotalYield * x.MarketCap)-x.FinancialStatements.CashFlowStatement.CommonStockIssuance.TwelveMonths)/ x.MarketCap)
            self.sorted_by_npy = [i.Symbol for i in top_npy] 
            return self.sorted_by_npy
        else:
            return []

    def rebalance(self):
        # yearly rebalance
        self.months += 1
        if self.months%12 == 0:
            self.yearly_rebalance = True


    def OnData(self, data):
         if not self.yearly_rebalance: return 
         if self.sorted_by_npy:
            portfolio_size = int(len(self.sorted_by_npy)/10)
            #pick the upper decile to long
            long_stocks = self.sorted_by_npy[-portfolio_size:]
            stocks_invested = [x.Key for x in self.Portfolio]
            for i in stocks_invested:
                #liquidate the stocks not in the filtered 
                if i not in self.sorted_by_npy:
                    self.Liquidate(i) 
                #long the stocks in the list
                elif i in long_stocks:
                    self.SetHoldings(i, 1/(portfolio_size))
            self.yearly_rebalance = False