# region imports
from AlgorithmImports import *
import numpy as np
import pandas as pd
# endregion

class MacroFlowAdaptiveVolatility(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 6, 1) 
        self.SetEndDate(2026, 1, 1)
        self.SetCash(100000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # 1. Asset Universe
        self.symbols = {
            "xly": self.AddEquity("XLY", Resolution.Daily).Symbol,
            "xlp": self.AddEquity("XLP", Resolution.Daily).Symbol,
            "hyg": self.AddEquity("HYG", Resolution.Daily).Symbol,
            "tlt": self.AddEquity("TLT", Resolution.Daily).Symbol,
            "cop": self.AddEquity("COPX", Resolution.Daily).Symbol,
            "gld": self.AddEquity("GLD", Resolution.Daily).Symbol,
            "tech": self.AddEquity("XLK", Resolution.Daily).Symbol,
            "energy": self.AddEquity("XLE", Resolution.Daily).Symbol,
            "market": self.AddEquity("RSP", Resolution.Daily).Symbol,
            "cash": self.AddEquity("BIL", Resolution.Daily).Symbol
        }
        self.vix = self.AddData(CBOE, "VIX").Symbol
        
        # 2. Risk & Quantitative Parameters
        self.target_vol = 0.10        
        self.asset_highs = {}         
        self.stopped_out_this_month = False 
        self.vix_sma = self.SMA(self.vix, 5, Resolution.Daily)
        self.lookback = 60
        self.SetWarmUp(self.lookback + 20)

        # 3. Custom Risk Dashboard
        riskPlot = Chart('Strategy Risk Dashboard')
        riskPlot.AddSeries(Series('Realized Volatility', SeriesType.Line, '%', Color.Orange))
        riskPlot.AddSeries(Series('Target Volatility', SeriesType.Line, '%', Color.Blue))
        riskPlot.AddSeries(Series('VIX Level', SeriesType.Line, '', Color.Gray))
        self.AddChart(riskPlot)

        self.Schedule.On(self.DateRules.MonthStart(self.symbols["tech"]), 
                         self.TimeRules.AfterMarketOpen(self.symbols["tech"], 30), 
                         self.Rebalance)
        
        self.Schedule.On(self.DateRules.EveryDay(self.symbols["tech"]),
                         self.TimeRules.BeforeMarketClose(self.symbols["tech"], 15),
                         self.RiskManagement)

    def Rebalance(self):
        if self.IsWarmingUp: return
        self.stopped_out_this_month = False 
        
        # FIX: Ensure ALL symbols are requested in history
        h = self.History(list(self.symbols.values()), self.lookback + 10, Resolution.Daily)
        if h.empty: return
        
        # Unstack for easier access
        prices = h.unstack(level=0)['close']

        # I. Triple-Gate Validation
        consumer_ok = self.GetSmoothedROC(self.symbols["xly"], h) > self.GetSmoothedROC(self.symbols["xlp"], h)
        credit_ok = self.GetSmoothedROC(self.symbols["hyg"], h) > self.GetSmoothedROC(self.symbols["tlt"], h)

        # II. Volatility Analytics
        # Extract returns for growth assets
        growth_rets = prices[[self.symbols["tech"], self.symbols["energy"]]].pct_change().dropna()
        vol_t = growth_rets[self.symbols["tech"]].std() * np.sqrt(252)
        vol_e = growth_rets[self.symbols["energy"]].std() * np.sqrt(252)
        corr = growth_rets[self.symbols["tech"]].corr(growth_rets[self.symbols["energy"]])

        if consumer_ok and credit_ok:
            # III. Growth Allocation (Inverse Volatility / Risk Parity)
            w_t = (1/vol_t) / ((1/vol_t) + (1/vol_e))
            w_e = 1 - w_t
            est_port_vol = np.sqrt((w_t**2 * vol_t**2) + (w_e**2 * vol_e**2) + (2*w_t*w_e*vol_t*vol_e*corr))
            
            exposure = np.clip(self.target_vol / est_port_vol, 0, 1.2)
            
            self.SetHoldings(self.symbols["tech"], w_t * exposure)
            self.SetHoldings(self.symbols["energy"], w_e * exposure)
            self.SetHoldings(self.symbols["market"], 0.10)
            self.Liquidate(self.symbols["gld"])
            
            self.PlotRisk(est_port_vol)
        else:
            # IV. Defensive State (Gold + Cash)
            # FIX: Pull Gold volatility safely from the prices dataframe
            gold_rets = prices[self.symbols["gld"]].pct_change().dropna()
            vol_g = gold_rets.std() * np.sqrt(252)
            
            self.SetHoldings(self.symbols["gld"], 0.40)
            self.SetHoldings(self.symbols["cash"], 0.60)
            self.Liquidate(self.symbols["tech"])
            self.Liquidate(self.symbols["energy"])
            
            self.PlotRisk(vol_g)

        # Update Highs
        for symbol in self.Portfolio.Keys:
            if self.Portfolio[symbol].Invested:
                self.asset_highs[symbol] = self.Securities[symbol].Price

    def RiskManagement(self):
        vix_price = self.Securities[self.vix].Price
        dynamic_stop = 0.12 if vix_price < 15 else (0.04 if vix_price > 25 else 0.08)

        for symbol in list(self.asset_highs.keys()):
            if not self.Portfolio[symbol].Invested: continue
            
            price = self.Securities[symbol].Price
            if price > self.asset_highs[symbol]: self.asset_highs[symbol] = price
            
            if price < self.asset_highs[symbol] * (1 - dynamic_stop):
                self.Log(f"STOP TRIGGERED: {symbol}. Moving to Cash.")
                self.Liquidate(symbol)
                self.stopped_out_this_month = True
                self.SetHoldings(self.symbols["cash"], 0.95)

        if self.stopped_out_this_month and vix_price < self.vix_sma.Current.Value * 0.90:
            self.Log("RE-ENTRY: Fear collapsed. Restarting Engine.")
            self.Rebalance()

    def GetSmoothedROC(self, symbol, history):
        prices = history.loc[symbol]['close']
        roc = (prices / prices.shift(21)) - 1
        return roc.rolling(5).mean().iloc[-1]

    def PlotRisk(self, current_vol):
        self.Plot('Strategy Risk Dashboard', 'Realized Volatility', current_vol * 100)
        self.Plot('Strategy Risk Dashboard', 'Target Volatility', self.target_vol * 100)
        self.Plot('Strategy Risk Dashboard', 'VIX Level', self.Securities[self.vix].Price)