# region imports
from AlgorithmImports import *
# endregion

class FatVioletBarracuda(QCAlgorithm):

    def Initialize(self):
        # Backtesting Settings
        self.SetStartDate(2022, 1, 1)  
        self.SetEndDate(2022, 12, 31)
        self.SetCash(100000) 
        self.SetBenchmark("SPY")

        # Equity settings
        self.Data = dict()
        tickers = ["XLK", "XLY"]
        for ticker in tickers:
            symbol = self.AddEquity(ticker, Resolution.Daily, 
                            dataNormalizationMode = DataNormalizationMode.Raw).Symbol     
            self.Data[symbol] = SymbolData(symbol)
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)

        # Warm up
        self.SetWarmup(100, Resolution.Daily)

        # Trading Condition check:
        self.ATRM = 3
        self.Trailing_days = 0

    def OnData(self, slice: slice):   

        for symbol in self.Data.keys():
            symbolData = self.Data[symbol]
            if slice.Bars.ContainsKey(symbol):
                bar = slice.Bars[symbol]
                
                symbolData.CustomCdl = bar
                symbolData.cldWindow.Add(bar)
                symbolData.atr.Update(bar.Close)
                symbolData.bb.Update(bar.EndTime, Bar.Close)
                symbolData.rsi.Update(bar.EndTime, Bar.Close)
                symbolData.hma.Update(bar.EndTime, Bar.Close)


            if (self.IsWarmingUp or self.CustomCdl is None or not self.IsMarketOpen(symbol) or not self.hmaWindow.IsReady or
                not self.bbWindow.IsReady):
                return

            # Trading Condition
            if not self.Portfolio[symbol].Invested:

                # HMA check
                if symbolData.hmaWindow[1].Value > symbolData.hma.Current.Value:
                    self.hmacheck = True
                else:
                    self.hmacheck = False

                # BB LowerBand check
                if CustomCdl.Close > symbolData.bb.LowerBand.Current.Value and symbolData.cldWindow[1].Close < symbolData.bbLower["LowerBand"][1]: #or self.cldWindow[2].Close < self.bbLower["LowerBand"][2]
                    self.bbLowerBandcheck = True
                else:
                    self.bbLowerBandcheck = False

                # RSI/SMA check
                if symbolData.rsi.Current.Value > symbolData.rsisma.Current.Value:
                    self.rsicheck = True
                else:
                    self.rsicheck = False

                # Quantity, SL and TP setup
                Buy_SL_Price = round(CustomCdl.Close - (symbolData.atr.Current.Value * self.ATRM), 2)
                quantity = round(0.10 * self.Portfolio.TotalPortfolioValue / CustomCdl.Close)
                
                # BUY TRADING
                if self.hmacheck == True and self.bbLowerBandcheck == True and self.rsicheck == True: 
                    
                    #Buy 
                    self.BuyTicket = self.MarketOrder(symbol, quantity, tag = "LowerBand Buy")
                    self.Trailing_days = 0
                    
                    #Stop Loss
                    self.StopTicket = self.StopMarketOrder(symbol, 
                                        self.BuyTicket.Quantity * -1, Buy_SL_Price)   


class SymbolData(object):

    def __init__(self, symbol):
        self.symbol = symbol
        
        # Rolling Window candle
        self.cldWindow = RollingWindow[TradeBar](10)

        # Indicators
        self.atr = AverageTrueRange(14)
        self.hma = HullMovingAverage(100)
        self.bb = BollingerBands(50, 2)
        self.rsi = RelativeStrengthIndex(14)
        self.rsisma = IndicatorExtensions.SMA(self.rsi, 20)

        # Rolling Window Indicator
        self.hmaWindow = RollingWindow[IndicatorDataPoint](5)
        self.hma.Updated += self.HMAUpdated
        self.bbLower = {}
        self.bbMiddle = {}
        self.bbWindow = RollingWindow[IndicatorDataPoint](5)
        self.bbLower["LowerBand"] = RollingWindow[float](5)
        self.bbMiddle["MiddleBand"] = RollingWindow[float](5)
        self.bb.Updated += self.BBUpdated

    def HMAUpdated(self, sender, updated):
        #Adds updated values to rolling window
        self.hmaWindow.Add(updated)

    def BBUpdated(self, sender, updated):
        #Adds updated values to rolling window
        self.bbWindow.Add(updated)
        self.bbLower["LowerBand"].Add(self.bb.LowerBand.Current.Value)
        self.bbMiddle["MiddleBand"].Add(self.bb.MiddleBand.Current.Value)