from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class CalibratedOptimizedCompensator(QCAlgorithm):

    def Initialize(self):
        # self.SetStartDate(2019, 8, 16)  # Set Start Date
        # self.SetEndDate(2019, 9, 28)
        self.SetCash(100000)  # Set Strategy Cash
        # Set QC500 Universe Selection Model
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.UniverseSettings.Resolution = Resolution.Daily
        # Dictionary to hold SecurityData 
        self.data = {}
        self.SetWarmUp(365)
        self.lookback = 365
        self.num_long = 5
        self.SetBenchmark("SPY")

        self.rebalance = True
        self.last_trade = 0

    def OnData(self, data):
        self.last_trade += 1
        if self.last_trade == 30:
            self.last_trade = 0
            self.rebalance = True
            
        if not self.rebalance: return
    
        # Sorted by MOM
        sortedByMOM = sorted(list(self.data.keys()), key=lambda k : self.data[k].mom.Current.Value, reverse=True)
        # Select top self.num_long symbols
        selected = sortedByMOM[:self.num_long]
        # 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, 1/self.num_long)
        self.rebalance = False
                
    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
        self.mom = algorithm.MOM(symbol, 365, Resolution.Daily)
        
        # Initialize MOM indicator with historical data
        history = algorithm.History(symbol, 365, 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)