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

class InstitutionalTQQQComposite(QCAlgorithm):

    def initialize(self):
        # 1. Setup: 5 Years (Captures 2020 Crash, 2022 Bear, 2023 Bull)
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2024, 6, 1)
        self.set_cash(100000)
        self.settings.rebalance_portfolio_on_insight_changes = False
        
        # 2. The Engine: TQQQ (3x Nasdaq)
        self.tqqq = self.add_equity("TQQQ", Resolution.MINUTE).symbol
        
        # 3. The Safe Haven: BIL (T-Bills)
        # Instead of 0% Cash, we earn ~5% interest when out of the market
        self.safe = self.add_equity("BIL", Resolution.MINUTE).symbol
        
        # 4. The Canary Signals (Intermarket Analysis)
        # XLI (Industrials) vs XLU (Utilities) -> Economic Cycle
        self.xli = self.add_equity("XLI", Resolution.DAILY).symbol 
        self.xlu = self.add_equity("XLU", Resolution.DAILY).symbol
        
        # 5. Volatility Controls
        self.vol_lookback = 20        # 1 Month Volatility Check
        self.target_vol = 0.40        # Target 40% Annualized Volatility (Aggressive but safe)
        
        # 6. Logic Parameters
        self.lookback = 60            # 3 Months trend comparison
        self.wait_days = 5            # "Cool off" period
        self.out_day_counter = 0 
        self.be_in = False
        
        # 7. Schedule: Trade 10 mins before close to capture daily data
        self.schedule.on(self.date_rules.every_day(self.tqqq), 
                         self.time_rules.before_market_close(self.tqqq, 10), 
                         self.rebalance)
                         
        self.set_warm_up(100)

    def on_data(self, data):
        pass

    def rebalance(self):
        if self.is_warming_up: return
        
        # --- PART 1: REGIME SIGNAL (The Light Switch) ---
        history = self.history([self.xli, self.xlu], self.lookback, Resolution.DAILY)
        if history.empty: return
        
        # Get closing prices
        try:
            xli_closes = history.loc[self.xli].close
            xlu_closes = history.loc[self.xlu].close
        except:
            return # Data missing
            
        if len(xli_closes) < self.lookback or len(xlu_closes) < self.lookback: return
        
        # Calculate Returns (Momentum)
        xli_perf = (xli_closes[-1] / xli_closes[0]) - 1
        xlu_perf = (xlu_closes[-1] / xlu_closes[0]) - 1
        
        # Logic: If Industrials > Utilities, Economy is Healthy
        if xli_perf > xlu_perf:
            self.out_day_counter = 0 
            self.be_in = True
        else:
            self.out_day_counter += 1
            if self.out_day_counter >= self.wait_days:
                self.be_in = False
        
        # --- PART 2: VOLATILITY TARGETING (The Sizing) ---
        # If we are "In", how much do we buy?
        
        target_weight = 0.0
        
        if self.be_in:
            # Calculate TQQQ's recent volatility (Standard Deviation)
            tqqq_hist = self.history(self.tqqq, self.vol_lookback + 1, Resolution.DAILY)
            if not tqqq_hist.empty:
                # Calculate daily returns
                returns = tqqq_hist.close.pct_change().dropna()
                # Calculate Annualized Volatility (StdDev * Sqrt(252))
                current_vol = np.std(returns) * np.sqrt(252)
                
                # Volatility Scaling Math:
                # If Target is 40% and Current is 20%, we go >100% (or capped at 1.0)
                # If Target is 40% and Current is 80%, we go 50% Size.
                if current_vol > 0:
                    vol_weight = self.target_vol / current_vol
                else:
                    vol_weight = 1.0
                
                # Cap leverage at 1.0 (No margin on top of TQQQ)
                target_weight = min(1.0, vol_weight)
                
                self.plot("Data", "TQQQ Vol", current_vol)
                self.plot("Data", "Target Weight", target_weight)
        else:
            # If we are OUT, weight is 0
            target_weight = 0.0

        # --- PART 3: EXECUTION ---
        # 1. Invest calculated weight into TQQQ
        self.set_holdings(self.tqqq, target_weight)
        
        # 2. Invest the REST into BIL (Safe Yield)
        # E.g., if TQQQ is 0.6, BIL is 0.4. If TQQQ is 0, BIL is 1.0.
        # This ensures our "Cash" is actually earning 5% interest
        safe_weight = 1.0 - target_weight
        self.set_holdings(self.safe, safe_weight)
        
        if self.be_in:
            self.plot("State", "Regime", 1)
        else:
            self.plot("State", "Regime", 0)