# region imports
from AlgorithmImports import *
# endregion

class MeasuredRedOrangeBat(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        
        # Warm up to pre-load historical data
        self.set_warm_up(100, Resolution.DAILY)
        
        # Per-position stop loss
        self._stop_loss_pct = 0.02  # 2% of portfolio value loss per position
        
        # Gold for cash allocation
        self._hedge_etf = "GLD"  # SPDR Gold Shares
        self.add_equity(self._hedge_etf)
        
        # Momentum parameters
        self._lookback_period = 90
        self._num_positions = 10
        self._momentum_indicators = {}
        self._volatility_indicators = {}  # Track volatility for position sizing
        
        # Add tech universe
        self._universe = self.add_universe(self._select_tech_stocks)
        
        # Schedule monthly rebalancing on first day of month at market open
        self.schedule.on(self.date_rules.month_start(),
                        self.time_rules.after_market_open("SPY", 1),
                        self._rebalance)
        
        # Check drawdown daily
        self.schedule.on(self.date_rules.every_day(),
                        self.time_rules.after_market_open("SPY", 30),
                        self._check_drawdown)
        
        # Allocate excess cash to bonds daily
        self.schedule.on(self.date_rules.every_day(),
                        self.time_rules.after_market_open("SPY", 31),
                        self._allocate_excess_cash)
    
    def _select_tech_stocks(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        # Filter for technology sector stocks
        tech_stocks = [f for f in fundamentals 
                      if f.has_fundamental_data
                      and f.asset_classification.morningstar_sector_code == MorningstarSectorCode.TECHNOLOGY
                      and f.market_cap > 1e9  # Market cap > $1B
                      and f.volume > 500000]  # Minimum liquidity
        
        # Sort by market cap and take top 20 for selection
        tech_stocks = sorted(tech_stocks, key=lambda f: f.market_cap, reverse=True)[:20]
        
        return [f.symbol for f in tech_stocks]
    
    def on_securities_changed(self, changes: SecurityChanges):
        # Initialize momentum and volatility indicators for new securities
        for security in changes.added_securities:
            symbol = security.symbol
            if symbol not in self._momentum_indicators:
                self._momentum_indicators[symbol] = self.momp(symbol, self._lookback_period, Resolution.DAILY)
                self._volatility_indicators[symbol] = self.std(symbol, 20, Resolution.DAILY)
        
        # Clean up removed securities
        for security in changes.removed_securities:
            symbol = security.symbol
            if symbol in self._momentum_indicators:
                del self._momentum_indicators[symbol]
            if symbol in self._volatility_indicators:
                del self._volatility_indicators[symbol]
    
    def _rebalance(self):
        # Skip if warming up or no universe selections
        if self.is_warming_up or not self._universe.selected:
            return
        
        # Calculate momentum scores for active universe
        momentum_scores = {}
        for symbol in self._universe.selected:
            if (symbol in self._momentum_indicators and 
                self._momentum_indicators[symbol].is_ready and
                symbol in self._volatility_indicators and
                self._volatility_indicators[symbol].is_ready):
                # Filter out negative momentum stocks
                momentum = self._momentum_indicators[symbol].current.value
                if momentum > 0:
                    momentum_scores[symbol] = momentum
        
        # Skip if not enough data
        if len(momentum_scores) < self._num_positions:
            return
        
        # Select top momentum stocks
        sorted_by_momentum = sorted(momentum_scores.items(), key=lambda x: x[1], reverse=True)
        selected_symbols = [symbol for symbol, _ in sorted_by_momentum[:self._num_positions]]
        
        # Volatility-adjusted position sizing
        volatilities = {symbol: self._volatility_indicators[symbol].current.value 
                       for symbol in selected_symbols}
        
        # Inverse volatility weighting
        inv_vols = {symbol: 1.0 / max(vol, 0.01) for symbol, vol in volatilities.items()}
        total_inv_vol = sum(inv_vols.values())
        
        # Calculate weights with 100% allocation
        max_allocation = 1.0
        weights = {symbol: max_allocation * (inv_vol / total_inv_vol) 
                  for symbol, inv_vol in inv_vols.items()}
        
        # Calculate total weight and allocate remainder to leveraged bonds
        total_weight = sum(weights.values())
        remaining_weight = 1.0 - total_weight
        
        # Debug logging
        self.log(f"Date: {self.time} | Momentum stocks selected: {len(selected_symbols)} | Total weight: {total_weight:.2%} | Remaining: {remaining_weight:.2%}")
        
        targets = [PortfolioTarget(symbol, weight) for symbol, weight in weights.items()]
        
        # Add gold allocation if there's remaining weight
        if remaining_weight > 0.01:  # Only allocate if meaningful amount
            targets.append(PortfolioTarget(self.symbol(self._hedge_etf), remaining_weight))
            self.log(f"GOLD ALLOCATION: {remaining_weight:.2%} to {self._hedge_etf}")
        else:
            self.log(f"No gold allocation needed (remaining: {remaining_weight:.2%})")
        
        # Set holdings and liquidate positions not in targets
        self.set_holdings(targets, liquidate_existing_holdings=True)
        
        self.log(f"Rebalanced to {len(selected_symbols)} positions | Portfolio value: ${self.portfolio.total_portfolio_value:,.0f}")
    
    def _check_drawdown(self):
        # Check per-position stop losses based on PnL
        portfolio_value = self.portfolio.total_portfolio_value
        max_loss_per_position = portfolio_value * self._stop_loss_pct
        symbols_to_liquidate = []
        
        # Debug: log current portfolio allocation
        allocated_value = sum(holding.absolute_holdings_value for holding in self.portfolio.values() if holding.quantity != 0)
        cash_value = self.portfolio.cash
        allocation_pct = (allocated_value / portfolio_value * 100) if portfolio_value > 0 else 0
        
        # Check all positions in portfolio
        for symbol in list(self.portfolio.keys()):
            holding = self.portfolio[symbol]
            if holding.quantity > 0:
                # Get unrealized loss in dollars
                unrealized_profit = holding.unrealized_profit
                
                # Liquidate if position loss exceeds 2% of portfolio value
                if unrealized_profit <= -max_loss_per_position:
                    symbols_to_liquidate.append(symbol)
                    self.log(f"STOP LOSS: {symbol} unrealized loss ${abs(unrealized_profit):.2f}, max allowed loss: ${max_loss_per_position:.2f}")
        
        # Liquidate positions that hit stop loss and log portfolio state
        if symbols_to_liquidate:
            self.log(f"Portfolio before liquidation - Allocated: {allocation_pct:.1f}% (${allocated_value:,.0f}) | Cash: ${cash_value:,.0f}")
            for symbol in symbols_to_liquidate:
                self.liquidate(symbol)
            self.log(f"Liquidated {len(symbols_to_liquidate)} positions due to stop loss")
    
    def _allocate_excess_cash(self):
        # Allocate excess cash to gold between rebalances
        portfolio_value = self.portfolio.total_portfolio_value
        if portfolio_value <= 0:
            return
        
        # Calculate current allocation excluding cash and gold
        equity_value = sum(holding.absolute_holdings_value for symbol, holding in self.portfolio.items() 
                          if holding.quantity > 0 and symbol.value != self._hedge_etf)
        equity_allocation = equity_value / portfolio_value
        
        # If equity allocation is below 100%, allocate excess to gold
        excess_allocation = 1.0 - equity_allocation
        
        if excess_allocation > 0.05:  # Only rebalance if more than 5% excess
            gold_symbol = self.symbol(self._hedge_etf)
            current_gold_value = self.portfolio[gold_symbol].absolute_holdings_value
            current_gold_allocation = current_gold_value / portfolio_value if portfolio_value > 0 else 0
            
            if excess_allocation > current_gold_allocation + 0.02:  # More than 2% opportunity
                self.set_holdings([PortfolioTarget(gold_symbol, excess_allocation)], liquidate_existing_holdings=False)
                self.log(f"Allocated {excess_allocation:.2%} to {self._hedge_etf} (excess cash reallocation)")