import numpy as np
from AlgorithmImports import *

class AssetWeightCalculator:
    def __init__(self, algorithm: QCAlgorithm):
        
        self.algorithm = algorithm

        self.risk_free = self.algorithm.add_equity("BIL", Resolution.HOUR)
        
    def coarse_selection(self, coarse):
        """
        Selects stonks, first filter
        """

        # Sorts by dollar volume before taking top 200 
        sorted_by_volume = sorted([x for x in coarse if x.price > 10 and x.has_fundamental_data],
                                key=lambda x: x.dollar_volume, 
                                reverse=True)
        return [x.symbol for x in sorted_by_volume][:200]

    def fine_selection(self, fine):
        """
        Selects stonks, second filter
        """
        filtered = [x.symbol for x in fine if x.market_cap is not None and x.market_cap > 10e9]
        self.algorithm.debug(f"Fine Selection: {len(filtered)} symbols passed filters")

        # Doing it this way makes it so that stocks are ranked on each universe update and then the macds can be redone with the scheduler in main
        ranked_symbols = self.rank_stocks(filtered)
        return ranked_symbols

    def calculate_sharpe_ratio(self, symbol, period=4914): # This is 3 yrs worth of trading days
        """
        Calculates the sharpe
        """
        try:
            # If a KeyValuePair was recieved only take the symbol
            if hasattr(symbol, "Key"):
                symbol = symbol.Key

            history = self.algorithm.history([symbol], period, Resolution.HOUR) 

            if history.empty:
                self.algorithm.debug(f"No history for {symbol.value}")
                return None
            
            # Get risk-free rate
            rf_history = self.algorithm.history(self.risk_free.symbol, 1, Resolution.HOUR)
            risk_free_rate = rf_history['close'].iloc[-1]/100 if not rf_history.empty else 0.02  # Default to 2% if no data
            
            # Sharpe ratio logic
            returns = history['close'].pct_change().dropna()
            excess_returns = returns - (risk_free_rate/1638)
            mean_excess_return = excess_returns.mean() * 1638
            std_dev = excess_returns.std() * np.sqrt(1638)
            return mean_excess_return / std_dev if std_dev != 0 else None
            
        except Exception as e:
            self.algorithm.debug(f"Error calculating Sharpe for {symbol.value}: {str(e)}")
            return None

    def rank_stocks(self, symbols):
        """
        Ranks da top 50 stocks based on sharpe
        """
        if not symbols:
            self.algorithm.debug("No symbols to rank")
            return []
            
        self.algorithm.debug(f"Ranking {len(symbols)} symbols")

        # Converting from key pair if neccessary
        symbols = [s.Key if hasattr(s, 'Key') else s for s in symbols]
        scores = {symbol: self.calculate_sharpe_ratio(symbol) for symbol in symbols}
        valid_scores = {k: v for k, v in scores.items() if v is not None}
        
        self.algorithm.debug(f"Valid Sharpe ratios: {len(valid_scores)} out of {len(symbols)}")
        
        if not valid_scores:
            return []
            
        sorted_scores = sorted(valid_scores, key=valid_scores.get, reverse=True)[:20]

        self.algorithm.log(f"All symbols before ranking: {[s.value for s in symbols]}")
        self.algorithm.log(f"Symbols after filtering: {[s.value for s in valid_scores.keys()]}")

        return sorted_scores

    def normalize_scores(self, scores):
        """
        The list of scores from the ranking method are
        normalized using a z score so that an additive
        operation may be used in WeightCombiner()
        """
        values = np.array(list(scores.values()))
        mean = np.mean(values)
        std_dev = np.std(values)

        if std_dev == 0:
            # If no variation in scores, assign equal normalized scores
            return {symbol: 0 for symbol in scores.keys()}

        normalized_scores = {symbol: (score - mean) / std_dev for symbol, score in scores.items()}
        print(normalized_scores) #To see output for debugging
        return normalized_scores

from AlgorithmImports import *

class MACDSignalGenerator:

    def __init__(self, algorithm: QCAlgorithm, symbols: list, cash_buffer: float = 0.05):
        self.algorithm = algorithm
        self.symbols = symbols
        self.cash_buffer = cash_buffer
        self.macd_indicators = {}  # {symbol: {variant: MACD}}
            
        # Define MACD parameters for different variants
        self.macd_variants = {
            "slow": {"fast": 12, "slow": 26, "signal": 9},
            "slow-med": {"fast": 9, "slow": 19, "signal": 5},
            "med-fast": {"fast": 7, "slow": 15, "signal": 3},
            "fast": {"fast": 5, "slow": 12, "signal": 2},
        }

    def remove_symbols(self, symbols: list):
        """
        Removes MACD indicators for the specified symbols.
        """
        for symbol in symbols:

            # Liquidate position before removing indicator
            self.algorithm.liquidate(symbol)

            # Unregister and delete indicators tied to each symbol
            if symbol in self.macd_indicators:
                for macd in self.macd_indicators[symbol].values():  # Better: gets MACD objects directly
                    self.algorithm.unregister_indicator(macd)
                del self.macd_indicators[symbol]
                

    def add_symbols(self, new_symbols):
            """
            Add in the new symbols that are given by AssetWeightCalculator.
            """
            # Log initial attempt
            self.algorithm.debug(f"Attempting to add symbols: {[s.value for s in new_symbols]}")

            # Get historical data for new symbols
            history = self.algorithm.history([s for s in new_symbols], 
                                        35,  # Longest MACD period needed
                                        Resolution.HOUR)

            # Log history data availability
            self.algorithm.debug(f"History data available for: {history.index.get_level_values(0).unique()}")
            
            self.symbols.extend(new_symbols)
            for symbol in new_symbols:

                security = self.algorithm.securities[symbol]

                # Detailed security check logging
               # self.algorithm.debug(f"Security {symbol.value} check:"
                               # f" has_data={security.has_data},"
                               # f" is_tradable={security.is_tradable},"
                               # f" price={security.price}")

                # Checking if price is 0
                if not (security.has_data and security.is_tradable and security.price > 0):
                    self.algorithm.debug(f"Waiting for valid price data: {symbol.value}")
                    continue

                # Adding the symbol
                if symbol not in self.macd_indicators:
                    self.macd_indicators[symbol] = {}

                    # Get symbol's historical data
                    if symbol not in history.index.get_level_values(0):
                        self.algorithm.debug(f"No history data for: {symbol.value}")
                        continue
                        
                    symbol_history = history.loc[symbol]
                    self.algorithm.debug(f"History rows for {symbol.value}: {len(symbol_history)}")

                    for variant, params in self.macd_variants.items():
                        macd = self.algorithm.macd(
                            symbol=symbol,
                            fast_period=params["fast"], 
                            slow_period=params["slow"], 
                            signal_period=params["signal"], 
                            type=MovingAverageType.EXPONENTIAL,
                            resolution=Resolution.HOUR,
                            selector=Field.CLOSE
                        )
                        self.macd_indicators[symbol][variant] = macd

                        # Warm up MACD with historical data
                        for time, row in symbol_history.iterrows():
                            macd.update(time, row['close'])
                            
                        self.macd_indicators[symbol][variant] = macd

    def calculate_position_sizes(self):
        position_sizes = {}
        max_position_limit = 0.1

        # Check if we have any symbols to process
        if not self.symbols or not self.macd_indicators:
            self.algorithm.debug("No symbols available for position calculation")
            return position_sizes
        
        # Calculating the maximum one variant can be in size
        max_position = (1 - self.cash_buffer) / (len(self.symbols) * len(self.macd_variants))


        for symbol in self.macd_indicators:
            position_sizes[symbol] = {}

            for variant, macd in self.macd_indicators[symbol].items():
                if macd.is_ready:

                    security = self.algorithm.securities[symbol]

                    # Detailed security check logging
                    # self.algorithm.debug(f"Position Check for {symbol.value}:"
                                   # f" has_data={security.has_data},"
                                   # f" is_tradable={security.is_tradable},"
                                   # f" price={security.price},"
                                   # f" last_data={security.get_last_data() is not None},")
                    
                    # More comprehensive check
                    # if not (security.has_data and 
                           # security.is_tradable and 
                           # security.price > 0 and
                           # security.get_last_data() is not None):
                       # self.algorithm.debug(f"Security not ready: {symbol.value}")
                       # continue

                    # Distance between fast and slow
                    distance = macd.fast.current.value - macd.slow.current.value

                    # Normalize the distance as a percentage difference and then as a fraction of max position
                    position_size = max_position * (distance / macd.slow.current.value) * 70 # Scalar value of max_position, the scalar integer can be though of as a form of leverage setting
                    
                    # Only allow positive positions, cap at maximum
                    position_size = max(0, min(position_size, max_position_limit))
                    position_sizes[symbol][variant] = position_size
                    #self.algorithm.debug(f"Calculated position for {symbol.value} {variant}: {position_size}")
                
                else:
                  position_sizes[symbol][variant] = 0 

        # Running daily cause the logging is too heavy hourly 
        if self.algorithm.time.hour == 10 and self.algorithm.time.minute == 0:
            rounded_positions = [(s.value, {k: round(v, 5) for k, v in sizes.items()}) for s, sizes in position_sizes.items()]
            #self.algorithm.debug(f"Daily position sizes proposed: {rounded_positions}")

        return position_sizes

# region imports
from AlgorithmImports import *
# endregion
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data import *
from QuantConnect.Indicators import *
from datetime import timedelta
import numpy as np
import pandas as pd
import torch
import os
import torch.nn as nn


from strategy import KQTStrategy

def calculate_ema(prices, span):
    # Using adjust=False mimics how many trading platforms calculate EMA
    return pd.Series(prices).ewm(span=span, adjust=False).mean().values
def calculate_rsi(prices, period=14):
    """Calculate Relative Strength Index"""
    deltas = np.diff(prices)
    seed = deltas[:period+1]
    
    up = seed[seed >= 0].sum()/period
    down = -seed[seed < 0].sum()/period
    
    if down == 0: return 100  # No downward movement means RSI == 100
    rs = up/down
    rsi = np.zeros_like(prices)
    rsi[:period] = 100. - 100./(1. + rs)
    
    for i in range(period, len(prices)):
        delta = deltas[i-1]
        
        if delta > 0:
            upval = delta
            downval = 0.
        else:
            upval = 0.
            downval = -delta
            
        up = (up * (period-1) + upval) / period
        down = (down * (period-1) + downval) / period
        
        rs = up/down if down != 0 else float('inf')
        rsi[i] = 100. - 100./(1. + rs)
        
    return rsi

def calculate_macd(prices, fast=12, slow=26, signal=9):
    """Calculate MACD line and histogram"""
    # Convert to numpy array if not already
    prices = np.array(prices)
    
    # Calculate EMAs
    ema_fast = pd.Series(prices).ewm(span=fast, adjust=False).mean().values
    ema_slow = pd.Series(prices).ewm(span=slow, adjust=False).mean().values
    
    # Calculate MACD line and signal line
    macd_line = ema_fast - ema_slow
    signal_line = pd.Series(macd_line).ewm(span=signal, adjust=False).mean().values
    
    # Calculate histogram
    histogram = macd_line - signal_line
    
    return macd_line, signal_line, histogram

def calculate_atr(high, low, close, period=14):
    """Calculate Average True Range"""
    if len(high) != len(low) or len(high) != len(close):
        raise ValueError("Input arrays must have the same length")
    
    tr = np.zeros(len(high))
    tr[0] = high[0] - low[0]  # Initial TR = high - low of first bar
    
    for i in range(1, len(tr)):
        tr[i] = max(
            high[i] - low[i],
            abs(high[i] - close[i-1]),
            abs(low[i] - close[i-1])
        )
    
    # Calculate ATR
    atr = np.zeros_like(tr)
    atr[0] = tr[0]
    for i in range(1, len(atr)):
        atr[i] = (atr[i-1] * (period-1) + tr[i]) / period
        
    return atr

class KQTAlgorithm(QCAlgorithm):
    def Initialize(self):
        """Initialize the algorithm"""
        # Set start date, end date, and cash
        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2025, 4, 6)
        self.SetCash(1000000)
        self.previous_portfolio_value = 0

        # Set benchmark to SPY
        self.SetBenchmark("SPY")
        
        # Initialize the KQT strategy
        self.strategy = KQTStrategy()
        self.lookback = 60  # Need enough data for technical indicators
        self.tickers = []
        self.symbols = {}
        self.sector_mappings = {}
        self.strategy.sector_mappings = self.sector_mappings  # Share the dictionary
        self._universe = self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)        

        # Add SPY for market data
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # Add bond ETF for market downturns
        self.bond_etf = self.AddEquity("BIL", Resolution.Daily).Symbol
        
        # Initialize moving averages for market regime detection
        self.spy_sma50 = SimpleMovingAverage(50)
        self.spy_sma200 = SimpleMovingAverage(200)
        self.RegisterIndicator(self.spy, self.spy_sma50, Resolution.Daily)
        self.RegisterIndicator(self.spy, self.spy_sma200, Resolution.Daily)
        self.in_bond_position = False
        
        # Storage for historical data and predictions
        self.stock_data = {}
        self.current_predictions = {}
        self.previous_positions = {}
        
        # Track stopped out positions
        self.stopped_out = set()
        
        # Schedule the trading function to run before market close
        self.Schedule.On(self.DateRules.EveryDay(), 
                self.TimeRules.At(10, 0),  # 10:00 AM Eastern
                self.TradeExecute)
        
        # Initialize model with feature count
        feature_count = 18  # Adjust based on your actual feature count
        
        # Load pre-trained model weights if available
        self.TryLoadModelWeights()
        
    def CoarseSelectionFunction(self, coarse):
        sorted_by_dollar_volume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sorted_by_dollar_volume[:500]]
    def FineSelectionFunction(self, fine):
        sorted_by_market_cap = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        selected = sorted_by_market_cap[:100]
        
        # Debug the first item to understand available properties
        if len(selected) > 0:
            f = selected[0]
            self.Debug(f"Available fundamental properties: {[attr for attr in dir(f) if not attr.startswith('_')]}")
            if hasattr(f, 'AssetClassification'):
                self.Debug(f"AssetClassification properties: {[attr for attr in dir(f.AssetClassification) if not attr.startswith('_')]}")
        
        for f in selected:
            ticker = f.Symbol.Value
            # Try multiple ways to get sector information
            sector = "Unknown"
            
            try:
                if hasattr(f, 'AssetClassification') and f.AssetClassification is not None:
                    # Try commonly used sector properties
                    if hasattr(f.AssetClassification, 'MorningstarSectorCode'):
                        sector = str(f.AssetClassification.MorningstarSectorCode)
                    elif hasattr(f.AssetClassification, 'MorningstarIndustryCode'):
                        sector = str(f.AssetClassification.MorningstarIndustryCode)
                    elif hasattr(f.AssetClassification, 'GicsCode'):
                        sector = str(f.AssetClassification.GicsCode)
                    # Additional fallbacks
                    elif hasattr(f.AssetClassification, 'Sector'):
                        sector = f.AssetClassification.Sector
                    elif hasattr(f.AssetClassification, 'Industry'):
                        sector = f.AssetClassification.Industry
            except Exception as e:
                self.Debug(f"Error getting sector for {ticker}: {str(e)}")
                
            self.sector_mappings[ticker] = sector
            
        return [f.Symbol for f in selected]
    def OnSecuritiesChanged(self, changes):
        self.Debug(f"Universe changed: Added {len(changes.AddedSecurities)}, Removed {len(changes.RemovedSecurities)}")
        for added in changes.AddedSecurities:
            self.Debug(f"Added: {added.Symbol.Value}")
        for removed in changes.RemovedSecurities:
            self.Debug(f"Removed: {removed.Symbol.Value}")
        for added in changes.AddedSecurities:
            ticker = added.Symbol.Value
            if ticker not in self.tickers:
                self.tickers.append(ticker)
            self.symbols[ticker] = added.Symbol
        for removed in changes.RemovedSecurities:
            ticker = removed.Symbol.Value
            if ticker in self.tickers:
                self.tickers.remove(ticker)
            if ticker in self.symbols:
                del self.symbols[ticker]
            if ticker in self.sector_mappings:
                del self.sector_mappings[ticker]
            if ticker in self.stock_data:
                del self.stock_data[ticker]
    def TryLoadModelWeights(self):
        """Try to load model weights from ObjectStore"""
        try:
            if self.ObjectStore.ContainsKey("kqt_model_weights"):
                self.Debug("Found model weights in ObjectStore, loading...")
                # Get base64 encoded string
                encoded_bytes = self.ObjectStore.Read("kqt_model_weights")
                
                # Decode back to binary
                import base64
                model_bytes = base64.b64decode(encoded_bytes)
                
                # Save temporarily to file
                import tempfile
                with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as temp:
                    temp_path = temp.name
                    temp.write(model_bytes)
                
                # Extract input size from saved weights and reinitialize model
                try:
                    state_dict = torch.load(temp_path)
                    input_shape = state_dict['embedding.0.weight'].shape
                    actual_input_size = input_shape[1]
                    self.Debug(f"Detected input size from weights: {actual_input_size}")
                    
                    # Reinitialize model with correct input size
                    self.Debug("Successfully loaded model weights")
                except Exception as inner_e:
                    self.Debug(f"Error examining weights: {str(inner_e)}")
                
                # Clean up
                import os
                os.unlink(temp_path)
            else:
                self.Debug("No model weights found in ObjectStore")
        except Exception as e:
            self.Debug(f"Error loading model weights: {str(e)}")
    
    def OnData(self, data):
        """OnData event is the primary entry point for your algorithm."""
        # We're using scheduled events instead of processing each data point
        pass
    
    def TradeExecute(self):
        """Execute trading logic daily before market close"""
        # *** FIXED Market Open Check ***
        if not self.Securities.ContainsKey(self.spy) or not self.Securities[self.spy].Exchange.ExchangeOpen:
            # self.Debug(f"Skipping TradeExecute: Market closed or SPY not ready. Time: {self.Time}")
            return # Skip if market is closed or SPY data isn't loaded yet
        
        self.Debug(f"TradeExecute running on {self.Time}")
        
        # Check if SMAs are ready
        if not self.spy_sma50.IsReady or not self.spy_sma200.IsReady:
            self.Debug("Moving averages not ready yet, running normal strategy")
            self.ExecuteNormalStrategy()
            return
            
        # Check market regime
        spy_price = self.Securities[self.spy].Price
        is_bear_market = spy_price < self.spy_sma50.Current.Value and spy_price < self.spy_sma200.Current.Value
        
        if is_bear_market:
            # Bearish regime - move to bonds if not already there
            if not self.in_bond_position:
                self.Debug(f"BEAR MARKET DETECTED - Moving to bonds. SPY: {spy_price}, SMA50: {self.spy_sma50.Current.Value}, SMA200: {self.spy_sma200.Current.Value}")
                self.Liquidate()  # Sell all positions
                self.MarketOrder(self.bond_etf, int(self.Portfolio.Cash * 0.95 / self.Securities[self.bond_etf].Price))  # Use 95% of cash for bonds
                self.in_bond_position = True
            else:
                self.Debug("Maintaining bond position during bear market")
        else:
            # Bullish regime - exit bonds if in them
            if self.in_bond_position:
                self.Debug(f"BULL MARKET DETECTED - Exiting bonds and returning to normal strategy. SPY: {spy_price}, SMA50: {self.spy_sma50.Current.Value}, SMA200: {self.spy_sma200.Current.Value}")
                self.Liquidate(self.bond_etf)  # Sell bond position
                self.in_bond_position = False
            
            # Execute normal strategy
            self.ExecuteNormalStrategy()
            
        # Update portfolio return for regime detection
        daily_return = self.CalculatePortfolioReturn()
        self.strategy.update_portfolio_returns(daily_return)
        
        # Store today's value for tomorrow's calculation
        self.previous_portfolio_value = self.Portfolio.TotalPortfolioValue

    def ExecuteNormalStrategy(self):
        """Execute the original trading strategy"""
        self.Debug(f"Current universe size: {len(self.tickers)}")
        self.Debug(f"Using sectors: {list(self.sector_mappings.keys())[:5]}...")
        
        # 1. Update historical data for all stocks
        self.UpdateHistoricalData()
        
        # 2. Generate predictions for each stock
        self.current_predictions = self.GeneratePredictions()

        # 3. Check for stop losses
        self.ProcessStopLosses()
        
        # 4. Generate new position sizes
        market_returns = self.GetMarketReturns()
        target_positions = self.strategy.generate_positions(self.current_predictions, market_returns)
        self.Debug(f"Target positions before execution: {target_positions}")
        self.Debug(f"Market returns: {market_returns}")

        # 5. Execute trades to reach target positions
        self.ExecuteTrades(target_positions)
        
        self.Debug(f"Generated {len(target_positions)} positions: {target_positions}")
    
    def CalculatePortfolioReturn(self):
        """Calculate today's portfolio return"""
        # Get the portfolio value change
        current_value = self.Portfolio.TotalPortfolioValue
        
        # Use our stored previous value instead of the non-existent property
        if self.previous_portfolio_value > 0:
            return (current_value / self.previous_portfolio_value - 1) * 100  # as percentage
        
        # On first day, just store the value and return 0
        self.previous_portfolio_value = current_value
        return 0
    
    def UpdateHistoricalData(self):
        """Fetch and update historical data for all symbols"""
        for ticker in self.tickers:
            # Use the stored Symbol object instead of accessing Securities by ticker
            if ticker not in self.symbols:
                self.Debug(f"Symbol not found for ticker {ticker}, skipping")
                continue
                
            symbol = self.symbols[ticker]
            
            # Request sufficient history for features
            history = self.History(symbol, self.lookback, Resolution.Daily)
            
            if history.empty or len(history) < self.lookback:
                self.Debug(f"Not enough historical data for {ticker}, skipping")
                continue
                    
            # Store historical data
            if isinstance(history.index, pd.MultiIndex):
                history_reset = history.reset_index()
                symbol_data = history_reset[history_reset['symbol'] == symbol]
                self.stock_data[ticker] = symbol_data
            else:
                self.stock_data[ticker] = history
    
    def GetMarketReturns(self):
        """Get recent market returns for regime detection"""
        spy_history = self.History(self.spy, 10, Resolution.Daily)
        
        if spy_history.empty:
            return []
            
        # Handle both MultiIndex and regular index formats
        if isinstance(spy_history.index, pd.MultiIndex):
            spy_history_reset = spy_history.reset_index()
            spy_history_filtered = spy_history_reset[spy_history_reset['symbol'] == self.spy]
            spy_prices = spy_history_filtered['close'].values
        else:
            spy_prices = spy_history['close'].values
            
        # Calculate returns
        spy_returns = []
        for i in range(1, len(spy_prices)):
            daily_return = (spy_prices[i] / spy_prices[i-1] - 1) * 100
            spy_returns.append(daily_return)
            
        return spy_returns
    
    def GeneratePredictions(self):
        """Generate predictions for all stocks"""
        predictions = {}
        self.Debug(f"Generating predictions with {len(self.stock_data)} stocks in data")

        for ticker, history in self.stock_data.items():
            try:
                if len(history) < self.lookback:
                    self.Debug(f"Skipping {ticker}: Not enough history ({len(history)} < {self.lookback})")
                    continue

                # Prepare data for this stock
                X, stock_df = self.strategy.prepare_stock_data(history, ticker)
                if X is None or len(X) == 0:

                    # FALLBACK: Simple prediction if ML fails
                    if isinstance(history.index, pd.MultiIndex):
                        history_reset = history.reset_index()
                        closes = history_reset[history_reset['symbol'] == self.symbols[ticker]]['close'].values
                    else:
                        closes = history['close'].values
                    
                    if len(closes) > 20:
                        # Simple momentum strategy for fallback
                        short_ma = np.mean(closes[-5:])
                        long_ma = np.mean(closes[-20:])
                        momentum = closes[-1] / closes[-10] - 1 if len(closes) > 10 else 0
                        
                        pred_score = momentum + 0.5 * (short_ma/long_ma - 1)
                        predictions[ticker] = {
                            "pred_return": pred_score * 2,
                            "composite_score": pred_score * 5
                        }
                        self.Debug(f"Used fallback prediction for {ticker}: {pred_score}")
                    continue
                # If we reach here, X must be valid data for the ML model
                # Generate prediction with ML model
                pred_return = self.strategy.predict_returns(X, ticker) # Now X is guaranteed not to be None

                # Check if prediction itself returned None (handle potential issues in predict_returns)
                if pred_return is None:
                     self.Debug(f"ML prediction for {ticker} returned None. Skipping.")
                     continue # Skip this ticker if ML prediction failed internally

                # Store ML prediction
                predictions[ticker] = {
                    "pred_return": pred_return,
                    "composite_score": pred_return / self.strategy.adaptive_threshold if self.strategy.adaptive_threshold != 0 else pred_return # Avoid division by zero
                }
                self.Debug(f"ML prediction for {ticker}: {pred_return}")

            except Exception as e:
                # Catch errors during data prep or ML prediction call itself
                self.Debug(f"Error processing {ticker} in GeneratePredictions loop: {str(e)}")
                import traceback
                self.Debug(traceback.format_exc()) # Log full traceback for debugging
                continue # Skip to next ticker on any error in the main try block

        self.Debug(f"Generated {len(predictions)} predictions")
        return predictions
    
    def ProcessStopLosses(self):
        """Check and process stop loss orders"""
        stop_loss_level = self.strategy.get_stop_loss_level()
        
        for ticker in self.tickers:
            if ticker not in self.symbols or not self.Portfolio[self.symbols[ticker]].Invested:
                continue
                
            symbol = self.symbols[ticker]
            position = self.Portfolio[symbol]
            
            # Get today's return
            history = self.History(symbol, 2, Resolution.Daily)
            if history.empty or len(history) < 2:
                continue
                
            # Handle both MultiIndex and regular index formats
            if isinstance(history.index, pd.MultiIndex):
                history_reset = history.reset_index()
                symbol_data = history_reset[history_reset['symbol'] == symbol]
                if len(symbol_data) < 2:
                    continue
                close_prices = symbol_data['close'].values
            else:
                close_prices = history['close'].values
                
            daily_return = (close_prices[-1] / close_prices[-2] - 1) * 100
            
            position_type = "long" if position.Quantity > 0 else "short"
            hit_stop = False
            
            if position_type == "long" and daily_return < stop_loss_level:
                hit_stop = True
                self.Debug(f"Stop loss triggered for {ticker} (long): {daily_return:.2f}%")
            elif position_type == "short" and daily_return > -stop_loss_level:
                hit_stop = True
                self.Debug(f"Stop loss triggered for {ticker} (short): {daily_return:.2f}%")
                
            if hit_stop:
                self.stopped_out.add(ticker)
        
    def ExecuteTrades(self, target_positions):
        """Execute trades to reach target positions with improved execution"""
        if not target_positions:
            self.Debug("No target positions received")
            return
        
        self.Debug(f"Executing trades for {len(target_positions)} positions")
        
        # Calculate total allocation first to check for overallocation
        portfolio_value = self.Portfolio.TotalPortfolioValue
        total_allocation = sum(abs(weight) for weight in target_positions.values())
        
        # Scale down if total allocation exceeds 80% (reduced from 95%)
        if total_allocation > 0.8:
            scaling_factor = 0.8 / total_allocation  # Reduced allocation ceiling
            self.Debug(f"Scaling positions by {scaling_factor:.2f} to prevent overallocation")
            for ticker in target_positions:
                target_positions[ticker] *= scaling_factor
        
        # Sort positions by conviction (absolute value of weight) - highest first
        sorted_positions = sorted(target_positions.items(), key=lambda x: abs(x[1]), reverse=True)
        
        # Available buying power tracking - FIXED LINE
        available_buying_power = self.Portfolio.MarginRemaining  # Use MarginRemaining instead of GetBuyingPower()
        self.Debug(f"Starting available margin: ${available_buying_power:.2f}")
        
        # Execute trades to reach target positions
        for ticker, target_weight in sorted_positions:
            symbol = self.symbols[ticker]
            current_security = self.Securities[symbol]
            
            # Calculate target share amount
            price = current_security.Price
            target_value = portfolio_value * target_weight
            target_shares = int(target_value / price) if price > 0 else 0
            
            # Get current holdings
            holding = self.Portfolio[symbol]
            current_shares = holding.Quantity
            
            # Calculate shares to trade
            shares_to_trade = target_shares - current_shares
            
            # Skip tiny orders
            if abs(shares_to_trade) > 0:
                try:
                    # Check if we have enough buying power for this order
                    required_buying_power = abs(shares_to_trade) * price * 1.01  # Add 1% buffer
                    
                    if shares_to_trade < 0 or required_buying_power <= available_buying_power:
                        # Place the order
                        if shares_to_trade > 0:
                            self.MarketOrder(symbol, shares_to_trade)
                            self.Debug(f"BUY {shares_to_trade} shares of {ticker} (${shares_to_trade * price:.2f})")
                            available_buying_power -= required_buying_power
                        else:
                            self.MarketOrder(symbol, shares_to_trade)  # Negative for sell
                            self.Debug(f"SELL {abs(shares_to_trade)} shares of {ticker}")
                            # Selling increases buying power
                            available_buying_power += abs(shares_to_trade) * price * 0.98  # Assume 2% impact/fees
                    else:
                        # Not enough buying power, try reduced size
                        max_affordable = int(available_buying_power / (price * 1.01))
                        if max_affordable > 0:
                            self.MarketOrder(symbol, max_affordable)
                            self.Debug(f"REDUCED BUY {max_affordable} shares of {ticker} due to buying power constraints")
                            available_buying_power -= max_affordable * price * 1.01
                        else:
                            self.Debug(f"Skipped {ticker} order: Insufficient buying power")
                except Exception as e:
                    self.Debug(f"Order error for {ticker}: {str(e)}")
        
        # Store current positions for next day
        self.previous_positions = target_positions.copy()

# region imports
from AlgorithmImports import *
# endregion
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data import *
from QuantConnect.Indicators import *
from datetime import timedelta
import numpy as np
import pandas as pd
import torch
import os
import torch.nn as nn
from sklearn.preprocessing import RobustScaler


class KQTStrategy:
    def __init__(self):
        self.model = None
        self.lookback = 30
        self.scalers = {}
        self.feature_cols = []
        self.stock_to_id = {}
        self.sector_mappings = {}

        self.adaptive_threshold = 0.2
        self.pred_std = 1.0
        self.current_regime = "neutral"
        self.portfolio_returns = []
        self.defensive_mode = False
        self.previous_day_hit_stops = []
        

    
        

    
    def create_sliding_sequences(self, df, feature_cols, lookback, stride=1):
        X = []
        for i in range(0, len(df) - lookback + 1, stride):
            X.append(df.iloc[i:i+lookback][feature_cols].values.astype(np.float32))
        return np.array(X)
    
    def clip_outliers(self, df, cols, lower=0.01, upper=0.99):
        df_copy = df.copy()
        for col in cols:
            if col in df_copy.columns:
                q_low = df_copy[col].quantile(lower)
                q_high = df_copy[col].quantile(upper)
                df_copy.loc[df_copy[col] < q_low, col] = q_low
                df_copy.loc[df_copy[col] > q_high, col] = q_high
        return df_copy
    
    def filter_features_to_match_model(self, df, feature_cols, required_count=5):
        """Ensure we have exactly the required number of features"""
        if len(feature_cols) == required_count:
            return feature_cols
            
        # First, prioritize the lag returns (most important)
        lag_features = [col for col in feature_cols if 'return_lag' in col]
        
        # Next, add in the most predictive technical features in a fixed order
        tech_priority = ['roc_5', 'volatility_10', 'ma_cross', 'dist_ma20', 'momentum_1m',
                        'oversold', 'overbought', 'roc_diff', 'volatility_regime']
                        
        prioritized_features = lag_features.copy()
        for feat in tech_priority:
            if feat in feature_cols and len(prioritized_features) < required_count:
                prioritized_features.append(feat)
        
        # If still not enough, add remaining features
        remaining = [col for col in feature_cols if col not in prioritized_features]
        while len(prioritized_features) < required_count and remaining:
            prioritized_features.append(remaining.pop(0))
        
        # If too many, truncate
        return prioritized_features[:required_count]

    def add_technical_features(self, df):
        if 'Close' not in df.columns:
            return df
            
        df['ma5'] = df['Close'].rolling(5).mean() / df['Close'] - 1  # Relative to price
        df['ma20'] = df['Close'].rolling(20).mean() / df['Close'] - 1
        df['ma_cross'] = df['ma5'] - df['ma20']  # Moving average crossover signal
        
        df['volatility_10'] = df['Close'].pct_change().rolling(10).std()
        df['volatility_ratio'] = df['Close'].pct_change().rolling(5).std() / df['Close'].pct_change().rolling(20).std()
        
        df['roc_5'] = df['Close'].pct_change(5)
        df['roc_10'] = df['Close'].pct_change(10)
        df['roc_diff'] = df['roc_5'] - df['roc_10']
        
        df['dist_ma20'] = (df['Close'] / df['Close'].rolling(20).mean() - 1)
        
        return df.fillna(0)
    
    def add_enhanced_features(self, df):
        """Add enhanced technical features"""
        df['volatility_trend'] = df['volatility_10'].pct_change(5)
        df['volatility_regime'] = (df['volatility_10'] > df['volatility_10'].rolling(20).mean()).astype(int)
        
        if 'volume' in df.columns:
            df['vol_ma_ratio'] = df['volume'] / df['volume'].rolling(20).mean()
            df['vol_price_trend'] = df['vol_ma_ratio'] * df['roc_5']
        
        df['momentum_1m'] = df['Close'].pct_change(20)
        df['momentum_3m'] = df['Close'].pct_change(60)
        df['momentum_breadth'] = (
            (df['roc_5'] > 0).astype(int) + 
            (df['momentum_1m'] > 0).astype(int) + 
            (df['momentum_3m'] > 0).astype(int)
        ) / 3
        
        df['mean_rev_signal'] = -1 * df['dist_ma20'] * df['volatility_10']
        df['oversold'] = (df['dist_ma20'] < -2 * df['volatility_10']).astype(int)
        df['overbought'] = (df['dist_ma20'] > 2 * df['volatility_10']).astype(int)
        
        df['regime_change'] = (np.sign(df['ma_cross']) != np.sign(df['ma_cross'].shift(1))).astype(int)
        
        df['risk_adj_momentum'] = df['roc_5'] / (df['volatility_10'] + 0.001)
        
        return df

    def prepare_stock_data(self, stock_data, ticker, is_training=False):
        """Prepare data for a single stock"""
        if len(stock_data) < self.lookback + 5:  # Need enough data
            return None, None
        
        stock_df = pd.DataFrame({
            'Close': stock_data['close'].values,
            'time': stock_data['time'].values
        })
        
        if 'volume' in stock_data.columns:
            stock_df['volume'] = stock_data['volume'].values
            
        stock_df = stock_df.sort_values('time').reset_index(drop=True)
        
        stock_df['pct_return'] = stock_df['Close'].pct_change().shift(-1) * 100
        
        # In prepare_stock_data, replace the feature cols section with:
        feature_cols = []

        # Add basic lag features
        for i in range(1, 6):
            col_name = f'return_lag{i}'
            stock_df[col_name] = stock_df['pct_return'].shift(i)
            feature_cols.append(col_name)

        # Add technical features
        stock_df = self.add_technical_features(stock_df)
        stock_df = self.add_enhanced_features(stock_df)

        # Add all potential features
        additional_features = ['ma_cross', 'volatility_10', 'roc_5', 'roc_diff', 'dist_ma20']
        enhanced_features = ['volatility_trend', 'volatility_regime', 'momentum_1m', 
                            'momentum_breadth', 'mean_rev_signal', 'oversold', 
                            'overbought', 'regime_change', 'risk_adj_momentum']

        for col in additional_features + enhanced_features:
            if col in stock_df.columns:
                feature_cols.append(col)

        # Filter to the exact number of features expected by the model
        model_feature_count = 5  # Use the exact count from your model
        feature_cols = self.filter_features_to_match_model(stock_df, feature_cols, model_feature_count)

        if not self.feature_cols:
            self.feature_cols = feature_cols.copy()
        
        stock_df = stock_df.dropna().reset_index(drop=True)
        
        # Handle outliers
        stock_df = self.clip_outliers(stock_df, feature_cols)
        
        # Replace the scaling code in prepare_stock_data with this:
        # Scale features
        if ticker not in self.scalers or is_training:
            # Check if we have data
            if len(stock_df) == 0 or len(feature_cols) == 0:
                return None, stock_df  # Return early if no data
                
            # Check if any features are empty/nan
            if stock_df[feature_cols].isna().any().any() or stock_df[feature_cols].empty:
                # Fill NaNs with zeros
                stock_df[feature_cols] = stock_df[feature_cols].fillna(0)
                
            # Ensure we have data
            if len(stock_df[feature_cols]) > 0:
                try:
                    scaler = RobustScaler()
                    stock_df[feature_cols] = scaler.fit_transform(stock_df[feature_cols])
                    self.scalers[ticker] = scaler
                except Exception as e:
                    print(f"Scaling error for {ticker}: {str(e)}")
                    # Use a simple standardization as fallback
                    for col in feature_cols:
                        mean = stock_df[col].mean()
                        std = stock_df[col].std()
                        if std > 0:
                            stock_df[col] = (stock_df[col] - mean) / std
                        else:
                            stock_df[col] = 0
            else:
                return None, stock_df  # Return early if empty after processing
        else:
            # Use existing scaler
            scaler = self.scalers[ticker]
            try:
                stock_df[feature_cols] = scaler.transform(stock_df[feature_cols])
            except Exception as e:
                print(f"Transform error for {ticker}: {str(e)}")
                # Simple standardization fallback
                for col in feature_cols:
                    if col in stock_df.columns and len(stock_df[col]) > 0:
                        mean = stock_df[col].mean()
                        std = stock_df[col].std()
                        if std > 0:
                            stock_df[col] = (stock_df[col] - mean) / std
                        else:
                            stock_df[col] = 0
        
        # Create sequences for prediction
        X = self.create_sliding_sequences(stock_df, feature_cols, self.lookback, stride=1)
        
        if len(X) == 0:
            return None, stock_df
            
        return X, stock_df
    
        # Add to strategy.py in KQTStrategy class
    def calculate_portfolio_risk_score(self, market_returns):
        """Calculate a portfolio risk score (0-100) to scale overall exposure"""
        risk_score = 50  # Neutral starting point
        
        # VIX-like volatility measurement using SPY returns
        if len(market_returns) >= 5:
            recent_vol = np.std(market_returns[-5:]) * np.sqrt(252)  # Annualized
            longer_vol = np.std(market_returns[-10:]) * np.sqrt(252) if len(market_returns) >= 10 else recent_vol
            
            # Volatility spike detection
            vol_ratio = recent_vol / longer_vol if longer_vol > 0 else 1
            if vol_ratio > 1.5:  # Sharp volatility increase
                risk_score -= 30
            elif vol_ratio > 1.2:
                risk_score -= 15
                
        # Consecutive negative days
        if len(market_returns) >= 3:
            neg_days = sum(1 for r in market_returns[-3:] if r < 0)
            if neg_days == 3:  # Three consecutive down days
                risk_score -= 20
            elif neg_days == 2:
                risk_score -= 10
                
        # Trend direction
        if len(market_returns) >= 10:
            avg_recent = np.mean(market_returns[-5:])
            avg_older = np.mean(market_returns[-10:-5])
            trend_change = avg_recent - avg_older
            
            # Declining trend
            if trend_change < -0.3:
                risk_score -= 15
            # Accelerating uptrend
            elif trend_change > 0.3 and avg_recent > 0:
                risk_score += 10
                
        return max(10, min(100, risk_score))  # Constrain between 10-100
            
    def predict_returns(self, X, ticker):
        """Make predictions for a single stock"""
        if self.model is None:
            return 0
            
        if ticker not in self.stock_to_id:
            self.stock_to_id[ticker] = len(self.stock_to_id)
            
        stock_id = self.stock_to_id[ticker]
        
        try:
            X_tensor = torch.tensor(X, dtype=torch.float32)
            stock_ids = torch.tensor([stock_id] * len(X), dtype=torch.long)
            
            with torch.no_grad():
                predictions = self.model(X_tensor, stock_ids)
                
            # Convert to standard Python float for safety
            return float(predictions.detach().numpy().flatten()[-1])
        except Exception as e:
            print(f"Prediction error for {ticker}: {e}")
            return 0  # Return neutral prediction on error
        
    def detect_market_regime(self, daily_returns, lookback=10):
        """Detect current market regime based on portfolio returns"""
        if len(daily_returns) >= 1:
            market_return = np.mean(daily_returns)
            market_vol = np.std(daily_returns)
            
            if len(self.portfolio_returns) >= 3:
                recent_returns = self.portfolio_returns[-min(lookback, len(self.portfolio_returns)):]
                avg_recent_return = np.mean(recent_returns)
                
                if len(self.portfolio_returns) >= 5:
                    very_recent = np.mean(self.portfolio_returns[-3:])
                    less_recent = np.mean(self.portfolio_returns[-min(8, len(self.portfolio_returns)):-3])
                    trend_change = very_recent - less_recent
                    
                    if trend_change > 0.5 and avg_recent_return > 0.2:
                        return "breakout_bullish"
                    elif trend_change < -0.5 and avg_recent_return < -0.2:
                        return "breakdown_bearish"
                
                if avg_recent_return > 0.15:
                    if market_return > 0:
                        return "bullish_strong"
                    else:
                        return "bullish_pullback"
                elif avg_recent_return < -0.3:
                    if market_return < -0.2:
                        return "bearish_high_vol"
                    else:
                        return "bearish_low_vol"
                elif avg_recent_return > 0 and market_return > 0:
                    return "bullish"
                elif avg_recent_return < 0 and market_return < 0:
                    return "bearish"
            
            if market_return > -0.05:
                return "neutral"
            else:
                return "bearish"
        
        return "neutral"
        
    def detect_bearish_signals(self, recent_returns):
        """Detect early warning signs of bearish conditions"""
        bearish_signals = 0
        signal_strength = 0
        
        if len(self.portfolio_returns) >= 5:
            recent_portfolio_returns = self.portfolio_returns[-5:]
            pos_days = sum(1 for r in recent_portfolio_returns if r > 0)
            neg_days = sum(1 for r in recent_portfolio_returns if r < 0)
            
            if neg_days > pos_days:
                bearish_signals += 1
                signal_strength += 0.2 * (neg_days - pos_days)
        
        if len(self.portfolio_returns) >= 10:
            recent_vol = np.std(self.portfolio_returns[-5:])
            older_vol = np.std(self.portfolio_returns[-10:-5])
            if recent_vol > older_vol * 1.3:  # 30% volatility increase
                bearish_signals += 1
                signal_strength += 0.3 * (recent_vol/older_vol - 1)
        
        
        if len(self.portfolio_returns) >= 5:
            if self.portfolio_returns[-1] < 0 and self.portfolio_returns[-2] > 0.3:
                bearish_signals += 1
                signal_strength += 0.3
        
        return bearish_signals, signal_strength
            
    def generate_positions(self, prediction_data, current_returns=None):
        """Generate position sizing based on predictions with improved diversification"""
        if not prediction_data:
            return {}
            
        # Update market regime
        if current_returns is not None:
            self.current_regime = self.detect_market_regime(current_returns)
            bearish_count, bearish_strength = self.detect_bearish_signals(current_returns)
            self.defensive_mode = bearish_count >= 2 or bearish_strength > 0.5
        
        # Calculate portfolio risk score (0-100)
        portfolio_risk_score = self.calculate_portfolio_risk_score(current_returns if current_returns else [])
        # Convert to a scaling factor (0.1 to 1.0)
        risk_scaling = portfolio_risk_score / 100
        
        base_threshold = 0.25 * self.pred_std
        
        if self.current_regime in ["bullish_strong", "breakout_bullish"]:
            self.adaptive_threshold = base_threshold * 0.4
        elif self.current_regime in ["bearish_high_vol", "breakdown_bearish"]:
            self.adaptive_threshold = base_threshold * 2.5
        elif self.current_regime in ["bearish", "bearish_low_vol"]:
            self.adaptive_threshold = base_threshold * 1.6
        elif self.current_regime in ["bullish_pullback"]:
            self.adaptive_threshold = base_threshold * 0.9
        else:  # neutral or other regimes
            self.adaptive_threshold = base_threshold * 0.75
        
        positions = {}
        
        # Group stocks by sector
        sector_data = {}
        for ticker, data in prediction_data.items():
            pred_return = data["pred_return"]
            sector = self.sector_mappings.get(ticker, "Unknown")
            
            if sector not in sector_data:
                sector_data[sector] = []
                
            sector_data[sector].append({
                "ticker": ticker,
                "pred_return": pred_return,
                "composite_score": pred_return / self.adaptive_threshold
            })
        
        # Rank sectors by predicted return
        sector_avg_scores = {}
        for sector, stocks in sector_data.items():
            sector_avg_scores[sector] = np.mean([s["pred_return"] for s in stocks])
        
        # CHANGE: Include more sectors (3-4 instead of just 2)
        ranked_sectors = sorted(sector_avg_scores.keys(), key=lambda x: sector_avg_scores[x], reverse=True)
        top_sector_count = 3 if portfolio_risk_score > 60 else 2  # More diversification in lower risk periods
        top_sectors = ranked_sectors[:min(top_sector_count, len(ranked_sectors))]
        
        # CHANGE: Allow more stocks per sector in bull markets
        stocks_per_sector = 3 if self.current_regime in ["bullish_strong", "breakout_bullish"] else 2
        
        # Allocate within top sectors - focus on stocks with strongest signals
        for sector in top_sectors:
            sector_stocks = sorted(sector_data[sector], key=lambda x: x["pred_return"], reverse=True)
            
            # Take top N stocks in each selected sector
            top_stocks = sector_stocks[:min(stocks_per_sector, len(sector_stocks))]
                        
            # CHANGE: Make position size proportional to signal strength but limited by volatility
            for stock in top_stocks:
                ticker = stock["ticker"]
                signal_strength = stock["pred_return"] / (0.2 * self.pred_std)
                
                # Base size calculation
                base_size = min(0.3, max(0.05, 0.15 * signal_strength))
                
                # Scale by portfolio risk
                final_size = base_size * risk_scaling
                
                positions[ticker] = final_size
        
        # Defensive adjustments
        if self.defensive_mode or self.current_regime in ["bearish_high_vol", "bearish_low_vol", "breakdown_bearish"]:
            # 1. Reduce overall position sizes
            scaling_factor = 0.5 if self.defensive_mode else 0.7  # More aggressive reduction
            for ticker in positions:
                positions[ticker] *= scaling_factor
            
            # 2. Add inverse positions (shorts) as hedges if we have bearish predictions
            if len(positions) > 0 and portfolio_risk_score < 40:  # Only hedge in higher risk environments
                negative_preds = {t: data["pred_return"] for t, data in prediction_data.items() 
                                if data["pred_return"] < 0 and t not in positions}
                
                if negative_preds:
                    worst_stocks = sorted(negative_preds.items(), key=lambda x: x[1])[:2]
                    for ticker, pred in worst_stocks:
                        hedge_size = -0.15 if self.defensive_mode else -0.1
                        positions[ticker] = hedge_size
        
        return positions

    def get_stop_loss_level(self):
        """Get appropriate stop-loss level based on market regime"""
        if self.current_regime in ["bullish_strong", "breakout_bullish"]:
            if self.defensive_mode:
                return -2.0  # Tighter in defensive mode
            else:
                return -3.5  # More room for positions to breathe
        elif self.current_regime in ["bearish_high_vol", "breakdown_bearish"]:
            return -1.5  # Tighter stop-loss in bearish regimes
        else:
            if self.defensive_mode:
                return -1.8
            else:
                return -2.5
    
    def update_portfolio_returns(self, daily_return):
        """Update portfolio return history"""
        self.portfolio_returns.append(daily_return)
        if len(self.portfolio_returns) > 60:  # Keep a rolling window
            self.portfolio_returns = self.portfolio_returns[-60:]
    
    def update_model_calibration(self, all_predictions):
        """Update prediction standard deviation for threshold calibration"""
        all_pred_values = [p for p in all_predictions.values()]
        if len(all_pred_values) > 5:
            self.pred_std = np.std(all_pred_values)