from Execution.ImmediateExecutionModel import ImmediateExecutionModel

class VerticalHorizontalSplitter(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 12, 3)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash

        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 5
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddAlpha(MyAlphaModel())

        self.SetPortfolioConstruction(MarketCapWeightedPortfolioConstructionModel(self))
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.market_cap_by_symbol = {}


    def OnData(self, data):
        pass


    def CoarseSelectionFunction(self, coarse):
        filtered = [c for c in coarse if c.HasFundamentalData]
        sortedByDollarVolume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]
    

    def FineSelectionFunction(self, fine):
        sortedByPeRatio = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        self.market_cap_by_symbol.clear()
        symbols = []
        for f in sortedByPeRatio[:self.__numberOfSymbolsFine]:
            symbols.append(f.Symbol)
            self.market_cap_by_symbol[f.Symbol] = f.MarketCap
        return symbols
        
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            self.market_cap_by_symbol.pop(security.Symbol, None)


class MyAlphaModel(AlphaModel):
    
    symbols = set()
    
    def Update(self, algorithm, data):
        insights = []
        for symbol in self.symbols:
            if symbol in data.Bars:
                insights.append(Insight(symbol, timedelta(days=1), InsightType.Price, InsightDirection.Up))
        return Insight.Group(insights)
    
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            self.symbols.add(security.Symbol)
        
        for security in changes.RemovedSecurities:
            self.symbols.remove(security.Symbol)


class MarketCapWeightedPortfolioConstructionModel(PortfolioConstructionModel):
    
    symbols = set()
    
    def __init__(self, algorithm):
        self.algo = algorithm
    
    
    def CreateTargets(self, algorithm, insights):
        market_cap_sum = 0
        for market_cap in self.algo.market_cap_by_symbol.values():
            market_cap_sum += market_cap
        
        targets = []
        for symbol in self.symbols:
            pct = self.algo.market_cap_by_symbol[symbol] / market_cap_sum
            target = PortfolioTarget.Percent(algorithm, symbol, pct)
            targets.append(target)
        return targets
    
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            self.symbols.add(security.Symbol)
        
        for security in changes.RemovedSecurities:
            self.symbols.remove(security.Symbol)