from AlgorithmImports import *


class BBandsAlgorithm(QCAlgorithm):

    def Initialize(self) -> None:
        # backtest settings
        self.SetCash(100000)
        self.SetStartDate(2014, 1, 1)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.lastSellTime = datetime.min

        self.BBstd = float(self.GetParameter("stdBB")) # 4
        self.BBvalue = float(self.GetParameter("BB")) # 15
        self.ROCvalue = float(self.GetParameter("ROC")) # 79

        # Initialize an empty dictionary to store last trade times (new addition)
        self.lastTradeTimes = {}

        # settings
        self.EnableAutomaticIndicatorWarmUp = True
        self.Settings.FreePortfolioValuePercentage = 0.05
        self.UniverseSettings.Resolution = Resolution.Daily

        # ETF universe
        self.etf = self.AddEquity("SPY", self.UniverseSettings.Resolution).Symbol
        self.AddUniverseSelection(
            ETFConstituentsUniverseSelectionModel(self.etf, self.UniverseSettings, self.ETFConstituentsFilter))

        # Alternative investments
        self.alternatives = {
            'UUP': self.AddEquity('UUP', self.UniverseSettings.Resolution).Symbol,
            'TLT': self.AddEquity('TLT', self.UniverseSettings.Resolution).Symbol,
            'GLD': self.AddEquity('GLD', self.UniverseSettings.Resolution).Symbol
        }

        self.SetBenchmark(self.etf)

        self.symbolData = {}
        self.universe = []
        self.buy_prices = {}

        # initialize flag for stop loss triggered
        self.stop_loss_triggered = False
    
    def ETFConstituentsFilter(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
        # validate in-data
        if constituents is None:
            return Universe.Unchanged

        selected = [c for c in constituents if c.Weight]
        sorted_selected = sorted([c for c in selected], key=lambda c: c.Weight, reverse=True)[:(int(self.GetParameter("selected")))] #20

        self.universe = [c.Symbol for c in sorted_selected]
        return self.universe


    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            symbol = order.Symbol  # Extract the symbol from the order
            fillPrice = orderEvent.FillPrice  # Extract the fill price of the order

            # Check if it's a sell order and tag accordingly
            if order.Direction == OrderDirection.Sell:
                # Tagging for stop loss
                if fillPrice <= 0.95 * self.buy_prices.get(symbol, float('inf')):
                    order.Tag = "STOP LOSS"
                # Tagging for profit sell
                elif fillPrice >= 1.05 * self.buy_prices.get(symbol, 0):
                    order.Tag = "5%+ PROFIT"
                # Tagging if held for more than 10 days
                elif (self.Time - self.lastTradeTimes.get(symbol, datetime.min)).days >= 10:
                    order.Tag = "HELD FOR 10 DAYS"

            # Check if it's a buy order and tag accordingly
            if order.Direction == OrderDirection.Buy:
                order.Tag = "BUY CONDITIONS MET"

            # After tagging, here you may want to update your records such as `lastTradeTimes` or `buy_prices`
            if order.Direction == OrderDirection.Buy:
                self.buy_prices[symbol] = fillPrice
                self.lastTradeTimes[symbol] = self.Time  # Record the time of the purchase
            elif order.Direction == OrderDirection.Sell and symbol in self.lastTradeTimes:
                del self.lastTradeTimes[symbol] 

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        # validate in-data
        if changes is None:
            return

        for security in changes.AddedSecurities:
            self.symbolData[security.Symbol] = SymbolData(self, security.Symbol)

        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)
            symbolData = self.symbolData.pop(security.Symbol, None)
            if symbolData:
                symbolData.dispose()

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

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

        # liquidate assets that we should not trade
        for symbol in [x.Key for x in self.Portfolio if x.Value.Invested]:
            if symbol not in self.universe or not self.Securities[symbol].IsTradable:
                self.Liquidate(symbol)

            # Time-based Exit
            if symbol in self.lastTradeTimes:
                holding_period = (self.Time - self.lastTradeTimes[symbol]).days
                if holding_period >= 10:
                    self.Liquidate(symbol, "HELD FOR 10 DAYS")
                    # After liquidation, it's a good practice to remove the symbol from lastTradeTimes
                    del self.lastTradeTimes[symbol]
                    continue

        # Calculate the ROC for all symbols and select the max ROC symbol involving only positive ROC values
        roc_values = {symbol: self.symbolData[symbol].roc.Current.Value for symbol in self.universe if
                    symbol in self.symbolData and self.symbolData[symbol].roc.Current.Value > 0}
        if roc_values:
            max_roc_symbol = max(roc_values, key=roc_values.get)
            max_roc = roc_values[max_roc_symbol]
        else:
            max_roc_symbol = None
            max_roc = None

        # Check if all ROCs are negative and liquidate/reallocate if necessary
        if all(value < 0 for value in roc_values.values()):
            self.Liquidate(tag = "Liquidate for New Max ROC")
            self.Reallocate()
            return

        # Implement the Buy Conditions
        if max_roc_symbol and self.CanInvest() and max_roc > 0:
            price = self.Securities[max_roc_symbol].Price
            symbolData = self.symbolData[max_roc_symbol]
            if symbolData.bb.MiddleBand.Current.Value < price < symbolData.bb.UpperBand.Current.Value:
                quantity = self.CalculateOrderQuantity(max_roc_symbol, 1)
                orderTicket = self.MarketOrder(max_roc_symbol, quantity)
                orderTicket.UpdateTag(f"BUY CONDITIONS MET")
                self.buy_prices[max_roc_symbol] = price
                self.lastTradeTimes[max_roc_symbol] = self.Time  # Update the last trade time when a new holding is set (new addition)

        # Implement the Sell Conditions
        for symbol in self.Portfolio.Keys:
            price = self.Securities[symbol].Price
            if self.Portfolio[symbol].Invested:
                # Profit taking
                if price >= 1.05 * self.buy_prices.get(symbol, 0):
                    self.Liquidate(symbol,"5%+ PROFIT")
                    # It's a good practice to remove the symbol from lastTradeTimes after liquidation
                    if symbol in self.lastTradeTimes:
                        del self.lastTradeTimes[symbol]
                    continue

                # Stop Loss
                if price <= 0.95 * self.buy_prices.get(symbol, self.Portfolio[symbol].AveragePrice):
                    self.Liquidate(symbol, "STOP LOSS")
                    # After liquidation, it's a good practice to remove the symbol from lastTradeTimes
                    if symbol in self.lastTradeTimes:
                        del self.lastTradeTimes[symbol]
                    continue

    def CanInvest(self) -> bool:
        return  sum(1 for x in self.Portfolio if x.Value.Invested) < 10

    def Reallocate(self) -> None:
        # Calculate ROC for UUP and TLT
        roc_uup = self.symbolData[self.alternatives['UUP']].roc.Current.Value
        roc_tlt = self.symbolData[self.alternatives['TLT']].roc.Current.Value
        
        # Decide which alternative to buy based on ROC
        if roc_uup > 0 or roc_tlt > 0:
            if roc_uup > roc_tlt:
                quantity = self.CalculateOrderQuantity(self.alternatives['UUP'], 1)
                orderTicket = self.MarketOrder(self.alternatives['UUP'], quantity)
                orderTicket.UpdateTag(f"UUP ROC > 0")

            else:
                quantity = self.CalculateOrderQuantity(self.alternatives['TLT'], 1)
                orderTicket = self.MarketOrder(self.alternatives['TLT'], quantity)
                orderTicket.UpdateTag(f"TLT ROC > 0")

        else:
            # If both UUP and TLT have negative ROC, invest in GLD
                quantity = self.CalculateOrderQuantity(self.alternatives['GLD'], 1)
                orderTicket = self.MarketOrder(self.alternatives['GLD'], quantity)
                orderTicket.UpdateTag(f"GLD ROC > 0")

        # If any equities meet buy conditions, liquidate alternatives
        for symbol in self.universe:
            if self.Securities[symbol].Invested:
                symbol_data = self.symbolData[symbol]
                price = self.Securities[symbol].Price
                if symbol_data.roc.Current.Value > 0 and symbol_data.bb.MiddleBand.Current.Value < price < symbol_data.bb.UpperBand.Current.Value:
                    # Liquidate UUP, TLT, and GLD
                    self.Liquidate(self.alternatives['UUP'], "Liquidating for New Max ROC")
                    self.Liquidate(self.alternatives['TLT'], "Liquidating for New Max ROC")
                    self.Liquidate(self.alternatives['GLD'], "Liquidating for New Max ROC")
                    break  # Exit the loop as we only need to liquidate once

        
class SymbolData(object):
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol

        self.stdBB = float(algorithm.BBstd)
        self.BBi = int(algorithm.BBvalue)
        self.ROCi = int(algorithm.ROCvalue)

        # Assuming 'BB' and 'ROC' are methods provided by QuantConnect's QCAlgorithm class
        self.bb = algorithm.BB(symbol, self.BBi, self.stdBB, MovingAverageType.Simple, Resolution.Daily)
        self.roc = algorithm.ROC(symbol, self.ROCi, Resolution.Daily)

    def dispose(self):
        # deregister indicators and remove consolidator
        pass