class TheSqueeze(AlphaModel):
    
    def __init__(self, period = 20, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolData = {}
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)
        
    def Update(self, algorithm, data):
        
        insights = []
        for key, sd in self.symbolData.items():
            if sd.bollinger.IsReady and \
                sd.keltner.IsReady and \
                sd.momentum.IsReady and \
                sd._bollinger["UpperBand"].IsReady and \
                sd._bollinger["LowerBand"].IsReady and \
                sd._keltner["UpperBand"].IsReady and \
                sd._keltner["LowerBand"].IsReady and \
                sd.momWindow.IsReady: 
        
                if algorithm.Portfolio[key].Invested: continue
            
                if sd._bollinger["UpperBand"][0] > sd._keltner["UpperBand"][0] and \
                    sd._bollinger["LowerBand"][0] < sd._keltner["LowerBand"][0] and \
                    sd.momWindow[1] < sd.momWindow[0]:
                        
                        insights.append(Insight.Price(sd.Security.Symbol, self.insightPeriod, InsightDirection.Up))
                        
                if sd._bollinger["UpperBand"][0] > sd._keltner["UpperBand"][0] and \
                    sd._bollinger["LowerBand"][0] < sd._keltner["LowerBand"][0] and \
                    sd.momWindow[1] > sd.momWindow[0]:
                        
                        insights.append(Insight.Price(sd.Security.Symbol, self.insightPeriod, InsightDirection.Down)) 
    
        #self.Debug(insights)
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.period, self.resolution)
            
        for removed in changes.RemovedSecurities:
            if removed.Symbol in self.symbolData:
                self.symbolData[removed.Symbol].Dispose()
                data = self.symbolData.pop(removed.Symbol, None)


class SymbolData:
    def __init__(self, algorithm, security, period, resolution):
        self.algorithm = algorithm
        self.Security = security
        self.bollinger = BollingerBands(period, 2, MovingAverageType.Simple)
        self.keltner = KeltnerChannels(period, 1.5, MovingAverageType.Simple)
        self.momentum = Momentum(period)
        self._bollinger = {}
        self._keltner = {}
        
        history = algorithm.History(security.Symbol, period, resolution)
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.bollinger.Update(bar.Index[1], bar.close)
            self.keltner.Update(tradeBar)
            self.momentum.Update(bar.Index[1], bar.close)
        
        self.consolidator = TradeBarConsolidator(timedelta(1))
        algorithm.RegisterIndicator(security.Symbol, self.bollinger, self.consolidator)
        algorithm.RegisterIndicator(security.Symbol, self.keltner, self.consolidator)
        algorithm.RegisterIndicator(security.Symbol, self.momentum, self.consolidator)
        
        self.bollinger.Updated += self.BollingerUpdated
        self.keltner.Updated += self.KeltnerUpdated
        self.momentum.Updated += self.MomentumUpdated
        
        self.bollingerWindow = RollingWindow[IndicatorDataPoint](2)
        self._bollinger["UpperBand"] = RollingWindow[float](2)
        self._bollinger["LowerBand"] = RollingWindow[float](2)
        
        self.keltnerWindow = RollingWindow[IndicatorDataPoint](2)
        self._keltner["UpperBand"] = RollingWindow[float](2)
        self._keltner["LowerBand"] = RollingWindow[float](2)
        
        self.momWindow = RollingWindow[IndicatorDataPoint](2)
        
    def BollingerUpdated(self, sender, updated):
        self.bollingerWindow.Add(updated)
        self._bollinger["UpperBand"].Add(self.bollinger.UpperBand.Current.Value)
        self._bollinger["LowerBand"].Add(self.bollinger.LowerBand.Current.Value)
        
    def KeltnerUpdated(self, sender, updated):
        self.keltnerWindow.Add(updated)
        self._keltner["UpperBand"].Add(self.keltner.UpperBand.Current.Value)
        self._keltner["LowerBand"].Add(self.keltner.LowerBand.Current.Value)
        
    def MomentumUpdated(self, sender, updated):
        self.momWindow.Add(updated)
        
    def Dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Security.Symbol, self.consolidator)

class EmaCrossAlphaModel(AlphaModel):
    '''Alpha model that uses an EMA cross to create insights'''

    def __init__(self,
                 fastPeriod = 12,
                 slowPeriod = 26,
                 resolution = Resolution.Daily):
        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.symbolDataBySymbol = {}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{})'.format(self.__class__.__name__, fastPeriod, slowPeriod, resolutionString)


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady:

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))

            symbolData.FastIsOverSlow = symbolData.Fast > symbolData.Slow

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                # create fast/slow EMAs
                symbolData = SymbolData(added)
                symbolData.Fast = algorithm.EMA(added.Symbol, self.fastPeriod, self.resolution)
                symbolData.Slow = algorithm.EMA(added.Symbol, self.slowPeriod, self.resolution)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol
        self.Fast = None
        self.Slow = None

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow

# Your New Python File

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from TheSqueeze2 import TheSqueeze
#from TheSqueeze import TheSqueeze
#from SqueezeAlternateStructure import TheSqueeze
#from TestAlpha2 import TheSqueeze
from RiskManagement import ATRRiskManagementModel
#from Universe import LiquidUniverseSelection


class LogicalFluorescentOrangeDinosaur(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 1)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction) 
        self.indicators = { }
        
        self.SetAlpha(TheSqueeze())
        
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None))
        
        #self.AddRiskManagement(ATRRiskManagementModel())
        
        self.SetExecution(ImmediateExecutionModel())
        #time flag for weekly selection
        #self.week = 0
        self.lastMonth = -1
        
    def CoarseSelectionFunction(self, universe):  
            current_week = self.Time.isocalendar()[1]
            #if current_week == self.week:
            #    return Universe.Unchanged
            if self.Time.month == self.lastMonth:
                return Universe.Unchanged
            self.lastMonth = self.Time.month
            
            selected = []
            universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)  
            universe = [c for c in universe if c.Price > 0.50 and c.HasFundamentalData]
            
            #check time flag
            #if self.week == self.Time.week: return
            #if not self.Time.weekday() == 0: return
    
            for coarse in universe:  
                symbol = coarse.Symbol
                
                if symbol not in self.indicators:
                    # 1. Call history to get an array of 20 days of history data
                    history = self.History(symbol, 21, Resolution.Daily)
                    
                    #2. Adjust SelectionData to pass in the history result
                    self.indicators[symbol] = SelectionData(history) 
    
                self.indicators[symbol].update(self.Time, coarse.AdjustedPrice)
                
                if  self.indicators[symbol].is_ready() and \
                    self.indicators[symbol].bollinger.UpperBand.Current.Value < self.indicators[symbol].keltner.UpperBand.Current.Value and \
                    self.indicators[symbol].bollinger.LowerBand.Current.Value > self.indicators[symbol].keltner.LowerBand.Current.Value:
                    
                    selected.append(symbol)
            
            #update time flag        
            #self.week = current_week
            
            #self.Log(selected)        
            
            return selected
            
class SelectionData():
    #3. Update the constructor to accept a history array
    def __init__(self, history):
        self.bollinger = BollingerBands(20, 2, MovingAverageType.Simple)
        self.keltner = KeltnerChannels(20, 1.5, MovingAverageType.Simple)
        #4. Loop over the history data and update the indicatorsc
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.bollinger.Update(bar.Index[1], bar.close)
            self.keltner.Update(tradeBar)
    
   # @property
    #def BollingerUpper(self):
     #   return float(self.bollinger.UpperBand.Current.Value)
        
    #@property
    #def BollingerLower(self):
     #   return float(self.bollinger.LowerBand.Current.Value)
        
    #@property
    #def KeltnerUpper(self):
     #   return float(self.keltner.UpperBand.Current.Value)
        
    #@property
    #def KeltnerLower(self):
    #    return float(self.keltner.LowerBand.Current.Value)
    
    
    def is_ready(self):
        return self.bollinger.IsReady and self.keltner.IsReady
    
    def update(self, time, value):
        return self.bollinger.Update(time, value)

class TheSqueeze(AlphaModel):
    
    def __init__(self, period = 20, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), period)        
        self.symbolDataBySymbol = {}
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)
        
    def Update(self, algorithm, data):
        
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.bollinger.IsReady and \
                symbolData.keltner.IsReady and \
                symbolData.momentum.IsReady: 
        
                if algorithm.Portfolio[key].Invested: continue
            
                if symbolData.bollinger.UpperBand > symbolData.keltner.UpperBand and \
                    symbolData.bollinger.LowerBand < symbolData.keltner.LowerBand and \
                    symbolData.momentum > 0:
                        
                        insights.append(Insight.Price(symbolData.Security.Symbol, self.predictionInterval, InsightDirection.Up))
                        
                if symbolData.bollinger.UpperBand > symbolData.keltner.UpperBand and \
                    symbolData.bollinger.LowerBand < symbolData.keltner.LowerBand and \
                    symbolData.momentum < 0:
                        
                        insights.append(Insight.Price(symbolData.Security.Symbol, self.predictionInterval, InsightDirection.Down)) 
    
        #self.Debug(insights)
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            
            if symbolData is None:
                symbolData = SymbolData(added)
                symbolData.bollinger = algorithm.BollingerBands(added.Symbol, self.period, 2, MovingAverageType.Simple, self.resolution)
                symbolData.keltner = algorithm.KeltnerChannels(added.Symbol, self.period, 1.5, MovingAverageType.Simple, self.resolution)
                symbolData.momentum = algorithm.Momentum(added.Symbol, self.period, self.resolution)
                self.symbolDataBySymbol[added.Symbol] = symbolData
                
            else:
                symbolData.bollinger.Reset()
                symbolData.keltner.Reset()
                symbolData.momentum.Reset()

class SymbolData:
    def __init__(self, algorithm, security, period, resolution):
        self.Security = security
        self.Symbol = security.Symbol
        self.bollinger = None
        self.keltner = None
        self.momentum = None

class ATRRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits drawdown based upon 2 x 20 Day ATR'''

    def __init__(self, atrPeriod = 20):
        '''Initializes a new instance of the ATRRiskManagementModel class
        Args:
            ATRPeriod: 20'''
        self.atrPeriod = atrPeriod

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        targets = []
        
        # check if security is long/short
        # get the value of the security
        # subtract/add 2*ATR 
        # append targets/liquidate
        # dont forget history request for the indicator
        
        for kvp in algorithm.Securities:
            #kvp.Value = symbol ticker
            security = kvp.Value
            symbol = kvp.Key
            
            if not security.Invested:
                continue

            if algorithm.Portfolio[symbol].IsLong:
                lossLevel = security.Price - (2*self.atr.Current.Value)
                if lossLevel == security.Price:
                    # liquidate
                    targets.append(PortfolioTarget(security.Symbol, 0))
                    
            if algorithm.Portfolio[symbol].IsShort:
                lossLevel = security.Price + (2*self.atr.Current.Value)
                if lossLevel == security.Price:
                    # liquidate
                    targets.append(PortfolioTarget(security.Symbol, 0))

        return targets

    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            if not added.Invested: continue
        
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.atrPeriod)
            
        for removed in changes.RemovedSecurities:
            if removed.Invested: continue
        
            data = self.symbolData.pop(removed.Symbol, None)

class SymbolData:
    def __init__(self, algorithm, security, atrPeriod):
        self.Security = security
        self.atr = AverageTrueRange(atrPeriod, MovingAverageType.Simple)
        
        history = algorithm.History(security.Symbol, 20, Resolution.Daily)
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.momentum.Update(bar.Index[1], bar.close)

class TheSqueeze(AlphaModel):
    
    def __init__(self, period = 20, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolData = {}
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)
        
    def Update(self, algorithm, data):
        
        insights = []
        for key, sd in self.symbolData.items():
            if sd.bollinger.IsReady and \
                sd.keltner.IsReady and \
                sd.momentum.IsReady and \
                sd._bollinger["UpperBand"].IsReady and \
                sd._bollinger["LowerBand"].IsReady and \
                sd._keltner["UpperBand"].IsReady and \
                sd._keltner["LowerBand"].IsReady and \
                sd.momWindow.IsReady: 
        
                if algorithm.Portfolio[key].Invested: continue
            
                if sd._bollinger["UpperBand"][0] > sd._keltner["UpperBand"][0] and \
                    sd._bollinger["LowerBand"][0] < sd._keltner["LowerBand"][0] and \
                    sd.momWindow[1] < sd.momWindow[0]:
                        
                        insights.append(Insight.Price(sd.Security.Symbol, self.insightPeriod, InsightDirection.Up))
                        
                if sd._bollinger["UpperBand"][0] > sd._keltner["UpperBand"][0] and \
                    sd._bollinger["LowerBand"][0] < sd._keltner["LowerBand"][0] and \
                    sd.momWindow[1] > sd.momWindow[0]:
                        
                        insights.append(Insight.Price(sd.Security.Symbol, self.insightPeriod, InsightDirection.Down)) 
    
        #self.Debug(insights)
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.period, self.resolution)
            #self.Log(self.self.symbolData[added.Symbol])
            
        for removed in changes.RemovedSecurities:
            data = self.symbolData.pop(removed.Symbol, None)
            
            if data is not None:
                algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator)


class SymbolData:
    def __init__(self, algorithm, security, period, resolution):
        self.Security = security
        self.bollinger = BollingerBands(period, 2, MovingAverageType.Simple)
        self.keltner = KeltnerChannels(period, 1.5, MovingAverageType.Simple)
        self.momentum = Momentum(period)
        self._bollinger = {}
        self._keltner = {}
        
        self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        algorithm.RegisterIndicator(security.Symbol, self.bollinger, self.Consolidator)
        algorithm.RegisterIndicator(security.Symbol, self.keltner, self.Consolidator)
        algorithm.RegisterIndicator(security.Symbol, self.momentum, self.Consolidator)
        
        self.bollinger.Updated += self.BollingerUpdated
        self.keltner.Updated += self.KeltnerUpdated
        self.momentum.Updated += self.MomentumUpdated
        
        self.bollingerWindow = RollingWindow[IndicatorDataPoint](2)
        self._bollinger["UpperBand"] = RollingWindow[float](2)
        self._bollinger["LowerBand"] = RollingWindow[float](2)
        
        self.keltnerWindow = RollingWindow[IndicatorDataPoint](2)
        self._keltner["UpperBand"] = RollingWindow[float](2)
        self._keltner["LowerBand"] = RollingWindow[float](2)
        
        self.momWindow = RollingWindow[IndicatorDataPoint](2)
        
    def BollingerUpdated(self, sender, updated):
        self.bollingerWindow.Add(updated)
        self._bollinger["UpperBand"].Add(self.bollinger.UpperBand.Current.Value)
        self._bollinger["LowerBand"].Add(self.bollinger.LowerBand.Current.Value)
        
    def KeltnerUpdated(self, sender, updated):
        self.keltnerWindow.Add(updated)
        self._keltner["UpperBand"].Add(self.keltner.UpperBand.Current.Value)
        self._keltner["LowerBand"].Add(self.keltner.LowerBand.Current.Value)
        
    def MomentumUpdated(self, sender, updated):
        self.momWindow.Add(updated)
        
        

#THIS WORKS
class TheSqueeze(AlphaModel):
    
    def __init__(self, period = 20, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolData = {}
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)
        
    def Update(self, algorithm, data):
        
        insights = []
        for symbol, symbolData in self.symbolData.items():
            if symbolData.bollinger.IsReady and \
                symbolData.keltner.IsReady and \
                symbolData.momentum.IsReady: 
        
                #if algorithm.Portfolio[symbol].Invested: continue
            
                if symbolData.bollinger.UpperBand.Current.Value > symbolData.keltner.UpperBand.Current.Value and \
                    symbolData.bollinger.LowerBand.Current.Value < symbolData.keltner.LowerBand.Current.Value and \
                    symbolData.momentum.Current.Value > 0:
                        
                        insights.append(Insight.Price(symbolData.Security.Symbol, self.insightPeriod, InsightDirection.Up))
                        
                if symbolData.bollinger.UpperBand.Current.Value > symbolData.keltner.UpperBand.Current.Value and \
                    symbolData.bollinger.LowerBand.Current.Value < symbolData.keltner.LowerBand.Current.Value and \
                    symbolData.momentum.Current.Value < 0:
                        
                        insights.append(Insight.Price(symbolData.Security.Symbol, self.insightPeriod, InsightDirection.Down))
    
        #self.Debug(insights)
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.period, self.resolution)
            
        for removed in changes.RemovedSecurities:
            data = self.symbolData.pop(removed.Symbol, None)


class SymbolData:
    def __init__(self, algorithm, security, period, resolution):
        self.Security = security
        self.bollinger = BollingerBands(period, 2, MovingAverageType.Simple)
        self.keltner = KeltnerChannels(period, 1.5, MovingAverageType.Simple)
        self.momentum = Momentum(period)
        self._bollinger = {}
        self._keltner = {}
        
        history = algorithm.History(security.Symbol, 20, resolution)
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.bollinger.Update(bar.Index[1], bar.close)
            self.keltner.Update(tradeBar)
            self.momentum.Update(bar.Index[1], bar.close)
        
        

class LiquidUniverseSelection(QCAlgorithm):
    
    def __init__(self, algorithm):
        self.algorithm = algorithm
        self.securities = []
       
    def CoarseSelectionFunction(self, coarse):
        # sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        coarseSelection = [x for x in coarse if x.HasFundamentalData and x.DollarVolume > 5000000]
        
        universe = [x.Symbol for x in coarseSelection]
        #self.algorithm.Securities = universe
        #self.Log(universe)
        return universe

    #def OnData(self, data):
       #if self._changes is None: return
    
        #for security in self._changes.RemovedSecurities:
            #if security.Invested:
                #self.securities.remove(security.Symbol)
                
        #for security in self._changes.AddedSecurities:
            #pass
        
        #self._changed = None
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.securities.append(added)
            
        for removed in changes.RemovedSecurities:
            if removed in self.securities:
                self.securities.remove(removed)
                
        for invested in self.securities.Invested:
            self.securities.remove(invested)
        
        
        #self.Log(f"OnSecuritiesChanged({self.UtcTime}):: {changes}")