# region imports
from AlgorithmImports import *
# endregion

"""
9-Sig Strategy for TQQQ
Based on Kelly Letter's 9-Sig Strategy

A rules-based, quarterly rebalancing strategy that uses value averaging to buy low 
and sell high, targeting 9% quarterly growth in TQQQ value.

Core Rules:
- Start with 60/40 allocation (60% TQQQ, 40% bonds)
- Rebalance quarterly against a 9% growth target (signal line)
- Sell excess when above target, buy more when below target
- Skip 2 sell signals after a 30%+ drop to stay invested during recovery
- No intra-quarter trading - discipline over emotion

For more information, visit: https://kellyletter.com
"""

class NineSigTQQQStrategy(QCAlgorithm):
    """9-Sig Strategy for TQQQ with quarterly rebalancing based on 9% growth target."""

    def initialize(self):
        self.set_start_date(2010, 3, 1)
        self.set_cash(100000)
        
        # Add securities
        self._tqqq = self.add_equity("TQQQ", Resolution.DAILY).symbol
        self._bonds = self.add_equity("AGG", Resolution.DAILY).symbol
        
        # Strategy parameters
        self._quarterly_growth_rate = 0.09  # 9% quarterly growth target
        self._initial_tqqq_allocation = 0.60
        self._initial_bond_allocation = 0.40
        
        # State variables
        self._signal_line = 0  # Target TQQQ value
        self._is_initialized = False
        self._peak_tqqq_value = 0  # Track peak for 30% drop detection
        self._sell_signals_to_skip = 0  # Skip sell signals after 30%+ drop
        self._last_rebalance_date = None
        
        # Schedule quarterly rebalancing at start of each quarter
        self.schedule.on(self.date_rules.month_start("TQQQ"), 
                        self.time_rules.after_market_open("TQQQ", 30),
                        self._check_and_rebalance)
        
        # Warmup for price data
        self.set_warmup(1)
        self.settings.free_portfolio_value_percentage = 0.01

    def _check_and_rebalance(self):
        """Check if it's a quarter start and rebalance if needed."""
        if self.is_warming_up:
            return
            
        # Initialize on first call
        if not self._is_initialized:
            self._initialize_positions()
            return
        
        current_month = self.time.month
        
        # Only rebalance at the start of quarters (Jan, Apr, Jul, Oct)
        if current_month not in [1, 4, 7, 10]:
            return
            
        # Prevent double-rebalancing in the same quarter
        if self._last_rebalance_date is not None:
            if self._last_rebalance_date.year == self.time.year and self._last_rebalance_date.month == self.time.month:
                return
        
        self._rebalance()
    
    def _rebalance(self):
        """Quarterly rebalancing based on 9% growth target."""
        
        # Get current TQQQ value
        current_tqqq_value = self.portfolio[self._tqqq].holdings_value
        
        if current_tqqq_value <= 0:
            self.log("No TQQQ holdings, skipping rebalance")
            return
        
        # Check for 30%+ drop from peak
        self._peak_tqqq_value = max(self._peak_tqqq_value, current_tqqq_value)
        drawdown = (self._peak_tqqq_value - current_tqqq_value) / self._peak_tqqq_value
        
        if drawdown >= 0.30 and self._sell_signals_to_skip == 0:
            self._sell_signals_to_skip = 2
            self.log(f"30%+ drawdown detected. Will skip next 2 sell signals.")
        
        # Rebalancing logic
        difference = current_tqqq_value - self._signal_line
        
        if difference > 0:
            # TQQQ above signal line - sell excess
            if self._sell_signals_to_skip > 0:
                self._sell_signals_to_skip -= 1
                self.log(f"Skipping sell signal ({self._sell_signals_to_skip} remaining). Current: ${current_tqqq_value:.2f}, Target: ${self._signal_line:.2f}")
            else:
                # Sell TQQQ, buy bonds
                self._sell_tqqq_buy_bonds(difference)
        else:
            # TQQQ below signal line - buy more
            self._buy_tqqq_from_bonds(abs(difference))
        
        # Reset signal line for next quarter (9% growth)
        self._signal_line = self.portfolio[self._tqqq].holdings_value * (1 + self._quarterly_growth_rate)
        self._last_rebalance_date = self.time
        
        self.log(f"Rebalanced. TQQQ: ${self.portfolio[self._tqqq].holdings_value:.2f}, Bonds: ${self.portfolio[self._bonds].holdings_value:.2f}, New Signal: ${self._signal_line:.2f}")
    
    def _initialize_positions(self):
        """Initial 60/40 allocation."""
        total_value = self.portfolio.total_portfolio_value
        
        self.set_holdings(self._tqqq, self._initial_tqqq_allocation)
        self.set_holdings(self._bonds, self._initial_bond_allocation)
        
        # Set initial signal line
        self._signal_line = self.portfolio[self._tqqq].holdings_value * (1 + self._quarterly_growth_rate)
        self._peak_tqqq_value = self.portfolio[self._tqqq].holdings_value
        self._is_initialized = True
        self._last_rebalance_date = self.time
        
        self.log(f"Initialized: TQQQ ${self.portfolio[self._tqqq].holdings_value:.2f}, Bonds ${self.portfolio[self._bonds].holdings_value:.2f}, Signal ${self._signal_line:.2f}")
    
    def _sell_tqqq_buy_bonds(self, amount: float):
        """Sell TQQQ excess and buy bonds."""
        if amount <= 0:
            return
        
        tqqq_price = self.securities[self._tqqq].price
        bonds_price = self.securities[self._bonds].price
        
        if tqqq_price <= 0 or bonds_price <= 0:
            return
        
        # Sell TQQQ
        tqqq_shares_to_sell = int(amount / tqqq_price)
        if tqqq_shares_to_sell > 0:
            self.market_order(self._tqqq, -tqqq_shares_to_sell)
            
            # Buy bonds with the proceeds
            bond_shares_to_buy = int(amount / bonds_price)
            if bond_shares_to_buy > 0:
                self.market_order(self._bonds, bond_shares_to_buy)
                self.log(f"Sold {tqqq_shares_to_sell} TQQQ, bought {bond_shares_to_buy} AGG (${amount:.2f})")
    
    def _buy_tqqq_from_bonds(self, amount: float):
        """Buy TQQQ using bond allocation."""
        if amount <= 0:
            return
        
        bonds_value = self.portfolio[self._bonds].holdings_value
        if bonds_value <= 0:
            self.log("No bonds available, temporary buy-and-hold")
            return
        
        # Use bond allocation to buy TQQQ
        amount_to_use = min(amount, bonds_value)
        
        bonds_price = self.securities[self._bonds].price
        tqqq_price = self.securities[self._tqqq].price
        
        if bonds_price <= 0 or tqqq_price <= 0:
            return
        
        # Sell bonds
        bond_shares_to_sell = int(amount_to_use / bonds_price)
        if bond_shares_to_sell > 0:
            self.market_order(self._bonds, -bond_shares_to_sell)
            
            # Buy TQQQ with the amount
            tqqq_shares_to_buy = int(amount_to_use / tqqq_price)
            if tqqq_shares_to_buy > 0:
                self.market_order(self._tqqq, tqqq_shares_to_buy)
                self.log(f"Bought {tqqq_shares_to_buy} TQQQ shares (${amount_to_use:.2f}) from bonds")