from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class DeltaMomentumUniverseSelectionModel(FundamentalUniverseSelectionModel):
    symboldict = {}

    def __init__(self):
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        topDollarVolume = sorted(coarse, key=lambda k : k.DollarVolume, reverse=True)[:10]

        # When we get new symbols, we add them to the dict and warm up the indicator
        symbols = [x.Symbol for x in topDollarVolume if x.Symbol not in self.symboldict]
        history = algorithm.History(symbols, 136, Resolution.Daily)
        if not history.empty:
            history = history.close.unstack(0)
            for symbol in symbols:
                if str(symbol) not in history:
                    continue

                df = history[symbol].dropna()
                if not df.empty:
                    self.symboldict[symbol] = SymbolData(self, df)

        # Now, we update the dictionary with the latest data
        for x in coarse:
            symbol = x.Symbol
            if symbol in self.symboldict:
                self.symboldict[symbol].Update(x.EndTime, x.AdjustedPrice)

        topMOM = sorted(self.symboldict.items(), key=lambda x: x[1].DeltaMOM, reverse=True)[:200]
        return [x[0] for x in topMOM]
        

class SymbolData:
    def __init__(self, symbol, history):
        self.mom10 = Momentum(10)
        self.mom136 = Momentum(136)

        for time, close in history.iteritems():
            self.Update(time, close)

    def Update(self, time, close):
        self.mom10.Update(time, close)
        self.mom136.Update(time, close)


    @property
    def DeltaMOM(self):
        return self.mom10.Current.Value - self.mom136.Current.Value
    
    def __repr__(self):
        return f'{self.DeltaMOM}'

from universe import DeltaMomentumUniverseSelectionModel

class DynamicVentralCoil(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        # Set universe resolution
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Set universe selection and alpha
        self.AddUniverseSelection(DeltaMomentumUniverseSelectionModel())

    def OnSecuritiesChanged(self, changes):
        
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)

        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.005)