from AlgorithmImports import *

class IntelligentSkyRodent(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 2)  # Set Start Date
        self.SetEndDate(2023, 1, 11)
        self.SetCash(100000)  # Set Strategy Cash
        # self.SetSecurityInitializer(self.CustomSecurityInitializer)

        # Define a list of equities to be used by the algorithm
        self.equities = ["SPY", "QQQ"]
        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)

            # 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", 3),
                self.FunctionBeforeMarketClose)
            # 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])


    # 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[-1].Close / history[0].Close) - 1


    def OnDataConsolidated(self, sender, bar):
        pass
    
    def FunctionBeforeMarketClose(self):

        # Iterate over each equity
        for equity in self.equities:

            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)

            self.indicators[equity]['RSI']['Period-5'].Update(self.Time, self.Securities[equity].Close)

        # Implement the main decision tree for the algorithm

            if self.indicators['QQQ']['RSI']['Period-5'].Current.Value < self.indicators['SPY']['RSI']['Period-5'].Current.Value:
                    self.SetHoldings("QQQ", 1)
            else:
                    self.SetHoldings("SPY", 1)