from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class CalibratedOptimizedCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 8, 16)  # 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):
        # Sort by ROIC
        sortedByROIC = sorted(list(self.data.keys()), key=lambda k : self.data[k].roic, reverse=True)
        # Sorted by MOM
        sortedByMOM = sorted(sortedByROIC[:100], key=lambda k : self.data[k].mom.Current.Value, reverse=True)
        # Select top 25 symbols
        selected = sortedByMOM[:25]
        # Liquidate outdated symbols
        for kvp in self.Portfolio:
            symbol = kvp.Key
            holding = kvp.Value
            if symbol not in selected and holding.Invested:
                self.Liquidate(symbol)
        # Set holdings to selected symbols
        for symbol in selected:
            if not self.Portfolio[symbol].Invested:
                self.SetHoldings(symbol, 0.04)
                
    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.roic = security.Fundamentals.OperationRatios.ROIC.OneYear
        self.mom = algorithm.MOM(symbol, 14, Resolution.Daily)
        
        # Initialize MOM indicator with historical data
        history = algorithm.History(symbol, 14, 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)