from clr import AddReference
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")

from QuantConnect import *
from QuantConnect.Indicators import *

class EMAMomentumUniverse(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2021, 8, 15)
        self.SetCash(25000)
        self.SetTimeZone("America/New_York")
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.Coarse, self.Fine) 
        self.averages = { }
        self.selected = []
        self.fine_universe = []
        self.SetWarmUp(200)
        self.ATRscan = { }

    def OnSecuritiesChanged(self, changes):
        for x in changes.AddedSecurities:
            if x.Symbol not in self.ATRscan:
                SPC_history = self.History(x.Symbol, 14, Resolution.Minute)
                self.ATRscan[x.Symbol] = SPC_Data(x.Symbol, SPC_history)
        for x in changes.RemovedSecurities:
            data = self.ATRscan.pop(x.Symbol,None)
            
        self.Debug(f"Added Securities: {[security.Symbol.Value for security in changes.AddedSecurities]} Time: {self.Time}")
        self.Debug(f"Removed Securities: {[security.Symbol.Value for security in changes.RemovedSecurities]} Time: {self.Time}")
        
            
    def Coarse(self, coarse): 

        self.selected = []

        coarse = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)  
        self.selected = [c.Symbol for c in coarse if (c.HasFundamentalData 
                                        and c.Price > 100 and c.Price < 1000)][:100]
        
        return self.selected
        
    
    def Fine(self, fine):

        sortedfine = sorted(fine, key=lambda x: x.MarketCap)
        self.fine_universe = [x.Symbol for x in sortedfine if x.MarketCap > 700000000000 ][:5]
        
        return self.fine_universe
        
        
    def OnData(self, data):
        if self.IsWarmingUp:
            return

        
            
        for symbol, spc_data in self.ATRscan.items():
            if data.Bars.ContainsKey(symbol):
                spc_data.update(data.Bars[symbol])
            
            if spc_data.ATR14.IsReady: 
                ATR = spc_data.ATR14.Current.Value
                debug_symbol = symbol
                x = 'stop'
            else:
                self.Debug("ATR not ready")

            # Self.Debug(f"Symbol: {symbol} ATR: {ATR} Time: {self.Time}")
            

class SPC_Data:
    def __init__(self, symbol, SPC_history):
        self.symbol = symbol
        self.ATR14 = AverageTrueRange(14, MovingAverageType.Simple)
        # self.ATR14 = NormalizedAverageTrueRange(14)
        
        for bar in SPC_history.itertuples():
            tbar = TradeBar(bar.Index[1], self.symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
            self.ATR14.Update(tbar)

    def is_ready(self):
        return self.ATR14.IsReady
    
    def update(self, bar):
        self.ATR14.Update(bar)