from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil
class VentralDynamicContainmentField(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 12, 2)  # Set Start Date
        self.SetEndDate(2019, 1, 8)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = 500
        self.dollarVolumeBySymbol = {}
        self.symbols = []
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday),
            self.TimeRules.At(16, 00),
            self.SelectSymbols)
        self.select = False


    def OnData(self, data):
        if self.select:
            self.Log(self.Time)

    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
        if not self.select: return self.symbols
        
        filtered = [x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0]
        sortedByDollarVolume = sorted(filtered, key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.symbols.clear()
        self.dollarVolumeBySymbol.clear()
        for x in sortedByDollarVolume:
            self.symbols.append(x.Symbol)
            self.dollarVolumeBySymbol[x.Symbol] = x.DollarVolume

        # return the symbol objects our sorted collection
        return self.symbols
    
    # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
    def FineSelectionFunction(self, fine):
        if not self.select: return self.symbols
        
        filteredFine = [x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and (x.CompanyReference.PrimaryExchangeID == "NYS" or x.CompanyReference.PrimaryExchangeID == "NAS")
                                        and (self.Time - x.SecurityReference.IPODate).days > 180
                                        and x.EarningReports.BasicAverageShares.ThreeMonths * x.EarningReports.BasicEPS.TwelveMonths * x.ValuationRatios.PERatio > 5e8]

        sortedByDollarVolume = []
        sortedBySector = sorted(filteredFine, key = lambda x: x.CompanyReference.IndustryTemplateCode)

        percent = self.numberOfSymbolsFine/float(len(sortedBySector))

        # select stocks with top dollar volume in every single sector
        for code, g in groupby(sortedBySector, lambda x: x.CompanyReference.IndustryTemplateCode):
            y = sorted(g, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
            c = ceil(len(y) * percent)
            sortedByDollarVolume.extend(y[:c])

        sortedByDollarVolume = sorted(sortedByDollarVolume, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)
        self.symbols = [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]
        self.Log(f"selection complete: {self.Time}")
        self.select = False
        return self.symbols
        
    def SelectSymbols(self):
        self.select = True