from AlgorithmImports import *
import numpy as np

class VIXSpikeDefenseStrategyWithFedRate(QCAlgorithm):

    def Initialize(self):
        self.SetCash(1000000)

        # Backtesting Period Selection
        self.TrainingPeriod = "IS"  # Set this to your desired period ("IS", "OOSA", "OOSB", "LT", "ST")
        self.SetBacktestingPeriod()

        # Warm-up the algorithm for 252 days 
        self.SetWarmUp(252, Resolution.Daily)

        # Add assets
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.tlt = self.AddEquity("TLT", Resolution.Daily).Symbol
        self.iei = self.AddEquity("IEI", Resolution.Daily).Symbol
        self.gld = self.AddEquity("GLD", Resolution.Daily).Symbol
        self.dbc = self.AddEquity("DBC", Resolution.Daily).Symbol
        self.qqq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.ieo = self.AddEquity("IEO", Resolution.Daily).Symbol
        # Shares U.S. Oil & Gas Exploration & Production ETF (IEO) - This ETF 
        self.shv = self.AddEquity("SHV", Resolution.Daily).Symbol
        self.sh = self.AddEquity("SH", Resolution.Daily).Symbol

        # Add new ETFs for the Fed Rate Strategy
        self.gsg = self.AddEquity("GSG", Resolution.Daily).Symbol
        self.xme = self.AddEquity("XME", Resolution.Daily).Symbol
        self.moo = self.AddEquity("MOO", Resolution.Daily).Symbol
        self.pall = self.AddEquity("PALL", Resolution.Daily).Symbol

        # Add VIX index
        self.vix = self.AddData(CBOE, "VIX").Symbol

        # Add Fed Funds Rate data from FRED
        self.fed_rate = self.AddData(Fred, "DFF").Symbol

        # Parameters for Z-Score calculation
        self.lookback_period = 252  # 252-day rolling window
        self.z_score_threshold = 3.0  # Z-score threshold for detecting a spike

        # Rolling window to store VIX data for calculation
        self.vix_window = RollingWindow[float](self.lookback_period)

        # Store the previous Fed Funds Rate
        self.previous_fed_rate = None

        # Initial allocation (before any spikes)
        self.initial_allocation = {
            self.spy: 0.20,
            self.tlt: 0.30,
            self.iei: 0.15,
            self.gld: 0.075,
            self.dbc: 0.075,
            self.ieo: 0.10,
            self.qqq: 0.10,
            self.shv: 0.00,
            self.sh: 0.00,
            self.gsg: 0.00,
            self.xme: 0.00,
            self.moo: 0.00,
            self.pall: 0.00,
        }

        # Flag to track if the Fed Rate strategy has been executed this month
        self.fed_rate_strategy_executed = False

        # Set the combined strategy execution after the warm-up
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.AfterMarketOpen(self.spy), 
                         self.ExecuteCombinedStrategy)

    def SetBacktestingPeriod(self):
        """Sets the backtesting period based on the selected training period."""
        if self.TrainingPeriod == "IS":
            self.SetStartDate(2017, 1, 1)
            self.SetEndDate(2021, 1, 1)
        elif self.TrainingPeriod == "OOSA": 
            self.SetStartDate(2022, 1, 1)
            self.SetEndDate(2022, 11, 1)
        elif self.TrainingPeriod == "OOSB":
            self.SetStartDate(2016, 1, 1)
            self.SetEndDate(2017, 1, 1)
        elif self.TrainingPeriod == "LT":
            self.SetStartDate(2024, 8, 1)
            self.SetEndDate(2024, 8, 15)
        elif self.TrainingPeriod == "ST":
            self.SetStartDate(2020, 3, 1)
            self.SetEndDate(2020, 3, 31)
        elif self.TrainingPeriod == "all":
            self.SetStartDate(2010, 1, 1)
            self.SetEndDate(2024, 8, 1)
        elif self.TrainingPeriod == "test":
            self.SetStartDate(2023, 2, 1)
            self.SetEndDate(2023, 4, 1)

    def OnData(self, data: Slice):
        """Handle data feed updates."""
        if self.IsWarmingUp:
            return

        # Ensure all symbols have data
        if not all([self.Securities[symbol].HasData for symbol in self.initial_allocation.keys()]):
            return

        # Perform combined strategy check and adjust allocations accordingly
        self.ExecuteCombinedStrategy()
        self.plot_pnl_and_prices()

    def ExecuteCombinedStrategy(self):
        """Check Fed Funds Rate strategy first, then VIX strategy if the first condition isn't met."""
        if not self.portfolio.invested:
            self.ExecuteVIXStrategy()
        elif self.ShouldExecuteFedRateStrategy():
            self.ExecuteFedRateStrategy()
            self.fed_rate_strategy_executed = True  # Mark that the strategy has been executed for this month
        elif not self.fed_rate_strategy_executed:
            self.ExecuteVIXStrategy()

        

    def ShouldExecuteFedRateStrategy(self):
        """Determine if the Fed Funds Rate strategy should be executed."""
        # Only execute on the first trading day of the month
        if self.Time.day == 1:
            current_rate = self.Securities[self.fed_rate].Price

            # Get the rate from 45 days ago
            history = self.History(self.fed_rate, 45, Resolution.Daily)
            if not history.empty:
                last_month_rate = history['value'].iloc[0]
                # for index, row in history.iterrows():
                #     self.Debug(f"Date: {index}, Fed Rate Value: {row['value']}")
            else:
                self.Debug("Insufficient data to calculate last month's rate. Fed Ex=FALSE")
                self.fed_rate_strategy_executed = False
                return False  # Insufficient data, do not execute the strategy

            rate_increase = current_rate - last_month_rate

            # If the rate has increased by more than the threshold (e.g., 0.25%)
            if rate_increase >= 0.25:
                self.Debug(f"Fed Funds Rate increased by {rate_increase:.2f}% on {self.Time.date()}, Fed EX=TRUE")
                return True
            else:
                self.Debug(f"No significant increase ({rate_increase:.2f}%) on {self.Time.date()} in Fed Funds Rate. Fed Ex=FALSE")
                self.fed_rate_strategy_executed = False
                return False
        return False

    def ExecuteFedRateStrategy(self):
        """Execute the Fed Funds Rate strategy."""
        self.Debug("Executing Fed Rate Strategy")
        self.SetHoldings(self.spy, 0)  #
        self.SetHoldings(self.tlt, 0)  # TLT
        self.SetHoldings(self.iei, 0.15)  # Increase IEI to 10%
        self.SetHoldings(self.gld, 0.1)  # Hold 5% Gold (GLD)
        self.SetHoldings(self.dbc, 0.1)
        self.SetHoldings(self.ieo, 0.2)  # 
        self.SetHoldings(self.qqq, 0.15)  # Increase QQQ to 25%
        self.SetHoldings(self.shv, 0.10)
        self.SetHoldings(self.sh, 0.1)
        self.SetHoldings(self.gsg, 0)  # gsg like dbc
        self.SetHoldings(self.xme, 0)  # SPDR S&P Metals & Mining ETF
        self.SetHoldings(self.moo, 0.1)  
        self.SetHoldings(self.pall, 0) 
        # self.initial_allocation = {
        # IEO QQQ SHV IEI GLD MOO DBC SH
        # }

    def ExecuteVIXStrategy(self):
        """Execute the VIX Spike Defense Strategy."""
        # self.Debug("Executing VIX Strategy")
        self.CheckVIXSpike()

    def SetInitialAllocation(self):
        """Sets the initial portfolio allocation."""
        for symbol, weight in self.initial_allocation.items():
            self.SetHoldings(symbol, weight)

    def CheckVIXSpike(self):
        """Check for VIX spike using Z-Score method and plot data."""
        if self.vix_window.IsReady:
            # Calculate moving average and standard deviation
            mean_vix = sum(self.vix_window) / self.lookback_period
            std_vix = np.std([x for x in self.vix_window])

            # Calculate Z-score
            current_vix = self.Securities[self.vix].Price
            z_score = (current_vix - mean_vix) / std_vix

            # Detect VIX spike
            if z_score > self.z_score_threshold:
                self.Debug(f"VIX Spike Detected: Z-Score {z_score:.2f} on {self.Time.date()}")
                self.DefensiveAllocation()
            else:
                self.SetInitialAllocation()  # Ensure the portfolio returns to the initial allocation if no spike is detected

        # Update the rolling window with the latest VIX value
        self.vix_window.Add(self.Securities[self.vix].Price)

    def DefensiveAllocation(self):
        """Switch to a more defensive portfolio allocation."""
        self.SetHoldings(self.spy, 0)
        self.SetHoldings(self.tlt, 0.417)
        self.SetHoldings(self.iei, 0.213)
        self.SetHoldings(self.gld, 0.157)
        self.SetHoldings(self.dbc, 0)
        self.SetHoldings(self.ieo, 0)
        self.SetHoldings(self.qqq, 0)
        self.SetHoldings(self.shv, 0.126)
        self.SetHoldings(self.sh, 0.087)
        self.SetHoldings(self.gsg, 0)  
        self.SetHoldings(self.xme, 0)  
        self.SetHoldings(self.moo, 0)  
        self.SetHoldings(self.pall, 0) 

    def plot_pnl_and_prices(self):
        """Plot the PnL of the portfolio, VIX, and the price of each ETF."""
        self.Plot("PnL", "Portfolio Value", self.Portfolio.TotalPortfolioValue)
        self.Plot("VIX", "VIX", self.Securities[self.vix].Price)
        self.Plot("SPY Prices", "SPY", self.Securities[self.spy].Price)
        self.Plot("TLT Prices", "TLT", self.Securities[self.tlt].Price)
        self.Plot("IEI Prices", "IEI", self.Securities[self.iei].Price)
        self.Plot("GLD Prices", "GLD", self.Securities[self.gld].Price)
        self.Plot("DBC Prices", "DBC", self.Securities[self.dbc].Price)
        self.Plot("IEO Prices", "IEO", self.Securities[self.ieo].Price)
        self.Plot("QQQ Prices", "QQQ", self.Securities[self.qqq].Price)
        self.Plot("SHV Prices", "SHV", self.Securities[self.shv].Price)
        self.Plot("SH Prices", "SH", self.Securities[self.sh].Price)
        self.Plot("GSG Prices", "GSG", self.Securities[self.gsg].Price)
        self.Plot("XME Prices", "XME", self.Securities[self.xme].Price)
        self.Plot("MOO Prices", "MOO", self.Securities[self.moo].Price)
        self.Plot("PALL Prices", "PALL", self.Securities[self.pall].Price)
        self.Plot("Fed Funds Rate", "Fed Rate", self.Securities[self.fed_rate].Price)