import numpy as np
import pandas as pd
class UniverseRollingAlgorithm(QCAlgorithm):
    def Initialize(self):  
        self.SetTimeZone(TimeZones.NewYork)   
        self.SetStartDate(2007, 1, 1)   
        self.SetEndDate(2020, 6,10)      
        self.SetCash(25000)            
        self.SetSecurityInitializer(lambda x: x.SetFeeModel(CustomFeeModel()))
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.SplitAdjusted
        self.AddEquity('SPY', Resolution.Minute).Symbol
        self.Portfolio.MarginCallModel = MarginCallModel.Null
        self.__numberOfSymbols     = 300
        #self.UniverseSettings.Leverage = 1.9
        for i in range(1, 388, 1): 
            self.Schedule.On(self.DateRules.EveryDay(), \
            self.TimeRules.AfterMarketOpen("SPY", i), \
            self.auto)
    def CoarseSelectionFilter(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True) 
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols]] 
    def FineSelectionFilter(self, fine):  # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
        stocks = fine
        stocks = [x for x in stocks if x.MarketCap>20000000]
        stocks = [x for x in stocks if x.CompanyReference.CountryId == "USA"]
        #stocks = [x for x in self.get_top_zscore(stocks)]
        stocks = [x for x in stocks if x.AssetClassification.GrowthGrade in ["A", "B"]]
        stocks = sorted(stocks, key=lambda x: x.ValuationRatios.PERatio>30)
        stocks = [x for x in stocks if x.SecurityReference.IsPrimaryShare == 1]
        #stocks = [x for x in stocks if x.Symbol.Value != "UNH"]
        #stocks = [x for x in stocks if x.Symbol.Value != "GOOG"]
        stocks = sorted(stocks, key=lambda x: x.MarketCap, reverse=True)[:20] 
        stocks = sorted(stocks, key=lambda x: x.AssetClassification.GrowthScore, reverse=True)[:10] 
        stocks = sorted(stocks, key=lambda x: x.MarketCap, reverse=True)[:5]
        self.universe = [x.Symbol for x in stocks]
        return self.universe
    def get_top_zscore(self, stocks):
        gs = [[x.AssetClassification.GrowthScore, x] for x in stocks if x.AssetClassification.GrowthScore != 0.0]
        a = np.array([x[0] for x in gs])
        mean, std = a.mean(), a.std()
        gs_zscore = [[(x[0] - mean) / std, x[1]] for x in gs]
        gs_zscore = sorted(gs_zscore, key=lambda x: x[0])
        gs_zscore = gs_zscore[round(len(gs_zscore) * 0.3):]
        return [x[1] for x in gs_zscore]
    def OnData(self, data):
        pass
    def auto(self):
            positions = [sec.Symbol for sec in self.Portfolio.Values if self.Portfolio[sec.Symbol].Invested]
            if len(positions)>=5:
                for symbol in positions:
                    if symbol not in self.universe:
                        self.SetHoldings(symbol, 0,False,"OUT")
            elif len(positions)<=4:
                for symbol in self.universe:
                    if not self.Portfolio[symbol].Invested:
                        self.SetHoldings(symbol, 0.199,False,"Take Long Position")
class CustomFeeModel:
    def GetOrderFee(self, parameters):
        fee = 0
        return OrderFee(CashAmount(fee, 'USD'))