from AlgorithmImports import *

class IntelligentSkyRodent(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)  # Set Start Date
        # self.SetEndDate(2020, 1, 10)
        self.SetCash(1000)  # Set Strategy Cash
        # self.SetSecurityInitializer(self.CustomSecurityInitializer)

        # Define a list of equities to be used by the algorithm
        self.equities = ["SPY", "TQQQ", "SPXL", "UVXY", "SQQQ", "BSV", "TECL", "SMH", "SOXL", "DIA", "UDOW", "URPO", "QQQ", "PSQ", "TLT", "XLP"]
        self.indicators = dict()  # Initialize a dictionary to store technical analysis indicators for each equity

        # Initialize technical analysis indicators for each equity
        for equity in self.equities:
            self.AddEquity(equity, Resolution.Minute)
            self.Securities[equity].SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.indicators[equity] = dict()

            # Initialize RSI indicators for each equity
            self.indicators[equity]['RSI'] = dict()
            self.indicators[equity]['RSI']['Period-5'] = RelativeStrengthIndex(5)
            self.indicators[equity]['RSI']['Period-10'] = RelativeStrengthIndex(10)
            self.indicators[equity]['RSI']['Period-14'] = RelativeStrengthIndex(14)
            self.indicators[equity]['RSI']['Period-60'] = RelativeStrengthIndex(60)
            self.indicators[equity]['RSI']['Period-70'] = RelativeStrengthIndex(70)

            # Initialize SMA indicators for each equity
            self.indicators[equity]['SMA'] = dict()
            self.indicators[equity]['SMA']['Period-20'] = SimpleMovingAverage(20)
            self.indicators[equity]['SMA']['Period-200'] = SimpleMovingAverage(200)

    # Schedule the FunctionBeforeMarketClose method to run once per day before market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.BeforeMarketClose("SPY", 5),
                self.FunctionBeforeMarketClose)

    # Define a method to calculate return of a given equity over a given period
    def CalculateReturn(self, equity: str, period: int) -> float:
        history = self.History(equity, period, Resolution.Daily)
        history = list(history)
        return (history[0].Close / history[-1].Close) - 1

        # Initialize trade bar consolidators for each equity
        self.consolidators = dict()
        for equity in self.equities:
            self.consolidators[equity] = TradeBarConsolidator(timedelta(minutes=5))
            self.consolidators[equity].DataConsolidated += self.OnDataConsolidated
            self.SubscriptionManager.AddConsolidator(equity, self.consolidators[equity])
        # Schedule the consolidators to run at the specified time before market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.BeforeMarketClose("SPY", 10),
                self.Consolidators_Ready)
    
    def Consolidators_Ready(self):
        # Tell the consolidators to start consolidating data
        for consolidator in self.consolidators.values():
            consolidator.Start()

    def OnDataConsolidated(self, sender, bar):
        equity = sender.Symbol.Value
        # Update all RSI indicators for the current equity
        for period, rsi in self.indicators[equity]['RSI'].items():
            # If the RSI indicator is not ready, update its history
            if not rsi.IsReady:
                history = self.History(equity, period, Resolution.Daily)
                for historical_bar in history:
                    rsi.Update(historical_bar.EndTime, historical_bar.Close)
            # Update the current RSI value
            rsi.Update(self.Time, bar.Close)
        # Update all SMA indicators for the current equity
        for period, sma in self.indicators[equity]['SMA'].items():
            # If the SMA indicator is not ready, update its history
            if not sma.IsReady:
                history = self.History(equity, period, Resolution.Daily)
                for historical_bar in history:
                    sma.Update(historical_bar.EndTime, historical_bar.Close)
            # Update the current SMA value
            sma.Update(self.Time, bar.Close)

    def FunctionBeforeMarketClose(self):

        # Iterate over each equity
       
        for equity in self.equities:
            
            if self.Portfolio[equity].HoldStock:
            # Liquidate the position
                self.Liquidate(equity)

            if not self.indicators[equity]['RSI']['Period-5'].IsReady:
                history = self.History(equity, 5, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['RSI']['Period-5'].Update(bar.EndTime, bar.Close)

            if not self.indicators[equity]['RSI']['Period-10'].IsReady:
                history = self.History(equity, 10, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['RSI']['Period-10'].Update(bar.EndTime, bar.Close)

            if not self.indicators[equity]['RSI']['Period-14'].IsReady:
                history = self.History(equity, 14, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['RSI']['Period-14'].Update(bar.EndTime, bar.Close)
                    
            if not self.indicators[equity]['RSI']['Period-60'].IsReady:
                history = self.History(equity, 60, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['RSI']['Period-60'].Update(bar.EndTime, bar.Close)


            if not self.indicators[equity]['RSI']['Period-70'].IsReady:
                history = self.History(equity, 70, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['RSI']['Period-70'].Update(bar.EndTime, bar.Close)

            if not self.indicators[equity]['SMA']['Period-20'].IsReady:
                history = self.History(equity, 20, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['SMA']['Period-20'].Update(bar.EndTime, bar.Close)

            if not self.indicators[equity]['SMA']['Period-200'].IsReady:
                history = self.History(equity, 200, Resolution.Daily)

                for bar in history:
                    self.indicators[equity]['SMA']['Period-200'].Update(bar.EndTime, bar.Close)

            self.indicators[equity]['RSI']['Period-5'].Update(self.Time, self.Securities[equity].Close)
            self.indicators[equity]['RSI']['Period-10'].Update(self.Time, self.Securities[equity].Close)
            self.indicators[equity]['RSI']['Period-14'].Update(self.Time, self.Securities[equity].Close)
            self.indicators[equity]['RSI']['Period-60'].Update(self.Time, self.Securities[equity].Close)
            self.indicators[equity]['RSI']['Period-70'].Update(self.Time, self.Securities[equity].Close)            
            self.indicators[equity]['SMA']['Period-20'].Update(self.Time, self.Securities[equity].Close)
            self.indicators[equity]['SMA']['Period-200'].Update(self.Time, self.Securities[equity].Close)
        # Implement the main decision tree for the algorithm

        # If the current price of SPY is above its 200-day SMA price
        if self.Securities['SPY'].Close > self.indicators['SPY']['SMA']['Period-200'].Current.Value:
            # If the 10-day RSI of TQQQ is above 79
            if self.indicators['TQQQ']['RSI']['Period-10'].Current.Value > 79:
                # Buy UVXY
                self.SetHoldings("UVXY", 1)
            else:
                # If the 10-day RSI of SPXL is above 80
                if self.indicators['SPXL']['RSI']['Period-10'].Current.Value > 80:
                    # Buy UVXY
                    self.SetHoldings("UVXY", 1)
                else:
                    # If the 5-day return of TQQQ is above 20%
                    if self.CalculateReturn('TQQQ', 5) > 0.2:
                        # If the 10-day RSI of TQQQ is below 31
                        if self.indicators['TQQQ']['RSI']['Period-10'].Current.Value < 31:
                            # Buy TQQQ
                            self.SetHoldings("TQQQ", 1)
                        else:
                            # Buy one asset from the list [UVXY, SQQQ] with the highest 10-day RSI
                            if self.indicators['UVXY']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
                                self.SetHoldings("UVXY", 1)
                            else:
                                self.SetHoldings("SQQQ", 1)
                    else:
                        # Buy TQQQ
                        self.SetHoldings("TQQQ", 1)
        else:
            # If the 10-day RSI of TQQQ is below 31
            if self.indicators['TQQQ']['RSI']['Period-10'].Current.Value < 31:
                  # Buy TECL
                self.SetHoldings("TECL", 1)
            else:
            # If the 10-day RSI of SMH is below 30
                if self.indicators['SMH']['RSI']['Period-10'].Current.Value < 30:
                # Buy SOXL
                    self.SetHoldings("SOXL", 1)
                else:
                # If the 10-day RSI of DIA is below 27
                    if self.indicators['DIA']['RSI']['Period-10'].Current.Value < 27:
                    # Buy UDOW
                        self.SetHoldings("UDOW", 1)
                    else:
                    # If the 14-day RSI of SPY is below 28
                        if self.indicators['SPY']['RSI']['Period-14'].Current.Value < 28:
                        # Buy URPO
                            self.SetHoldings("URPO", 1)
                        else:
                        # Allocate 50% of the portfolio to Group 1 and 50% to Group 1
                            self.Group1()
                            self.Group2()

    def Group1(self):
        # Group 1: If the 200-day return of QQQ is below -20%
        if self.CalculateReturn('QQQ', 200) < -0.2:
        # If the current price of QQQ is below its 20-day SMA price
            if self.Securities['QQQ'].Close < self.indicators['QQQ']['SMA']['Period-20'].Current.Value:
            # If the 60-day return of QQQ is below -12%
                if self.CalculateReturn('QQQ', 60) < -0.12:
                # Allocate 50% of the portfolio to Group 3 and 50% to Group 4
                    self.Group3()
                    self.Group4()
                else:
                        # If the 10-day RSI of TLT is above the 10-day RSI of SQQQ
                    if self.indicators['TLT']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
                    # Buy TQQQ
                        self.SetHoldings("TQQQ", 0.5)
                    else:
                    # Buy SQQQ
                        self.SetHoldings("SQQQ", 0.5)              
            else:
            # If the 10-day RSI of SQQQ is below 31
                if self.indicators['SQQQ']['RSI']['Period-10'].Current.Value < 31:
                # Buy PSQ
                    self.SetHoldings("PSQ", 0.5)
                else:
                # Choose one asset from the list [QQQ, SMH] with the highest 10-day RSI
                    if self.indicators['QQQ']['RSI']['Period-10'].Current.Value > self.indicators['SMH']['RSI']['Period-10'].Current.Value:
                        self.SetHoldings("QQQ", 0.5)
                    else:
                        self.SetHoldings("SMH", 0.5)
        else:
        # If the current price of QQQ is below its 20-day SMA price
            if self.Securities['QQQ'].Close < self.indicators['QQQ']['SMA']['Period-20'].Current.Value:
            # If the 10-day RSI of TLT is above the 10-day RSI of SQQQ
                if self.indicators['TLT']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
                # Buy TQQQ
                    self.SetHoldings("TQQQ", 0.5)
                else:
                # Buy SQQQ
                    self.SetHoldings("SQQQ", 0.5)
            else:
            # If the 10-day RSI of SQQQ is below 31
                if self.indicators['SQQQ']['RSI']['Period-10'].Current.Value < 31:
                # Buy SQQQ
                    self.SetHoldings("SQQQ", 0.5)
                else:
                # If the 10-day return of QQQ is above 5.5%
                    if self.CalculateReturn('QQQ', 10) > 0.055:
                    # Buy SQQQ
                        self.SetHoldings("SQQQ", 0.5)
                    else:
                    # Choose one asset from the list [TQQQ, SOXL] with the highest 10-day RSI
                        if self.indicators['TQQQ']['RSI']['Period-10'].Current.Value > self.indicators['SOXL']['RSI']['Period-10'].Current.Value:
                            self.SetHoldings("TQQQ", 0.5)
                        else:
                            self.SetHoldings("SOXL", 0.5)
    def Group2(self):
        # Group 2: If the current price of QQQ is below its 20-day SMA price
        if self.Securities['QQQ'].Close < self.indicators['QQQ']['SMA']['Period-20'].Current.Value:
            # If the 60-day return of QQQ is below -12%
            if self.CalculateReturn('QQQ', 60) < -0.12:
            # Allocate 50% of the portfolio to Group 3 and 50% to Group 4
                self.Group3()
                self.Group4()
            else:
                # If the 10-day RSI of TLT is above the 10-day RSI of SQQQ
                if self.indicators['TLT']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
                # Buy TQQQ
                    self.SetHoldings("TQQQ", 0.5)
                else:
                # Buy SQQQ
                    self.SetHoldings("SQQQ", 0.5)
        else:
            # If the 10-day RSI of SQQQ is below 31
            if self.indicators['SQQQ']['RSI']['Period-10'].Current.Value < 31:
            # Buy QQQ
                self.SetHoldings("QQQ", 0.5)
            else:
            # If the 70-day return of QQQ is below -15%
                if self.CalculateReturn('QQQ', 70) < -0.15:
                # Choose one asset from the list [TQQQ, SOXL] with the highest 10-day RSI
                    if self.indicators['TQQQ']['RSI']['Period-10'].Current.Value > self.indicators['SOXL']['RSI']['Period-10'].Current.Value:
                        self.SetHoldings("TQQQ", 0.5)
                    else:
                        self.SetHoldings("SOXL", 0.5)
                else:
                # Choose two assets from the list [SPY, QQQ, DIA, XLP] with the highest 15 day return and weigh the picks equally
                    equities = ["SPY", "QQQ", "DIA", "XLP"]
                    returns = {}

                    for equity in equities:
                        returns[equity] = self.CalculateReturn(equity, 15)

                    sorted_returns = sorted(returns.items(), key=lambda x: x[1], reverse=True)
                    top_2_equities = sorted_returns[:2]

                    self.SetHoldings(top_2_equities[0][0], 0.25)
                    self.SetHoldings(top_2_equities[1][0], 0.25)
    def Group3(self):
        # If the current price of SPY is above its 20-day SMA price
        if self.Securities['SPY'].Close > self.indicators['SPY']['SMA']['Period-20'].Current.Value:
            # Buy SPY
            self.SetHoldings("SPY", 0.25)
        else:
            # If the 10-day RSI of TLT is above the 10-day RSI of SQQQ
            if self.indicators['TLT']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
                # Buy QQQ
                self.SetHoldings("QQQ", 0.25)
            else:
                # Buy PSQ
                self.SetHoldings("PSQ", 0.25)

    def Group4(self):
        # If the 10-day RSI of TLT is above the 10-day RSI of SQQQ
        if self.indicators['TLT']['RSI']['Period-10'].Current.Value > self.indicators['SQQQ']['RSI']['Period-10'].Current.Value:
            # Buy QQQ
            self.SetHoldings("QQQ", 0.25)
        else:
            # Buy PSQ
            self.SetHoldings("PSQ", 0.25)