from AlgorithmImports import *

class CalibratedOptimizedCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # Set QC500 Universe Selection Model
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        # Dictionary to hold SecurityData 
        self.data = {}
    
    def OnData(self, data):
      
        # Sorted by MOM
        sortedByMOM = sorted(list(self.data.keys()), key=lambda k : self.data[k].mom.Current.Value, reverse=True)
        # Select top 15 symbols
        rankings = sortedByMOM[:10]
        #[x for x in a if x > 2]
       # sma_shortlist = [x for rankings  if self.data[x].sma45.Current.Value > self.data[x].sma44.Current.Value  ]
        # Liquidate outdated symbols
        for kvp in self.Portfolio:
            symbol = kvp.Key
            holding = kvp.Value
            if self.data[symbol].sma45.Current.Value < self.data[symbol].sma44.Current.Value:
              if not self.Portfolio[symbol].Invested:
                self.Liquidate(symbol)
        # Set holdings to selected symbols
        for symbol in rankings:
            if self.data[symbol].sma45.Current.Value > self.data[symbol].sma44.Current.Value:
              if not self.Portfolio[symbol].Invested:
                self.SetHoldings(symbol, 1/len(self.Portfolio))
                
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.data:
                self.data[symbol] = SecurityData(self, symbol, security) 


class SecurityData:
    
    def __init__(self, algorithm, symbol, security):
        self.algorithm = algorithm
        self.symbol = symbol
        self.security = security
        # Retrieves ROIC fundamental
        self.mom = algorithm.MOMP(symbol, 42, Resolution.Daily)
        self.sma45 = algorithm.SMA(symbol,45,Resolution.Daily)
        self.sma44 = algorithm.SMA(symbol,44,Resolution.Daily)
        # Initialize MOM indicator with historical data
        history = algorithm.History(symbol,46, Resolution.Daily)
        
        if not history.empty:
            history = history.close.unstack(0)
            if not history.empty:
                df = history[symbol]
                for time, close in df.iteritems():
                    self.mom.Update(time, close)
                    self.sma45.Update(time,close)
                    self.sma44.Update(time,close)