# region imports
from AlgorithmImports import *

class ROCMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023,1,1)  # Set Start Date
        self.SetCash(1000)  # Set Strategy Cash
        self.lastSellTime = datetime.min

        self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))
        
        # Dictionary to store buy prices
        self.buy_prices = {}

        # Add Equity individually
        self.SPY = self.AddEquity("SPY", Resolution.Daily).Symbol  # SPY
        self.equity1 = self.AddEquity("AAPL", Resolution.Daily).Symbol  # Apple
        self.equity2 = self.AddEquity("MSFT", Resolution.Daily).Symbol  # Microsoft
        self.equity3 = self.AddEquity("AMZN", Resolution.Daily).Symbol  # Amazon
        self.equity4 = self.AddEquity("NVDA", Resolution.Daily).Symbol  # NVIDIA
        self.equity5 = self.AddEquity("TSLA", Resolution.Daily).Symbol  # Tesla
        self.equity6 = self.AddEquity("GOOGL", Resolution.Daily).Symbol  # Alphabet Class A
        self.equity7 = self.AddEquity("META", Resolution.Daily).Symbol  # Meta
        self.equity8 = self.AddEquity("GOOG", Resolution.Daily).Symbol  # Alphabet Class C
        self.equity9 = self.AddEquity("AVGO", Resolution.Daily).Symbol  # Broadcom
        self.equity10 = self.AddEquity("ORCL", Resolution.Daily).Symbol  # Oracle
        
        self.equities = [self.equity1, self.equity2, self.equity3, self.equity4, self.equity5,
                        self.equity6, self.equity7, self.equity8, self.equity9, self.equity10]

        self.uup = self.AddEquity("UUP", Resolution.Daily).Symbol  # PowerShares DB US Dollar Index Bullish Fund
        self.tlt = self.AddEquity("TLT", Resolution.Daily).Symbol  # iShares 20+ Year Treasury Bond ETF
        self.gld = self.AddEquity("GLD", Resolution.Daily).Symbol  # SPDR Gold Trust ETF
        
        # Define Bollinger Band for each symbol with 20 periods and 2 standard deviation
        self.bbands_equities = {symbol: self.BB(symbol, 55, 2, MovingAverageType.Simple, Resolution.Daily) for symbol in self.equities}

        # define our daily roc(50) indicators for each symbol
        self.roc_equities = {symbol: self.ROC(symbol, 55, Resolution.Daily) for symbol in self.equities}
        self.roc_uup = self.ROC(self.uup, 55, Resolution.Daily)
        self.roc_tlt = self.ROC(self.tlt, 55, Resolution.Daily)
        self.roc_gld = self.ROC(self.gld, 55, Resolution.Daily)

        # define a rolling window for the ROC for each symbol
        self.window_equities = {symbol: RollingWindow[IndicatorDataPoint](55) for symbol in self.equities}
        self.window_uup = RollingWindow[IndicatorDataPoint](55)
        self.window_tlt = RollingWindow[IndicatorDataPoint](55)
        self.window_gld = RollingWindow[IndicatorDataPoint](55)

        # Set warm-up period for 50 bars
        self.SetWarmUp(55)

        self.SetBenchmark(self.SPY)

        # initialize flag for stop loss triggered
        self.stop_loss_triggered = False

        self.previous_closes = {symbol: RollingWindow[float](2) for symbol in self.equities}

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            if order.Type == OrderType.StopMarket or order.Direction == OrderDirection.Sell:
                self.stop_loss_triggered = True
                self.lastSellTime = self.Time  # Record the time of the last sell

            # Update buy price for the purchased asset
            if order.Direction == OrderDirection.Buy:
                self.buy_prices[order.Symbol] = orderEvent.FillPrice

    def OnData(self, data):
        # Check if we are still warming up
        if self.IsWarmingUp:
            return

        if self.Time - self.lastSellTime < timedelta(days=1):
            return

        if not (all(roc.IsReady for roc in self.roc_equities.values()) and 
                all(data.ContainsKey(symbol) for symbol in self.equities) and
                self.roc_uup.IsReady and self.roc_tlt.IsReady and self.roc_gld.IsReady and 
                data.ContainsKey(self.uup) and data.ContainsKey(self.tlt) and data.ContainsKey(self.gld)):
            return 

        # Check the Bollinger Bands and ROC conditions
        for symbol in self.equities:
            current_price = self.Securities[symbol].Close
            lower_band = self.bbands_equities[symbol].LowerBand.Current.Value
            middle_band = self.bbands_equities[symbol].MiddleBand.Current.Value
            upper_band = self.bbands_equities[symbol].UpperBand.Current.Value
            current_roc = self.roc_equities[symbol].Current.Value

            # Store the current price for the symbol
            self.previous_closes[symbol].Add(current_price)
            if self.previous_closes[symbol].Count < 2:
                continue

            prev_close = self.previous_closes[symbol][1]

            # Tagging the values for better visibility
            tag_message = f"Prev Close: {prev_close}, Current Price: {current_price}, Lower Band: {lower_band}, Middle Band: {middle_band}, Upper Band: {upper_band}, ROC: {current_roc}"

            # Buy conditions
            if (prev_close <= lower_band and current_price > middle_band) and current_roc > 0:
                if not self.Portfolio[symbol].Invested:
                    self.SetHoldings(symbol, 0.95, tag_message)

            # Explicit profit-taking mechanism
            if self.Portfolio[symbol].Invested:
                if current_price >= 1.05 * self.buy_prices.get(symbol, 0):
                    self.Liquidate(symbol, f"Take profit, {tag_message}")


        roc_uup = self.roc_uup.Current.Value
        roc_tlt = self.roc_tlt.Current.Value

        # buying logic for equity
        if any(roc.Current.Value > 0 for roc in self.roc_equities.values()):
            max_roc_symbol = max(self.roc_equities, key=lambda symbol: self.roc_equities[symbol].Current.Value)
            tag_message1 = self.get_tag_message(max_roc_symbol)
            if self.Portfolio.Invested:
                return
            if self.Portfolio[max_roc_symbol].Invested and not self.stop_loss_triggered:
                return
            if not self.Portfolio[max_roc_symbol].Invested:
                orderTickets = self.Liquidate()
                for ticket in orderTickets:
                    ticket.UpdateTag(f"Liquidate for new max ROC")
                quantity = self.CalculateOrderQuantity(max_roc_symbol, 0.95)
                orderTicket = self.MarketOrder(max_roc_symbol, quantity)
                orderTicket.UpdateTag(f"Buy with max ROC, {tag_message1}")

            self.stop_loss_triggered = False
                    
        elif all(roc.Current.Value < 0 for roc in self.roc_equities.values()):
            orderTickets = self.Liquidate()
            for ticket in orderTickets:
                ticket.UpdateTag(f"Negative ROC for all equities")

            if not self.Portfolio.Invested:
                target_symbol, tag_message = None, ""
                if roc_uup > 0 and roc_tlt > 0:
                    target_symbol = self.uup if roc_uup > roc_tlt else self.tlt
                    tag_message = f"Buy UUP, ROC: {roc_uup}" if roc_uup > roc_tlt else f"Buy TLT, ROC: {roc_tlt}"
                elif roc_uup < 0 and roc_tlt > 0:
                    target_symbol = self.tlt
                    tag_message = f"Buy TLT, ROC: {roc_tlt}"
                elif roc_uup > 0 and roc_tlt < 0:
                    target_symbol = self.uup
                    tag_message = f"Buy UUP, ROC: {roc_uup}"
                else:
                    target_symbol = self.gld
                    tag_message = "Buy GLD"

                quantity = self.CalculateOrderQuantity(target_symbol, 0.95)
                orderTicket = self.MarketOrder(target_symbol, quantity)
                orderTicket.UpdateTag(tag_message)

            if any(roc.Current.Value > 0 for roc in self.roc_equities.values()):
                orderTicket = self.MarketOrder(self.uup, -self.Portfolio[self.uup].Quantity)
                orderTicket.UpdateTag(f"Liquidate UUP, ROC: {roc_uup}")
                orderTicket = self.MarketOrder(self.tlt, -self.Portfolio[self.tlt].Quantity)
                orderTicket.UpdateTag(f"Liquidate TLT, ROC: {roc_tlt}")
                orderTicket = self.MarketOrder(self.gld, -self.Portfolio[self.gld].Quantity)
                orderTicket.UpdateTag("Liquidate GLD")

    def get_tag_message(self, symbol):
        prev_close = self.previous_closes[symbol][1] if self.previous_closes[symbol].Count >= 2 else None
        current_price = self.Securities[symbol].Close
        lower_band = self.bbands_equities[symbol].LowerBand.Current.Value
        middle_band = self.bbands_equities[symbol].MiddleBand.Current.Value
        upper_band = self.bbands_equities[symbol].UpperBand.Current.Value
        current_roc = self.roc_equities[symbol].Current.Value
        
        return f"Prev Close: {prev_close}, Current Price: {current_price}, Lower Band: {lower_band}, Middle Band: {middle_band}, Upper Band: {upper_band}, ROC: {current_roc}"