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.BBstd = float(self.GetParameter("stdBB")) #4
        self.BBvalue = float(self.GetParameter("BB")) #17.6
        self.ROCvalue = float(self.GetParameter("ROC")) #24.8

        # 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 = {}

    def OnData(self, data: Slice) -> None:
        # 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)

        # Calculate the ROC for all symbols and select the max ROC symbol
        roc_values = {symbol: self.symbolData[symbol].roc.Current.Value for symbol in self.universe if
                      symbol in self.symbolData}
        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()
            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:
                self.SetHoldings(max_roc_symbol, 1)
                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)
                    continue

                # Stop Loss
                if price <= 0.95 * self.buy_prices.get(symbol, self.Portfolio[symbol].AveragePrice):
                    self.Liquidate(symbol)
                  #  self.stop_loss_triggered = True
                    continue

                # Time-based Exit
                if symbol in self.lastTradeTimes:  # Check if the last trade time exists (new addition)
                    time_invested = (self.Time - self.lastTradeTimes[
                        symbol]).days  # Use the stored last trade time (modified line)
                    if time_invested > 10:
                        self.Liquidate(symbol)
                        continue

    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 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 CanInvest(self) -> bool:
        return  sum(1 for x in self.Portfolio if x.Value.Invested) < 10

    def Reallocate(self) -> None:
        for symbol, security in self.alternatives.items():
            self.SetHoldings(security, 1 / len(self.alternatives))

    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug(str(orderEvent))
        
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