forLong = []
forShort = []
class FundStrat(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 6, 14)  
        self.SetEndDate(2021, 6, 14)  
        self.SetCash(100000)  
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(
            FineFundamentalUniverseSelectionModel(self.SelectCoarse, 
                                                    self.SelectFine))
            
        self.AddAlpha(MyAlphaModel())
        
        self.num_fine = 5
        
    def SelectCoarse(self, coarse):
        
        hasFundamentalData = [c for c in coarse if c.HasFundamentalData]
        filteredByPrice = [c.Symbol for c in hasFundamentalData if c.Price > 10]
        
        return filteredByPrice

    def SelectFine(self, fine):
        filter_sector = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Energy]
        sortedByFactor = sorted(filter_sector, 
                            key=lambda f: f.ValuationRatios.CashReturn, 
                            reverse=True)
        self.forLong = [f.Symbol for f in sortedByFactor[:self.num_fine]]
        self.forShort = [f.Symbol for f in sortedByFactor[-self.num_fine:]]
        return self.forShort + self.forLong
    
class SymbolData:
    def __init__(self,security):
        self.symbol = security.Symbol
        self.factor = security.Fundamentals.ValuationRatios.CashReturn
        
class MyAlphaModel:
        def __init__(self):
            self.dataBySymbol = {}
            
        def Update(self, algorithm, slice):
            insights = []
            for symbol in self.forLong:
                if symbol in slice.Bars:
                    insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Up))
            for symbol in self.forShort:
                if symbol in slice.Bars:
                    insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Down))
            
            return insights
    
        def OnSecuritiesChanged(self, algorithm, changes):
            self.changes = changes
            self.forLong = []
            self.forShort = []
            for security in self.changes.RemovedSecurities:
                if security.Symbol in self.dataBySymbol:
                    self.dataBySymbol.pop(security.Symbol)
                
            algorithm.Debug(str(len(self.changes.RemovedSecurities)+len(self.changes.AddedSecurities)))
            for security in self.changes.AddedSecurities:
               
               if security.Symbol not in self.dataBySymbol:
                   self.dataBySymbol[security.Symbol] = SymbolData(security)
            selected = list(self.dataBySymbol.values())
            sortedByFactor = sorted(selected,key = lambda x : x.factor)

            self.forLong = [x.symbol for x in sortedByFactor[:5]]
            self.forShort = [x.symbol for x in sortedByFactor[-5:]]