"""ML-based strategy selector for intelligent autopilot."""

import asyncio
from typing import Dict, Any, Optional, List, Tuple

from src.core.logger import get_logger
from src.core.config import get_config
from .market_analyzer import MarketConditions, MarketRegime, get_market_analyzer
from .models import StrategySelectorModel
from .performance_tracker import get_performance_tracker
from .strategy_groups import (
    StrategyGroup, get_strategy_group, get_strategies_in_group,
    get_best_strategy_in_group, convert_strategy_to_group_label, get_all_groups
)
from src.strategies import get_strategy_registry

logger = get_logger(__name__)


class StrategySelector:
    """ML-based strategy selector."""
    
    def __init__(self):
        """Initialize strategy selector."""
        self.logger = get_logger(__name__)
        self.config = get_config()
        self.market_analyzer = get_market_analyzer()
        self.performance_tracker = get_performance_tracker()
        self.model = StrategySelectorModel(model_type="classifier")
        self.strategy_registry = get_strategy_registry()
        self._available_strategies: List[str] = []
        
        # Bootstrap configuration - improved defaults for better ML accuracy
        self.bootstrap_days = self.config.get("autopilot.intelligent.bootstrap.days", 365)
        self.bootstrap_timeframe = self.config.get("autopilot.intelligent.bootstrap.timeframe", "1h")
        self.min_samples_per_strategy = self.config.get("autopilot.intelligent.bootstrap.min_samples_per_strategy", 10)
        self.bootstrap_symbols = self.config.get("autopilot.intelligent.bootstrap.symbols", ["BTC/USD", "ETH/USD"])
        self.bootstrap_timeframes = ["15m", "1h", "4h"]  # Multi-timeframe for more data
        
        self._available_strategies: List[str] = []
        self._load_available_strategies()
        
        self._last_model_ts = 0.0
        
        # Auto-load persisted model if available
        self._try_load_saved_model()
    
    def _try_load_saved_model(self):
        """Try to load a previously saved model."""
        import os
        try:
            latest_model = self.model.get_latest_model_path()
            if latest_model:
                mtime = os.path.getmtime(latest_model)
                # Only load if it's new or we haven't loaded anything
                if mtime > self._last_model_ts:
                    if self.model.load(latest_model):
                        self._last_model_ts = mtime
                        self.logger.info(f"Loaded persisted model: {latest_model}")
                    else:
                         self.logger.warning(f"Failed to load model {latest_model}")
            else:
                if self._last_model_ts == 0:
                    self.logger.info("No saved model found, model needs training")
        except Exception as e:
            self.logger.warning(f"Failed to load saved model: {e}")
    
    def _load_available_strategies(self):
        """Load list of available strategies."""
        self._available_strategies = self.strategy_registry.list_available()
        self.logger.info(f"Available strategies: {self._available_strategies}")
    
    async def train_model(
        self,
        min_samples_per_strategy: int = 10,
        force_retrain: bool = False
    ) -> Dict[str, Any]:
        """Train the ML model on historical performance data.
        
        Args:
            min_samples_per_strategy: Minimum samples required per strategy
            force_retrain: Force retraining even if model exists
            
        Returns:
            Training metrics dictionary
        """
        # Try to load existing model first
        if not force_retrain:
            latest_model = self.model.get_latest_model_path()
            if latest_model and self.model.load(latest_model):
                self.logger.info("Loaded existing trained model")
                return self.model.training_metadata.get('metrics', {})
        
        # Prepare training data
        training_data = await self.performance_tracker.prepare_training_data(
            min_samples_per_strategy=min_samples_per_strategy
        )
        
        if training_data is None:
            self.logger.warning("Insufficient training data. Using fallback rule-based selection.")
            return {}
        
        X = training_data['X']
        y = training_data['y']
        strategy_names = training_data['strategy_names']
        
        # Convert individual strategy names to group labels for better ML accuracy
        # This reduces the number of classes from 14 to 5
        y_groups = [convert_strategy_to_group_label(name) for name in y]
        group_names = [g.value for g in get_all_groups()]
        
        self.logger.info(f"Training with {len(set(y_groups))} strategy groups (from {len(strategy_names)} strategies)")
        
        if len(X) < min_samples_per_strategy * len(group_names):
            self.logger.warning("Insufficient training data. Using fallback rule-based selection.")
            return {}
        
        # Train model with ensemble mode for better accuracy
        # Ensemble combines XGBoost + LightGBM + RandomForest via voting
        # Now predicts GROUPS instead of individual strategies
        metrics = self.model.train(
            X, 
            y_groups,  # Use group labels
            group_names,  # Use group names
            use_ensemble=True,
            training_symbols=training_data.get('training_symbols', [])
        )
        
        # Save model
        self.model.save()
        
        return metrics
    
    def select_best_strategy(
        self,
        market_conditions: MarketConditions,
        min_confidence: float = 0.5
    ) -> Optional[Tuple[str, float, Dict[str, float]]]:
        """Select best strategy for current market conditions.
        
        The ML model predicts a strategy GROUP, then we use rule-based logic
        to select the best individual strategy within that group.
        
        Args:
            market_conditions: Current market conditions
            min_confidence: Minimum confidence threshold
            
        Returns:
            Tuple of (strategy_name, confidence, all_predictions) or None
        """
        # If model not trained, use fallback
        if not self.model.is_trained:
            self.logger.debug("Model not trained, using rule-based fallback")
            return self._fallback_strategy_selection(market_conditions)
        
        try:
            # Predict strategy GROUP using ML model
            predicted_group, group_confidence, all_group_predictions = self.model.predict(
                market_conditions.features
            )
            
            if group_confidence < min_confidence:
                self.logger.debug(
                    f"ML group confidence {group_confidence:.2f} below threshold {min_confidence}, "
                    "using fallback"
                )
                return self._fallback_strategy_selection(market_conditions)
            
            # Convert group string to enum
            try:
                group_enum = StrategyGroup(predicted_group)
            except ValueError:
                self.logger.warning(f"Unknown group '{predicted_group}', using fallback")
                return self._fallback_strategy_selection(market_conditions)
            
            # Select best individual strategy within the predicted group
            best_strategy, strategy_confidence = get_best_strategy_in_group(
                group_enum,
                market_conditions.features,
                available_strategies=self._available_strategies
            )
            
            # Combined confidence: group confidence * strategy confidence
            combined_confidence = group_confidence * strategy_confidence
            
            # Build all_predictions dict with individual strategies
            all_predictions = {}
            for group_name, group_score in all_group_predictions.items():
                try:
                    group_e = StrategyGroup(group_name)
                    strategies_in_group = get_strategies_in_group(group_e)
                    for strat in strategies_in_group:
                        if strat in self._available_strategies:
                            all_predictions[strat] = group_score
                except ValueError:
                    pass
            
            self.logger.info(
                f"ML selected group: {predicted_group} (conf: {group_confidence:.2f}) "
                f"-> strategy: {best_strategy} (combined conf: {combined_confidence:.2f})"
            )
            
            return best_strategy, combined_confidence, all_predictions
            
        except Exception as e:
            self.logger.error(f"Error in ML prediction: {e}")
            return self._fallback_strategy_selection(market_conditions)
    
    def _fallback_strategy_selection(
        self,
        market_conditions: MarketConditions
    ) -> Optional[Tuple[str, float, Dict[str, float]]]:
        """Fallback rule-based strategy selection.
        
        Args:
            market_conditions: Current market conditions
            
        Returns:
            Tuple of (strategy_name, confidence, all_predictions)
        """
        features = market_conditions.features
        regime = market_conditions.regime
        
        # Rule-based selection based on market regime
        strategy_rules = {
            MarketRegime.TRENDING_UP: "moving_average",
            MarketRegime.TRENDING_DOWN: "moving_average",
            MarketRegime.RANGING: "rsi",
            MarketRegime.HIGH_VOLATILITY: "momentum",
            MarketRegime.LOW_VOLATILITY: "grid",
            MarketRegime.BREAKOUT: "momentum",
            MarketRegime.REVERSAL: "rsi"
        }
        
        # Map regime enum to string
        regime_str = regime.value if hasattr(regime, 'value') else str(regime)
        
        # Select strategy based on regime
        selected_strategy = strategy_rules.get(regime, "rsi")
        
        # Check if strategy is available
        if selected_strategy not in self._available_strategies:
            # Fallback to first available
            if self._available_strategies:
                selected_strategy = self._available_strategies[0]
            else:
                return None
        
        # Calculate confidence based on regime clarity
        rsi = features.get('rsi', 50.0)
        adx = features.get('adx', 0.0)
        
        # Higher confidence for clear signals
        if regime in [MarketRegime.TRENDING_UP, MarketRegime.TRENDING_DOWN] and adx > 25:
            confidence = 0.7
        elif regime == MarketRegime.RANGING and (rsi < 30 or rsi > 70):
            confidence = 0.65
        else:
            confidence = 0.5
        
        all_predictions = {selected_strategy: confidence}
        
        self.logger.info(
            f"Fallback selected strategy: {selected_strategy} "
            f"(confidence: {confidence:.2f}, regime: {regime_str})"
        )
        
        return selected_strategy, confidence, all_predictions
    
    def get_strategy_rankings(
        self,
        market_conditions: MarketConditions
    ) -> List[Tuple[str, float]]:
        """Get all strategies ranked by expected performance.
        
        Args:
            market_conditions: Current market conditions
            
        Returns:
            List of (strategy_name, score) tuples, sorted by score descending
        """
        result = self.select_best_strategy(market_conditions, min_confidence=0.0)
        
        if result is None:
            return []
        
        _, _, all_predictions = result
        
        # Sort by score descending
        rankings = sorted(
            all_predictions.items(),
            key=lambda x: x[1],
            reverse=True
        )
        
        return rankings
    
    def update_model(self, trade_result: Dict[str, Any]) -> bool:
        """Update model with new trade result (incremental learning).
        
        Args:
            trade_result: Trade result with performance metrics
            
        Returns:
            True if update successful
        """
        # For now, we'll retrain periodically rather than incremental updates
        # This can be enhanced later with online learning algorithms
        self.logger.debug("Model update requested (will retrain on next cycle)")
        return True
    
    async def bootstrap_training_data(
        self,
        symbol: str = "BTC/USD",
        timeframe: Optional[str] = None,
        days: Optional[int] = None,
        exchange_name: str = "Binance Public"
    ) -> Dict[str, Any]:
        """Bootstrap training data by running backtests on historical data.
        
        Args:
            symbol: Trading symbol
            timeframe: Timeframe (defaults to config value)
            days: Number of days of historical data (defaults to config value)
            exchange_name: Exchange name
            
        Returns:
            Dictionary with bootstrap results
        """
        # Use config values as defaults
        if timeframe is None:
            timeframe = self.bootstrap_timeframe
        if days is None:
            days = self.bootstrap_days
        from src.backtesting.engine import BacktestingEngine
        from src.exchanges.public_data import PublicDataAdapter
        from src.data.collector import get_data_collector
        from src.core.database import get_database, MarketData
        from datetime import datetime, timedelta
        from decimal import Decimal
        from sqlalchemy import select, func
        import time
        
        self.logger.info(f"Bootstrapping training data for {symbol} ({timeframe})")
        
        db = get_database()
        
        # Step 1: Check and fetch historical data
        async with db.get_session() as session:
            try:
                end_date = datetime.utcnow()
                start_date = end_date - timedelta(days=days)
                
                # Check if we have data
                stmt = select(func.count()).select_from(MarketData).where(
                    MarketData.exchange == exchange_name,
                    MarketData.symbol == symbol,
                    MarketData.timeframe == timeframe,
                    MarketData.timestamp >= start_date
                )
                result = await session.execute(stmt)
                existing_data = result.scalar()
                
                if existing_data < 100:  # Need at least 100 candles
                    self.logger.info(f"Fetching historical data ({days} days)...")
                    adapter = PublicDataAdapter()
                    if await adapter.connect():
                        collector = get_data_collector()
                        current_date = start_date
                        chunk_days = 30
                        
                        while current_date < end_date:
                            chunk_end = min(current_date + timedelta(days=chunk_days), end_date)
                            ohlcv = await adapter.get_ohlcv(
                                symbol=symbol,
                                timeframe=timeframe,
                                since=current_date,
                                limit=1000
                            )
                            if ohlcv:
                                await collector.store_ohlcv(exchange_name, symbol, timeframe, ohlcv)
                            current_date = chunk_end
                            time.sleep(1)
                        
                        await adapter.disconnect()
                        self.logger.info("Historical data fetched")
                    else:
                        self.logger.error("Failed to connect to exchange")
                        return {"error": "Failed to fetch historical data"}
            except Exception as e:
                self.logger.error(f"Error checking/fetching historical data: {e}")
                return {"error": f"Database error: {e}"}
        
        # Step 2: Run backtests for each strategy
        backtest_engine = BacktestingEngine()
        bootstrap_results = []
        
        self.logger.info(f"Available strategies for bootstrap: {self._available_strategies}")
        
        total_strategies = len(self._available_strategies)
        for strategy_idx, strategy_name in enumerate(self._available_strategies):
            try:
                strategy_class = self.strategy_registry._strategies.get(strategy_name.lower())
                if not strategy_class:
                    continue
                
                strategy = strategy_class(
                    name=strategy_name,
                    parameters={},
                    timeframes=[timeframe]
                )
                strategy.enabled = True
                
                self.logger.info(f"Running backtest for {strategy_name} ({strategy_idx + 1}/{total_strategies})...")
                
                # Run backtest
                results = await backtest_engine.run_backtest(
                    strategy=strategy,
                    symbol=symbol,
                    exchange=exchange_name,
                    timeframe=timeframe,
                    start_date=start_date,
                    end_date=end_date,
                    initial_capital=Decimal("10000.0"),
                    slippage=0.001
                )
                
                if "error" in results:
                    self.logger.warning(f"Backtest failed for {strategy_name}: {results['error']}")
                    continue
                
                self.logger.info(f"Backtest successful for {strategy_name}, return: {results.get('total_return', 0.0)}%")
                
                # Get market conditions for the period
                market_data = await backtest_engine._get_historical_data(
                    exchange_name, symbol, timeframe, start_date, end_date
                )
                
                if len(market_data) > 0:
                    # =====================================================
                    # IMPROVED SAMPLING: Regime-change detection + time spacing
                    # =====================================================
                    # This creates more diverse, independent training samples
                    # by sampling at meaningful market transitions rather than
                    # every N candles (which creates nearly-identical samples)
                    
                    # Minimum time spacing between samples (in candles)
                    # For 1h: 24 candles = 24 hours
                    # For 15m: 96 candles = 24 hours
                    # For 4h: 6 candles = 24 hours
                    timeframe_to_spacing = {
                        "1m": 1440,  # 24 hours
                        "5m": 288,
                        "15m": 96,
                        "30m": 48,
                        "1h": 24,
                        "4h": 6,
                        "1d": 1
                    }
                    min_spacing = timeframe_to_spacing.get(timeframe, 24)
                    
                    samples_recorded = 0
                    last_regime = None
                    last_sample_idx = -min_spacing  # Allow first sample immediately
                    
                    # Limit processing to prevent excessive computation
                    # For small datasets (5 days), process all points but yield periodically
                    data_points = len(market_data) - 50
                    self.logger.info(f"Processing {data_points} data points for {strategy_name}...")
                    
                    # Need at least 50 candles for feature calculation
                    for i in range(50, len(market_data)):
                        # Yield control periodically to prevent blocking (every 10 iterations)
                        if i % 10 == 0:
                            await asyncio.sleep(0)  # Yield to event loop
                        
                        sample_data = market_data.iloc[i-50:i]
                        conditions = self.market_analyzer.analyze_current_conditions(
                            symbol, timeframe, sample_data
                        )
                        current_regime = conditions.regime
                        
                        # Determine if we should sample at this point
                        # 1. Sample on regime change (market transition = valuable data)
                        # 2. Sample after minimum time spacing (ensure independence)
                        regime_changed = (last_regime is not None and current_regime != last_regime)
                        time_elapsed = (i - last_sample_idx) >= min_spacing
                        
                        should_sample = regime_changed or (time_elapsed and last_regime is None) or (time_elapsed and i == 50)
                        
                        # Also sample periodically even without regime change (every 2x min_spacing)
                        periodic_sample = (i - last_sample_idx) >= (min_spacing * 2)
                        
                        if should_sample or periodic_sample:
                            # Record as training data
                            trade_result = {
                                'return_pct': results.get('total_return', 0.0) / 100.0,
                                'sharpe_ratio': results.get('sharpe_ratio', 0.0),
                                'win_rate': results.get('win_rate', 0.5),
                                'max_drawdown': abs(results.get('max_drawdown', 0.0)) / 100.0,
                                'trade_count': results.get('total_trades', 0)
                            }
                            
                            await self.performance_tracker.record_trade(
                                strategy_name=strategy_name,
                                market_conditions=conditions,
                                trade_result=trade_result
                            )
                            samples_recorded += 1
                            last_sample_idx = i
                            
                            if regime_changed:
                                self.logger.debug(
                                    f"Sampled at regime change: {last_regime} -> {current_regime}"
                                )
                        
                        last_regime = current_regime
                    
                    self.logger.info(f"Recorded {samples_recorded} samples for {strategy_name}")
                    
                    bootstrap_results.append({
                        'strategy': strategy_name,
                        'trades_recorded': samples_recorded,
                        'backtest_return': results.get('total_return', 0.0)
                    })
                    
            except Exception as e:
                self.logger.error(f"Error bootstrapping {strategy_name}: {e}", exc_info=True)
                continue
        
        total_samples = sum(r['trades_recorded'] for r in bootstrap_results)
        self.logger.info(f"Bootstrap complete: {total_samples} training samples created")
        
        return {
            'status': 'success',
            'strategies': bootstrap_results,
            'total_samples': total_samples
        }
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get information about the current model.
        
        Returns:
            Dictionary with model information
        """
        # Always check if a newer model is available on disk
        self._try_load_saved_model()
        
        info = {
            'is_trained': self.model.is_trained,
            'model_type': self.model.model_type,
            'available_strategies': self._available_strategies,
            'feature_count': len(self.model.feature_names) if self.model.is_trained else 0
        }
        
        if self.model.is_trained:
            info.update({
                'training_metadata': self.model.training_metadata,
                'feature_importance': dict(
                    sorted(
                        self.model.get_feature_importance().items(),
                        key=lambda x: x[1],
                        reverse=True
                    )[:10]  # Top 10 features
                )
            })
        
        return info


# Global instance
_strategy_selector: Optional[StrategySelector] = None


def get_strategy_selector() -> StrategySelector:
    """Get global strategy selector instance."""
    global _strategy_selector
    if _strategy_selector is None:
        _strategy_selector = StrategySelector()
    return _strategy_selector

