# 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(self, 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, data):   

        if self.IsWarmingUp:
            return
        
        for symbol, symbolData in self.Data.items():
            if not (data.ContainsKey(symbol) and data[symbol] is not None and symbolData.IsReady):
                continue

            CustomCdl = data[symbol]

            # 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]: 
                    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, algorithm, symbol):
        
        # Rolling Window candle
        self.cldWindow = RollingWindow[TradeBar](10)

        # Indicators
        self.atr = algorithm.ATR(symbol, 14, Resolution.Daily)
        self.hma = algorithm.HMA(symbol, 100, Resolution.Daily)
        self.bb = algorithm.BB(symbol, 50, 2, Resolution.Daily)
        self.rsi = algorithm.RSI(symbol, 14, Resolution.Daily)
        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)

    def CldUpdated(self, sender, bar):
        self.cldWindow.Add(bar)

    @property 
    def IsReady(self):
        return self.hma.IsReady and self.hmaWindow.IsReady and self.bbWindow.IsReady and self.cldWindow.IsReady