from QuantConnect.Data.UniverseSelection import *
import operator
from math import ceil,floor


class CoarseFineFundamentalComboAlgorithm(QCAlgorithm):
    def Initialize(self):

        self.SetStartDate(2012, 1, 1)  #Set Start Date
        self.SetEndDate(2018, 10, 11)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        self.flag1 = 1
        self.flag2 = 0
        self.flag3 = 0
    

        self.UniverseSettings.Resolution = Resolution.Daily        
        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.AddEquity("SPY")
        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 50
        self.num_portfolios = 5
 
        self._changes = None
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), Action(self.Rebalancing))


    def CoarseSelectionFunction(self, coarse):
        if self.flag1 == 2:
            CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
            sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True) 
            top = sortedByDollarVolume[:self.__numberOfSymbols]
            return [x.Symbol for x in top]
        else:
            return [x.Symbol for x in self.topFine]


    def FineSelectionFunction(self, fine):
        if self.flag1 == 2:
            self.flag1 = 0
            self.flag2 = 1
        
            filtered_fine = [x for x in fine if x.OperationRatios.OperationMargin.Value
                                            and x.CompanyReference.IndustryTemplateCode == "N"
                                            and x.ValuationRatios.PriceChange1M 
                                            and x.ValuationRatios.PERatio]
    
            sortedByfactor1 = sorted(filtered_fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=True)
            sortedByfactor2 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PriceChange1M, reverse=True)
            sortedByfactor3 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
            
            num_stocks = floor(len(filtered_fine)/self.num_portfolios)
            
            self.Log(str(num_stocks))


            stock_dict = {}
            
            for i,ele in enumerate(sortedByfactor1):
                rank1 = i
                rank2 = sortedByfactor2.index(ele)
                rank3 = sortedByfactor3.index(ele)
                score = sum([rank1*0.2,rank2*0.4,rank3*0.4])
                stock_dict[ele] = score
            
            self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
            sorted_symbol = [self.sorted_stock[i][0] for i in range(len(self.sorted_stock))]
            self.topFine = sorted_symbol[:self.__numberOfSymbolsFine]
                
            self.flag3 += 1            
            
            return [x.Symbol for x in self.topFine]
        else:
            return [x.Symbol for x in self.topFine]


    def OnData(self, data):
        if self.flag3 > 0:
            if self.flag2 == 1:
                self.flag2 = 0

                # if we have no changes, do nothing
                if self._changes == None: return
                # liquidate removed securities
                for security in self._changes.RemovedSecurities:
                    if security.Invested:
                        self.Liquidate(security.Symbol)
                 
                for security in self._changes.AddedSecurities:
                    self.SetHoldings(security.Symbol, 1/float(len(self._changes.AddedSecurities)))    
         
                self._changes = None

    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        self._changes = changes
    
    def Rebalancing(self):
        self.flag1 += 1