class HedgedEquityAlgo(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2016, 7, 1)  
        self.SetEndDate(2021, 9, 3)  
        self.SetCash(1000000)  
        self.AddEquity("SPY", Resolution.Minute)
        self.SetBenchmark("SPY")
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.Settings.FreePortfolioValuePercentage = 0.1
       
        symbols_one_list = ["TQQQ"]
        self.symbols_one = [self.AddEquity(symbol, Resolution.Minute, Market.USA).Symbol for symbol in symbols_one_list]
        
        symbols_two_list = ["SQQQ"]
        self.symbols_two = [self.AddEquity(symbol, Resolution.Minute, Market.USA).Symbol for symbol in symbols_two_list]
        
        self.AddUniverseSelection(ManualUniverseSelectionModel(self.symbols_one))
        self.AddUniverseSelection(ManualUniverseSelectionModel(self.symbols_two))

        self.lookback = 65
        self.SetWarmup(self.lookback)
        self.AddAlpha(LongEquityAlphaModel(self.symbols_one, self.Time, self.lookback))
        self.AddAlpha(ShortEquityAlphaModel(self.symbols_two, self.Time, self.lookback))
        
    def RebalanceFunction(self, time):
        return None
                        
class LongEquityAlphaModel(AlphaModel):
                        
    def __init__(self, symbols, Time, lookback):
        self.symbols = symbols
        self.dataBySymbol = {}
        self.rebalanceTime = Time
        self.lookback = lookback
        
    def Update(self, algorithm, data):
        insights = []
         
        if algorithm.Time < self.rebalanceTime: return []
        
        for symbol, symbolData in self.dataBySymbol.items():
            if data.Bars.ContainsKey(symbol) and symbolData.IsReady():
                if symbol in self.symbols:
                    if symbolData.slow.Current.Value < symbolData.fast.Current.Value:
                        if not algorithm.Portfolio[symbol].IsLong:
                            insights.append(Insight(symbol, timedelta(days=30), InsightType.Price, InsightDirection.Up))
                    else:
                        if algorithm.Portfolio[symbol].IsLong:
                            insights.append(Insight(symbol, timedelta(days=30), InsightType.Price, InsightDirection.Flat))
                        
        self.rebalanceTime = Expiry.EndOfDay(algorithm.Time)
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for change in changes.AddedSecurities:
            if change.Symbol not in self.symbols:
                continue
            else:
                self.dataBySymbol[change.Symbol] = SymbolData(algorithm, change.Symbol, self.lookback) 
            
        for change in changes.RemovedSecurities:
            if change.Symbol not in self.symbols:
                continue
            else:
                if change.Symbol in self.dataBySymbol:
                    del self.dataBySymbol[change.Symbol]
                
                
class ShortEquityAlphaModel(AlphaModel):
                        
    def __init__(self, symbols, Time, lookback):
        self.symbols = symbols
        self.dataBySymbol = {}
        self.rebalanceTime = Time
        self.lookback = lookback
        
    def Update(self, algorithm, data):
        insights = []
         
        if algorithm.Time < self.rebalanceTime: return []
        
        for symbol, symbolData in self.dataBySymbol.items():
            if data.Bars.ContainsKey(symbol) and symbolData.IsReady():
                if symbol in self.symbols:
                    if symbolData.slow.Current.Value > symbolData.fast.Current.Value:
                        if not algorithm.Portfolio[symbol].IsShort:
                            insights.append(Insight(symbol, timedelta(days=30), InsightType.Price, InsightDirection.Down))
                    else:
                        if algorithm.Portfolio[symbol].IsShort:
                            insights.append(Insight(symbol, timedelta(days=30), InsightType.Price, InsightDirection.Flat))
                    
        self.rebalanceTime = Expiry.EndOfDay(algorithm.Time)
        return insights
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for change in changes.AddedSecurities:
            if change.Symbol not in self.symbols:
                continue
            else:
                self.dataBySymbol[change.Symbol] = SymbolData(algorithm, change.Symbol, self.lookback) 
            
        for change in changes.RemovedSecurities:
            if change.Symbol not in self.symbols:
                continue
            else:
                if change.Symbol in self.dataBySymbol:
                    del self.dataBySymbol[change.Symbol]
                
class SymbolData:
    def __init__(self, algorithm, symbol, lookback):
        
        algorithm.Consolidate(symbol, Resolution.Daily, self.DailyBarHandler)
        
        self.fast = algorithm.SMA(symbol, 5, Resolution.Daily, Field.Low)
        self.slow = algorithm.SMA(symbol, 60, Resolution.Daily, Field.High)
        
        history = algorithm.History(symbol, lookback, Resolution.Daily)
        
        if not history.empty:
            for index, tradebar in history.loc[symbol].iterrows():
                self.fast.Update(index, tradebar.low)
                self.slow.Update(index, tradebar.high)
                
            last_row = history.loc[symbol].iloc[-1]
            self.open = last_row.open
            self.close = last_row.close
            self.high = last_row.high
            self.low = last_row.low
        
    def DailyBarHandler(self, consolidated):
        self.open = consolidated.Open
        self.close = consolidated.Close
        self.high = consolidated.High
        self.low = consolidated.Low
            
    def IsReady(self):
        return self.fast.IsReady and self.slow.IsReady