# region imports
from AlgorithmImports import *

class ROCMomentumAlgorithm(QCAlgorithm):

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

        self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))
        self.Settings.FreePortfolioValuePercentage = 0.05
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        
        # Dictionary to store buy prices
        self.buy_prices = {}

        # Dictionary to store purchase dates
        self.purchase_dates = {}

        # 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 13 periods and 2 standard deviation
        self.bbands_equities = {symbol: self.BB(symbol, self.GetParameter("BB"), self.GetParameter("stdBB"), MovingAverageType.Simple, Resolution.Daily) for symbol in self.equities}

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

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

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

        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

                # Remove the equity from purchase_dates dictionary if it's sold
                if orderEvent.Symbol in self.purchase_dates:
                    del self.purchase_dates[orderEvent.Symbol]

            # Update buy price for the purchased asset and record its purchase date
            if order.Direction == OrderDirection.Buy:
                self.buy_prices[order.Symbol] = orderEvent.FillPrice
                self.purchase_dates[order.Symbol] = self.Time

    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 and buy logic for equity
        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]
            max_roc_symbol = max(self.roc_equities, key=lambda s: self.roc_equities[s].Current.Value)

            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}"
            
            # Now, run the 20-day check
            purchase_date = self.purchase_dates.get(symbol)
            if purchase_date and (self.Time - purchase_date).days >= 10:
                self.Liquidate(symbol, "Held for >= 10 days")
                del self.purchase_dates[symbol]

            # Check if the current equity is the one with max ROC
            if symbol == max_roc_symbol:
                # Now check if the current price of the equity with max ROC is above its middle band
                if current_price > middle_band and current_price < upper_band:
                    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, 1)
                        orderTicket = self.MarketOrder(max_roc_symbol, quantity)
                        orderTicket.UpdateTag(f"Buy with max ROC, {tag_message}")

                    # Explicit profit-taking mechanism
                    if self.Portfolio[max_roc_symbol].Invested:
                        if current_price >= 1.05 * self.buy_prices.get(max_roc_symbol, 0):
                            tag_message1 = f"Prev Close: {prev_close}, Current Price: {current_price}, Lower Band: {lower_band}, Middle Band: {middle_band}, Upper Band: {upper_band}, ROC: {current_roc}"
                            self.Liquidate(max_roc_symbol, f"Take profit >= 1.05, {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:
                    if isinstance(ticket, OrderTicket):
                        ticket.UpdateTag(f"Negative ROC for all equities")

                if not self.Portfolio.Invested:
                    target_symbol, tag_message = None, ""
                    if self.roc_uup.Current.Value > 0 and self.roc_tlt.Current.Value > 0:
                        target_symbol = self.uup if self.roc_uup.Current.Value > self.roc_tlt.Current.Value else self.tlt
                        tag_message = f"Buy UUP, ROC: {self.roc_uup.Current.Value}" if self.roc_uup.Current.Value > self.roc_tlt.Current.Value else f"Buy TLT, ROC: {self.roc_tlt.Current.Value}"
                    elif self.roc_uup.Current.Value < 0 and self.roc_tlt.Current.Value > 0:
                        target_symbol = self.tlt
                        tag_message = f"Buy TLT, ROC: {self.roc_tlt.Current.Value}"
                    elif self.roc_uup.Current.Value > 0 and self.roc_tlt.Current.Value < 0:
                        target_symbol = self.uup
                        tag_message = f"Buy UUP, ROC: {self.roc_uup.Current.Value}"
                    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: {self.roc_uup.Current.Value}")
                    orderTicket = self.MarketOrder(self.tlt, -self.Portfolio[self.tlt].Quantity)
                    orderTicket.UpdateTag(f"Liquidate TLT, ROC: {self.roc_tlt.Current.Value}")
                    orderTicket = self.MarketOrder(self.gld, -self.Portfolio[self.gld].Quantity)
                    orderTicket.UpdateTag("Liquidate GLD")