#region imports
from AlgorithmImports import *
#endregion       

class SectorMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        self.SetEndDate(2024, 5, 31)
        self.SetCash(1000000)

        # Set the benchmark to SPY
        self.SetBenchmark("SPY")

        # BrokerageModel --> Here we select IB (this model includes fees)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Define parameters
        self.cmo_window = int(self.GetParameter("CMO window"))
        self.sma_s = int(self.GetParameter("small SMA window"))
        self.sma_l = int(self.GetParameter("long SMA window"))
        self.roc_period = int(self.GetParameter("roc_period"))
        self.selected_symbol_count = int(self.GetParameter("selected_symbol_count")) 

        # CMO indicator for SPY
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.cmo = self.CMO(self.spy, self.cmo_window, Resolution.Daily)
        self.cmo.Updated += self.OnCMOUpdated
        self.cmoWindow = RollingWindow[IndicatorDataPoint](self.cmo_window)
        
        # Define the universe of symbols
        self.long_universe = [
            "XLK", "XLE", "XLV", "XLF", "XLI", "XLB", "XLY", "XLP", "XLU", "XLC",
            "XLRE", "URA", "OIH", "XOP", "CRAK", "ICLN", "TAN", "PBW",
            "VLUE", "INFL", "FTGC", "GNR", "XME", "ARKK", "QCLN", "USO", "SLV",
            "GLD", "LIT", "HACK", "SKYY", "CLOU", "IBB", "KWEB", "ASHR", "FXI",
            "VGK", "EWJ", "EEM", "SOXX", "KBE", "KIE", "GDX", "SIL", "ASHS",
            "METV", "BITQ", "CRPT", "TLT", "XRT", "HYG", "LQD", "MOO", "SPHB",
            "SPLV", "IYT", "ITA", "CGW", "NSPY", "FEZ", "DAX", "FNGS", "KRBN",
            "GSG", "MEME", "JETS", "INDA", "UUP", "KMET", "NRJ", "IPAY", "AIQ",
            "HERO"
        ]

        # Initialize indicators and rolling windows for each symbol in the universe
        self.data = {}
        for ticker in self.long_universe:
            equity = self.AddEquity(ticker, Resolution.Daily)
            roc = RateOfChange(self.roc_period)
            sma_short = SimpleMovingAverage(self.sma_s)
            sma_long = SimpleMovingAverage(self.sma_l)
            self.data[ticker] = {
                "roc": roc,
                "sma_short": sma_short,
                "sma_long": sma_long,
                "rocWindow": RollingWindow[IndicatorDataPoint](self.roc_period),
                "smaShortWindow": RollingWindow[IndicatorDataPoint](self.sma_s),
                "smaLongWindow": RollingWindow[IndicatorDataPoint](self.sma_l)
            }
            roc.Updated += lambda sender, updated, symbol=ticker: self.OnROCUpdated(sender, updated, symbol)
            sma_short.Updated += lambda sender, updated, symbol=ticker: self.OnSMAShortUpdated(sender, updated, symbol)
            sma_long.Updated += lambda sender, updated, symbol=ticker: self.OnSMALongUpdated(sender, updated, symbol)

        # Schedule the rebalance function 6 and 5 minutes before market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 6), self.UpdateIndicators)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 5), self.Rebalance)

        # Set warmup period
        self.SetWarmUp(timedelta(days=max(self.sma_l, self.roc_period)))


    def OnCMOUpdated(self, sender, updated):
        if self.cmo.IsReady:
            self.cmoWindow.Add(updated)
            #self.Debug(f"CMO Value at {self.Time}: {self.cmo.Current.Value}")

    def OnROCUpdated(self, sender, updated, symbol):
        if self.data[symbol]["roc"].IsReady:
            self.data[symbol]["rocWindow"].Add(updated)

    def OnSMAShortUpdated(self, sender, updated, symbol):
        if self.data[symbol]["sma_short"].IsReady:
            self.data[symbol]["smaShortWindow"].Add(updated)

    def OnSMALongUpdated(self, sender, updated, symbol):
        if self.data[symbol]["sma_long"].IsReady:
            self.data[symbol]["smaLongWindow"].Add(updated)

    def UpdateIndicators(self):
        # Manually update CMO with historical close prices
        history = self.History(["SPY"], self.cmo_window, Resolution.Daily)
        for time, row in history.loc["SPY"].iterrows():
            self.cmo.Update(time, row['close'])
        
        # Update CMO with the current price at 5 minutes before close
        self.cmo.Update(self.Time, self.Securities["SPY"].Price)

        # Update indicators for each symbol in the universe with the current price
        for ticker, indicators in self.data.items():
            current_price = self.Securities[ticker].Price
            indicators["roc"].Update(self.Time, current_price)
            indicators["sma_short"].Update(self.Time, current_price)
            indicators["sma_long"].Update(self.Time, current_price)

    def Rebalance(self):
        if self.IsWarmingUp:
            return

        if not self.cmo.IsReady or self.cmo.Current.Value < 0:
            self.Liquidate()
            return

        # Filter the long universe based on momentum and SMA delta
        sorted_by_momentum = sorted(
            [ticker for ticker in self.long_universe if self.data[ticker]["roc"].IsReady and
             self.data[ticker]["sma_short"].IsReady and self.data[ticker]["sma_long"].IsReady and
             (self.data[ticker]["sma_short"].Current.Value - self.data[ticker]["sma_long"].Current.Value) >= 0],
            key=lambda ticker: self.data[ticker]["roc"].Current.Value,
            reverse=True
        )

        if len(sorted_by_momentum) < self.selected_symbol_count:
            self.Liquidate()
            return

        long_positions = sorted_by_momentum[:self.selected_symbol_count]

        # Adjust holdings
        for holding in self.Portfolio.Values:
            if holding.Invested and holding.Symbol.Value not in long_positions:
                self.Liquidate(holding.Symbol)

        weight = 1 / len(long_positions)
        for ticker in long_positions:
            self.SetHoldings(ticker, weight)

    def OnData(self, data: Slice):
        if self.IsWarmingUp:
            return
        pass

# https://quantpedia.com/strategies/sector-momentum-rotational-system/
#
# Use ten sector ETFs. Pick 4 ETFs with the strongest 12-week momentum into your portfolio and weight them equally. Hold them for one week and then rebalance.

#region imports
from AlgorithmImports import *
#endregion       

class SectorMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  
        self.SetEndDate(2024, 5, 31)
        self.SetCash(1000000)

        # Set the benchmark to SPY
        self.SetBenchmark("SPY")

        # BrokerageModel --> Here we select IB (this model includes fees)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # CMO indicator
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.cmo = self.CMO(self.spy, int(self.GetParameter("CMO window")))
        
        # daily data
        self.data = {}
        self.sma_s = int(self.GetParameter("small SMA window"))
        self.sma_l = int(self.GetParameter("long SMA window"))
        self.roc_period = int(self.GetParameter("roc_period"))
        if self.roc_period >= self.sma_l: 
            self.SetWarmUp(self.roc_period)
        else: 
            self.SetWarmUp(self.sma_l)
        self.selected_symbol_count = int(self.GetParameter("selected_symbol_count")) # long symbol count
        self.long_universe = [
            "XLK",  # Technology Select Sector SPDR Fund
            "XLE",  # Energy Select Sector SPDR Fund
            "XLV",  # Health Care Select Sector SPDR Fund
            "XLF",  # Financial Select Sector SPDR Fund
            "XLI",  # Industrials Select Sector SPDR Fund
            "XLB",  # Materials Select Sector SPDR Fund
            "XLY",  # Consumer Discretionary Select Sector SPDR Fund
            "XLP",  # Consumer Staples Select Sector SPDR Fund
            "XLU",  # Utilities Select Sector SPDR Fund
            "XLC",  # Communication Services Select Sector SPDR Fund
            "XLRE", # Real Estate Select Sector SPDR Fund
            "URA",
            "OIH",
            "XLE",
            "XOP",
            "CRAK",
            "ICLN",
            "TAN",
            "PBW",
            "VLUE",
            "INFL",
            "FTGC",
            "GNR",
            "XME",
            "ARKK",
            "QCLN",
            "USO",
            "SLV",
            "GLD",
            "LIT",
            "HACK",
            "SKYY",
            "CLOU",
            "IBB",
            "KWEB",
            "ASHR",
            "FXI",
            "VGK",
            "EWJ",
            "EEM",
            "SOXX",
            "KBE",
            "KIE",
            "GDX",
            "SIL",
            "ASHS",
            "METV",
            "BITQ",
            "CRPT",
            "TLT",
            "XRT",
            "HYG",
            "LQD",
            "MOO",
            "SPHB",
            "SPLV",
            "IYT",
            "ITA",
            "CGW",
            "NSPY",
            "FEZ",
            "DAX",
            "FNGS",
            "KRBN",
            "GSG",
            "MEME",
            "JETS",
            "INDA",
            "UUP",
            "KMET",
            "NRJ",
            "IPAY",
            "AIQ",
            "HERO",
        ]

        for ticker in self.long_universe:
            data = self.AddEquity(ticker, Resolution.Daily)
            data.SetLeverage(1)
            self.data[ticker] = {
                "roc": self.ROC(ticker, self.roc_period, Resolution.Daily),
                "sma_short": self.SMA(ticker, self.sma_s, Resolution.Daily),
                "sma_long": self.SMA(ticker, self.sma_l, Resolution.Daily),
            }
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 5), self.Rebalance)

    def Rebalance(self):
        if self.IsWarmingUp: return
        
        # Check if CMO is ready and its value
        if self.cmo.IsReady and self.cmo.Current.Value < 0:
            self.Liquidate()
            return

        # Filter the long universe based on momentum and SMA delta
        sorted_by_momentum = sorted(
            [ticker for ticker in self.long_universe if self.data[ticker]["roc"].IsReady and 
             self.data[ticker]["sma_short"].IsReady and self.data[ticker]["sma_long"].IsReady and 
             (self.data[ticker]["sma_short"].Current.Value - self.data[ticker]["sma_long"].Current.Value) >= 0], 
            key=lambda ticker: self.data[ticker]["roc"].Current.Value, 
            reverse=True
        )

        if len(sorted_by_momentum) < self.selected_symbol_count:
            self.Liquidate()
            return

        long_positions = sorted_by_momentum[:self.selected_symbol_count]
        
        # Adjust holdings
        for holding in self.Portfolio.Values:
            if holding.Invested and holding.Symbol.Value not in long_positions:
                self.Liquidate(holding.Symbol)

        weight = 1 / len(long_positions)
        for ticker in long_positions:
            self.SetHoldings(ticker, weight)

    def OnData(self, data: Slice):
        if self.IsWarmingUp: return
        pass

#region imports
from AlgorithmImports import *
#endregion
import plotly.express as px
import plotly.graph_objects as go

# Your New Python File

#Plot stuff
equity_chart = backtest.Charts["Strategy Equity"]
drawdown_chart = backtest.Charts["Drawdown"]
benchmark_chart = backtest.Charts["Benchmark"]

equity = equity_chart.Series["Equity"].Values
drawdown = drawdown_chart.Series["Equity Drawdown"].Values
benchmark = benchmark_chart.Series["Benchmark"].Values

df = pd.DataFrame({
    "Equity": pd.Series({value.Time: value.Close for value in equity}),
    "Drawdown": pd.Series({value.Time: value.Y for value in drawdown}),
    "Benchmark": pd.Series({value.Time: value.Y for value in benchmark})
}).ffill()

# Create subplots to plot series on same/different plots
fig, ax = plt.subplots(2, 1, figsize=(12, 12), sharex=True, gridspec_kw={'height_ratios': [2, 1]})

# Plot the equity curve
ax[0].plot(df.index, df["Equity"])
ax[0].set_title("Strategy Equity Curve")
ax[0].set_ylabel("Portfolio Value ($)")

# Plot the benchmark on the same plot, scale by using another y-axis
ax2 = ax[0].twinx()
ax2.plot(df.index, df["Benchmark"], color="grey")
ax2.set_ylabel("Benchmark Price ($)", color="grey")