# region imports
from AlgorithmImports import *
# endregion

class TechMomentumStrategy(QCAlgorithm):
    """Tech stock momentum strategy - top 5 by market cap with best momentum"""

    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_cash(100000)
        
        self.set_warm_up(126)
        self.settings.seed_initial_prices = True
        
        # Strategy parameters
        self._num_positions = 5
        self._momentum_period = 126  # 6 months
        self._volatility_period = 30  # 30 days for volatility calculation
        self._correlation_period = 60  # 60 days for correlation calculation
        self._rebalance_flag = False
        self._momentum_indicators = {}
        self._volatility_indicators = {}
        self._return_history = {}  # Track daily returns for correlation
        
        # Trailing stop parameters
        self._trailing_stop_percent = 0.15  # 15% trailing stop
        self._hard_stop_percent = 0.25  # 25% hard stop from entry
        self._position_highs = {}  # Track highest price for each position
        self._position_entry_prices = {}  # Track entry prices
        
        # Add tech stock universe
        self._universe = self.add_universe(self._select_tech_stocks)
        
        # Schedule monthly rebalancing
        self.schedule.on(self.date_rules.month_start(),
                        self.time_rules.after_market_open("SPY", 30),
                        self._set_rebalance_flag)
        
        # Check trailing stops daily
        self.schedule.on(self.date_rules.every_day(),
                        self.time_rules.after_market_open("SPY", 15),
                        self._check_trailing_stops)
        
        # Track returns daily for correlation
        self.schedule.on(self.date_rules.every_day(),
                        self.time_rules.after_market_open("SPY", 1),
                        self._track_returns)
    
    def _select_tech_stocks(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        """Select tech stocks with positive momentum"""
        
        # 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  # > $1B market cap
            and f.dollar_volume > 5e6  # > $5M daily volume
        ]
        
        if len(tech_stocks) < 20:
            return []
        
        # Return top 50 tech stocks by market cap for momentum filtering
        sorted_by_cap = sorted(tech_stocks, key=lambda x: x.market_cap, reverse=True)
        return [x.symbol for x in sorted_by_cap[:50]]
    
    def on_securities_changed(self, changes: SecurityChanges):
        """Add momentum and volatility indicators for new securities"""
        for security in changes.added_securities:
            if security.symbol not in self._momentum_indicators:
                self._momentum_indicators[security.symbol] = self.momp(security.symbol, self._momentum_period, Resolution.DAILY)
            if security.symbol not in self._volatility_indicators:
                self._volatility_indicators[security.symbol] = self.std(security.symbol, self._volatility_period, Resolution.DAILY)
            if security.symbol not in self._return_history:
                self._return_history[security.symbol] = RollingWindow[float](self._correlation_period)
        
        for security in changes.removed_securities:
            if security.symbol in self._momentum_indicators:
                self.remove_security(security.symbol)
    
    def _track_returns(self):
        """Track daily returns for correlation calculation"""
        for symbol in self._return_history.keys():
            if symbol in self.securities and self.securities[symbol].has_data:
                history = self.history(symbol, 2, Resolution.DAILY)
                if len(history) == 2:
                    returns = (history.iloc[-1]['close'] / history.iloc[-2]['close']) - 1
                    self._return_history[symbol].add(returns)
    
    def _calculate_correlation(self, symbol1: Symbol, symbol2: Symbol) -> float:
        """Calculate correlation between two symbols"""
        if symbol1 not in self._return_history or symbol2 not in self._return_history:
            return 0.5  # Default moderate correlation
        
        window1 = self._return_history[symbol1]
        window2 = self._return_history[symbol2]
        
        if not window1.is_ready or not window2.is_ready:
            return 0.5
        
        # Get returns as lists
        returns1 = [window1[i] for i in range(window1.count)]
        returns2 = [window2[i] for i in range(window2.count)]
        
        # Calculate correlation
        import numpy as np
        corr = np.corrcoef(returns1, returns2)[0, 1]
        
        return corr if not np.isnan(corr) else 0.5
    
    def _set_rebalance_flag(self):
        """Set flag to trigger rebalancing"""
        self._rebalance_flag = True
    
    def _check_trailing_stops(self):
        """Check trailing stops daily and exit positions if triggered"""
        for symbol, holding in self.portfolio.items():
            if not holding.invested:
                continue
            
            current_price = self.securities[symbol].price
            
            # Update position high
            if symbol not in self._position_highs:
                self._position_highs[symbol] = current_price
            else:
                self._position_highs[symbol] = max(self._position_highs[symbol], current_price)
            
            # Check trailing stop (15% from peak)
            peak_price = self._position_highs[symbol]
            trailing_stop_price = peak_price * (1 - self._trailing_stop_percent)
            
            if current_price <= trailing_stop_price:
                self.liquidate(symbol)
                self.log(f"Trailing stop hit for {symbol}: {current_price:.2f} <= {trailing_stop_price:.2f} (peak: {peak_price:.2f})")
                if symbol in self._position_highs:
                    del self._position_highs[symbol]
                if symbol in self._position_entry_prices:
                    del self._position_entry_prices[symbol]
                continue
            
            # Check hard stop (25% from entry price)
            if symbol in self._position_entry_prices:
                entry_price = self._position_entry_prices[symbol]
                hard_stop_price = entry_price * (1 - self._hard_stop_percent)
                
                if current_price <= hard_stop_price:
                    self.liquidate(symbol)
                    self.log(f"Hard stop hit for {symbol}: {current_price:.2f} <= {hard_stop_price:.2f} (entry: {entry_price:.2f})")
                    if symbol in self._position_highs:
                        del self._position_highs[symbol]
                    if symbol in self._position_entry_prices:
                        del self._position_entry_prices[symbol]
    
    def on_data(self, data: Slice):
        """Rebalance portfolio monthly with volatility-adjusted weights"""
        if not self._rebalance_flag:
            return
        
        self._rebalance_flag = False
        
        # Get active tech stocks from universe
        active_symbols = list(self._universe.selected)
        
        if len(active_symbols) < 5:
            return
        
        # Calculate momentum for each stock
        momentum_scores = []
        for symbol in active_symbols:
            if symbol in self._momentum_indicators:
                indicator = self._momentum_indicators[symbol]
                if indicator.is_ready and indicator.current.value > 0:
                    security = self.securities[symbol]
                    momentum_scores.append({
                        'symbol': symbol,
                        'momentum': indicator.current.value,
                        'market_cap': security.fundamentals.market_cap
                    })
        
        if len(momentum_scores) < 5:
            return
        
        # Sort by momentum (highest first)
        sorted_by_momentum = sorted(momentum_scores, key=lambda x: x['momentum'], reverse=True)
        
        # Take top 20 momentum stocks for correlation filtering
        top_momentum = sorted_by_momentum[:20]
        sorted_by_cap = sorted(top_momentum, key=lambda x: x['market_cap'], reverse=True)
        candidates = [x['symbol'] for x in sorted_by_cap[:15]]
        
        # Select stocks with lowest average correlation (greedy algorithm)
        final_stocks = []
        
        # Start with highest market cap stock
        if len(candidates) > 0:
            final_stocks.append(candidates[0])
        
        # Add stocks with lowest average correlation to already selected stocks
        while len(final_stocks) < self._num_positions and len(candidates) > len(final_stocks):
            best_candidate = None
            lowest_avg_correlation = float('inf')
            
            for candidate in candidates:
                if candidate in final_stocks:
                    continue
                
                # Calculate average correlation with already selected stocks
                correlations = []
                for selected in final_stocks:
                    corr = self._calculate_correlation(candidate, selected)
                    correlations.append(abs(corr))
                
                avg_corr = sum(correlations) / len(correlations) if correlations else 0
                
                if avg_corr < lowest_avg_correlation:
                    lowest_avg_correlation = avg_corr
                    best_candidate = candidate
            
            if best_candidate:
                final_stocks.append(best_candidate)
                self.log(f"Added {best_candidate.value} with avg correlation {lowest_avg_correlation:.2f}")
            else:
                break
        
        # Calculate inverse volatility weights (more allocation to lower volatility stocks)
        weights = {}
        total_inverse_vol = 0
        
        for symbol in final_stocks:
            if symbol in self._volatility_indicators:
                vol_indicator = self._volatility_indicators[symbol]
                if vol_indicator.is_ready and vol_indicator.current.value > 0:
                    # Inverse volatility - lower vol gets higher weight
                    inverse_vol = 1.0 / vol_indicator.current.value
                    weights[symbol] = inverse_vol
                    total_inverse_vol += inverse_vol
                else:
                    # Default equal weight if volatility not ready
                    weights[symbol] = 1.0
                    total_inverse_vol += 1.0
            else:
                weights[symbol] = 1.0
                total_inverse_vol += 1.0
        
        # Normalize weights to sum to 1.0
        if total_inverse_vol > 0:
            for symbol in weights:
                weights[symbol] = weights[symbol] / total_inverse_vol
        
        # Create portfolio targets with volatility-adjusted weights
        targets = [PortfolioTarget(symbol, weights[symbol]) for symbol in final_stocks]
        
        # Execute rebalancing
        self.set_holdings(targets, liquidate_existing_holdings=True)
        
        # Update entry prices and reset highs for new positions
        for symbol in final_stocks:
            current_price = self.securities[symbol].price
            if symbol not in self.portfolio or not self.portfolio[symbol].invested:
                self._position_entry_prices[symbol] = current_price
                self._position_highs[symbol] = current_price
        
        # Log weights for monitoring
        weights_str = ", ".join([f"{s.value}: {w:.1%}" for s, w in weights.items()])
        self.log(f"Rebalanced with volatility weights: {weights_str}")