#region imports
from AlgorithmImports import *

class ParticleQuantumChamber(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 1, 1)  
        # self.SetEndDate(2020, 6, 15)
        self.SetCash(10000)

        self.upro = self.AddEquity('UPRO', Resolution.Daily).Symbol  # UPRO = 3x levered SPX
        self.shy = self.AddEquity('SHY', Resolution.Daily).Symbol  # SHY = short term Treasury ETF
        self.spy = self.AddEquity('SPY', Resolution.Daily).Symbol  # SPY = S&P 500 ETF
        self.uvxy = self.AddEquity('UVXY', Resolution.Daily).Symbol  # 
        self.tmf = self.AddEquity('TMF', Resolution.Daily).Symbol  # 
        self.vix = self.AddEquity('VIX', Resolution.Daily).Symbol  # 
        self.tecl = self.AddEquity('TECL', Resolution.Daily).Symbol  # 
        self.sh = self.AddEquity('SH', Resolution.Daily).Symbol  # 
        self.tlt = self.AddEquity('TLT', Resolution.Daily).Symbol  # 
        self.spxu = self.AddEquity('SPXU', Resolution.Daily).Symbol  # 
        self.sphb = self.AddEquity('SPHB', Resolution.Daily).Symbol
        self.bsv = self.AddEquity('BSV', Resolution.Daily).Symbol
        self.qqq = self.AddEquity('QQQ', Resolution.Daily).Symbol  # 
        self.iei = self.AddEquity('IEI', Resolution.Daily).Symbol  # 
        
        # Register to receive data for the 'SPY' symbol
        # self.AddEquity('SPY', Resolution.Daily)
        self.SetWarmup(200)
        self.investpercent = 1
        self.SetBenchmark(self.spy)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        # self.SetAlpha(EtfSmaAlphaModel(self.spy, self.upro, self.shy))
        # self.SetUniverseSelection(ManualUniverseSelectionModel([self.upro, self.shy, self.spy]))
        self.SetExecution(ImmediateExecutionModel())
        # self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        # Setting the account type in the brokerage model
        # self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash) # 3 days for portfolio to receive cash from sale of a stock
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) # avoids simulating order settlement with margin (or simply don't use a brokerage model but may have inaccurate fees)
        # self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) # all insights are equal weight
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel()) # insights' weighting determines portfolio allocation

        # Investigate backing out and resetting after 2 weeks?
        # self.SetRiskManagement(MaximumDrawdownPercentPortfolio(0.1))  # Run to cash at 10% drawdown
        # self.riskedSecurities = set()
        # self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.1))  # 10% drawdown
        # self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Friday), self.TimeRules.At(9, 30), self.riskedSecurities.clear())

        # Create a custom chart
        stockPlot = Chart("Stock Plot")
        stockPlot.AddSeries(Series("SPY", SeriesType.Line))
        stockPlot.AddSeries(Series("UPRO", SeriesType.Line))
        stockPlot.AddSeries(Series("SHY", SeriesType.Line))
        stockPlot.AddSeries(Series("200d EMA", SeriesType.Line))  # Add 200d EMA series
        stockPlot.AddSeries(Series("10d RSI", SeriesType.Line))
        stockPlot.AddSeries(Series("60d RSI", SeriesType.Line))
        stockPlot.AddSeries(Series("Cash", SeriesType.Line))
        stockPlot.AddSeries(Series("Unsettled Cash", SeriesType.Line))
        self.AddChart(stockPlot)

        # Initialize EMA
        self.ema30spy = self.EMA(self.spy, 30, Resolution.Daily)
        self.ema200spy = self.EMA(self.spy, 200, Resolution.Daily)
        self.rsi10spy = self.RSI(self.spy, 10, MovingAverageType.Simple, Resolution.Daily)
        self.rsi10qqq = self.RSI(self.qqq, 10, MovingAverageType.Simple, Resolution.Daily)
        self.rsi10tlt = self.RSI(self.tlt, 10, MovingAverageType.Simple, Resolution.Daily)
        self.rsi10sh = self.RSI(self.sh, 10, MovingAverageType.Simple, Resolution.Daily)
        self.rsi10spxu = self.RSI(self.spxu, 10, MovingAverageType.Simple, Resolution.Daily)

        self.rsi60spy = self.RSI(self.spy, 60, MovingAverageType.Simple, Resolution.Daily)
        self.rsi60tlt = self.RSI(self.tlt, 60, MovingAverageType.Simple, Resolution.Daily)
        self.rsi14iei = self.RSI(self.iei, 14, MovingAverageType.Simple, Resolution.Daily)
        self.rsi14sphb = self.RSI(self.sphb, 14, MovingAverageType.Simple, Resolution.Daily)
        self.rsi7bsv = self.RSI(self.bsv, 7, MovingAverageType.Simple, Resolution.Daily)
        self.rsi7sphb = self.RSI(self.sphb, 7, MovingAverageType.Simple, Resolution.Daily)

    def liquidate_all_positions_except(self, symbol_to_keep):
        for holding in self.Portfolio.Values:
            if holding.Symbol != symbol_to_keep and holding.Invested:
                self.Liquidate(holding.Symbol)

    def OnData(self, data):

        if self.IsWarmingUp:
            return

        # if no spy data, skip
        if not data.ContainsKey(self.spy):
            return
        if not data.ContainsKey(self.uvxy):
            return
        

        if self.spy in data and data[self.spy] is not None:
            self.Plot("Stock Plot", "SPY", data[self.spy].Close)
            self.Plot("Stock Plot", "200d EMA", self.ema200spy.Current.Value)  # Plot 200d EMA
            self.Plot("Stock Plot", "10d RSI", self.rsi10spy.Current.Value)
            self.Plot("Stock Plot", "60d RSI", self.rsi60spy.Current.Value)
        if self.upro in data and data[self.upro] is not None:
            self.Plot("Stock Plot", "UPRO", data[self.upro].Close)
        if self.shy in data and data[self.shy] is not None:
            self.Plot("Stock Plot", "SHY", data[self.shy].Close)

        self.Plot("Stock Plot", "Cash", self.Portfolio.Cash)
        self.Plot("Stock Plot", "Unsettled Cash", self.Portfolio.UnsettledCash)


        self.invested_value = sum([x.HoldingsValue for x in self.Portfolio.Values]) # Calculate the currently invested value
        self.percent_invested = (self.invested_value / self.Portfolio.TotalPortfolioValue) * 100 # Calculate the percentage invested

        # allocation = self.Portfolio[stock_symbol].Percentage
        # self.Log(f"{stock_symbol} allocation: {allocation * 100:.2f}%")

        # Plotting Cash Amounts
        self.Plot("Portfolio Info", "Cash", self.Portfolio.Cash)
        self.Plot("Portfolio Info", "Unsettled Cash", self.Portfolio.UnsettledCash)
        self.Plot("Portfolio Info", "Percent Invested", self.percent_invested)

         # Set timedelta as an instance variable
        self.timedelta = timedelta(days=1)
        insights = []

        if data[self.spy].Close > self.ema200spy.Current.Value:
            if self.rsi10spy.Current.Value > 80:  # 
                # invest in uxvy
                insights.append(Insight.Price(self.uvxy, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.uvxy}")

            else:
                if self.rsi60spy.Current.Value > 60:  # 
                    # invest in tmf or vixy (add inv volatility here later)
                    ########## ADD VOLATILITY LATER ###########
                    #### TMF VS VIXY HERE
                    insights.append(Insight.Price(self.tmf, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                    self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.tmf}")

                else:
                    # invest in upro
                    insights.append(Insight.Price(self.upro, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                    self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.upro}")

        else: # spy close <= ema
            if self.rsi10qqq.Current.Value < 30:
                # set to tecl
                insights.append(Insight.Price(self.tecl, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.tecl}")
                    
            else: 
                if self.rsi60tlt.Current.Value < 45:
                    #### HALF IN THIS BLOCK ####
                    if data[self.spy].Close > self.ema30spy.Current.Value:
                        if self.rsi10spy.Current.Value > 70:
                            # set to SH
                            insights.append(Insight.Price(self.sh, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.sh}")
                        else:
                            # set to upro
                            insights.append(Insight.Price(self.upro, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.upro}")

                    else: # between tlt and sh 10drsi select top1
                        if self.rsi10tlt.Current.Value > self.rsi10sh.Current.Value: 
                            #invest in tlt
                            insights.append(Insight.Price(self.tlt, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.tlt}")
                        else:
                            insights.append(Insight.Price(self.sh, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.sh}")
                    #### END HALF IN THIS BLOCK ####

                    #### HALF IN THIS BLOCK ####
                    if self.rsi14iei.Current.Value < self.rsi14sphb.Current.Value:
                        if self.rsi7bsv.Current.Value < self.rsi7sphb.Current.Value:
                            # invest in sh
                            insights.append(Insight.Price(self.sh, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.sh}")
                        else:
                            #invest in upro
                            insights.append(Insight.Price(self.upro, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                            self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.upro}")
                    else:
                        #invest in upro
                        insights.append(Insight.Price(self.upro, self.timedelta, InsightDirection.Up, weight=self.investpercent/2))
                        self.Log(f"Insight allocation: {self.investpercent/2 * 100:.2f}% in {self.upro}")
                    #### END HALF IN THIS BLOCK ####

                else: # 
                    if data[self.spy].Close > self.ema30spy.Current.Value:
                        if self.rsi10spy.Current.Value > 70:
                            # invest in sh
                            insights.append(Insight.Price(self.sh, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                            self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.sh}")

                        else: 
                            # invest in upro
                            insights.append(Insight.Price(self.upro, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                            self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.upro}")

                    else: # between tlt and spxu 10drsi select top1
                        if self.rsi10tlt.Current.Value > self.rsi10spxu.Current.Value: #invest in tlt
                            insights.append(Insight.Price(self.tlt, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                            self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.tlt}")

                        else:
                            insights.append(Insight.Price(self.spxu, self.timedelta, InsightDirection.Up, weight=self.investpercent))
                            self.Log(f"Insight allocation: {self.investpercent * 100:.2f}% in {self.spxu}")

        # Group insights and return
        if insights:
            self.EmitInsights(Insight.Group(insights))