from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from TheSqueeze import TheSqueeze
#from Universe import LiquidUniverseSelection


class LogicalFluorescentOrangeDinosaur(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction) 
        self.indicators = { }
        
        self.SetAlpha(TheSqueeze())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        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 200 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.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._bollinger["UpperBand"].IsReady and \
                sd._bollinger["LowerBand"].IsReady and \
                sd._keltner["UpperBand"].IsReady and \
                sd._keltner["LowerBand"].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:
            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 = {}
        
        #history = algorithm.History(security.Symbol, 20, resolution.Daily)
        #algorithm.RegisterIndicator(security.Symbol, self.bollinger, history)
        #algorithm.RegisterIndicator(security.Symbol, self.keltner, history)
        #algorithm.RegisterIndicator(security.Symbol, self.momentum, history)
        
        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)
        
        
# Your New Python File

# Your New Python File

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}")