"""ML model definitions and persistence for strategy selection.

Implements automated model selection with:
- XGBoost
- LightGBM  
- RandomForest
- MLP Neural Network

Uses TimeSeriesSplit cross-validation and walk-forward validation.
"""

import os
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
from datetime import datetime
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import TimeSeriesSplit, cross_val_score
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import accuracy_score, top_k_accuracy_score
import joblib
import warnings

from src.core.logger import get_logger
from src.core.config import get_config

logger = get_logger(__name__)

# Try importing optional dependencies
try:
    from xgboost import XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not installed. Run: pip install xgboost")

try:
    from lightgbm import LGBMClassifier
    HAS_LIGHTGBM = True
except ImportError:
    HAS_LIGHTGBM = False
    logger.warning("LightGBM not installed. Run: pip install lightgbm")


class ModelSelector:
    """Automated model selection for strategy prediction.
    
    Trains multiple models, compares via cross-validation, picks the best.
    """
    
    def __init__(self):
        """Initialize model selector."""
        self.logger = get_logger(__name__)
        self.config = get_config()
        
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
        self.feature_names: List[str] = []
        self.strategy_names: List[str] = []
        self.is_trained = False
        self.training_metadata: Dict[str, Any] = {}
        
        # Best model after training
        self.best_model = None
        self.best_model_name: str = ""
        
        # Model storage path
        self.model_dir = Path.home() / ".local" / "share" / "crypto_trader" / "models"
        self.model_dir.mkdir(parents=True, exist_ok=True)
    
    def _get_candidate_models(self) -> Dict[str, Any]:
        """Get dictionary of candidate models to compare.
        
        Returns:
            Dictionary of model name -> model instance
        """
        models = {}
        
        # RandomForest - stable, interpretable baseline
        models['random_forest'] = RandomForestClassifier(
            n_estimators=200,
            max_depth=15,
            min_samples_split=10,
            min_samples_leaf=4,
            class_weight='balanced',
            random_state=42,
            n_jobs=-1
        )
        
        # XGBoost - typically best for tabular data
        if HAS_XGBOOST:
            models['xgboost'] = XGBClassifier(
                n_estimators=200,
                max_depth=8,
                learning_rate=0.1,
                subsample=0.8,
                colsample_bytree=0.8,
                random_state=42,
                n_jobs=-1,
                use_label_encoder=False,
                eval_metric='mlogloss'
            )
        
        # LightGBM - fast, handles noisy data well
        if HAS_LIGHTGBM:
            models['lightgbm'] = LGBMClassifier(
                n_estimators=200,
                max_depth=10,
                learning_rate=0.1,
                subsample=0.8,
                colsample_bytree=0.8,
                class_weight='balanced',
                random_state=42,
                n_jobs=-1,
                verbose=-1
            )
        
        # MLP Neural Network - high ceiling with enough data
        models['mlp'] = MLPClassifier(
            hidden_layer_sizes=(128, 64, 32),
            activation='relu',
            solver='adam',
            alpha=0.01,  # L2 regularization
            batch_size='auto',
            learning_rate='adaptive',
            learning_rate_init=0.001,
            max_iter=500,
            early_stopping=True,
            validation_fraction=0.1,
            n_iter_no_change=20,
            random_state=42
        )
        
        return models
    
    def train(
        self,
        X: pd.DataFrame,
        y: np.ndarray,
        strategy_names: List[str],
        use_ensemble: bool = False,
        n_splits: int = 5,
        training_symbols: List[str] = None
    ) -> Dict[str, Any]:
        """Train and select best model via cross-validation.
        
        Args:
            X: Feature matrix (market conditions)
            y: Target values (strategy names)
            strategy_names: List of strategy names
            use_ensemble: If True, use voting ensemble of top models instead of single best
            n_splits: Number of cross-validation splits
            training_symbols: List of symbols used for training
            
        Returns:
            Training metrics dictionary
        """
        self.strategy_names = strategy_names
        self.feature_names = list(X.columns)
        self.training_symbols = training_symbols or []
        
        # Encode labels
        y_encoded = self.label_encoder.fit_transform(y)
        
        # Scale features
        X_scaled = self.scaler.fit_transform(X)
        
        # Time series cross-validation
        tscv = TimeSeriesSplit(n_splits=n_splits)
        
        # Get candidate models
        models = self._get_candidate_models()
        
        self.logger.info(f"Training {len(models)} candidate models with {len(X)} samples...")
        
        # Evaluate each model
        model_scores = {}
        model_cv_results = {}
        
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            
            for name, model in models.items():
                try:
                    self.logger.info(f"Evaluating {name}...")
                    scores = cross_val_score(
                        model, X_scaled, y_encoded, 
                        cv=tscv, scoring='accuracy', n_jobs=-1
                    )
                    # Filter out nan scores before calculating mean
                    valid_scores = [s for s in scores if not np.isnan(s)]
                    if valid_scores:
                        mean_score = np.mean(valid_scores)
                        std_score = np.std(valid_scores) if len(valid_scores) > 1 else 0.0
                        model_scores[name] = mean_score
                        model_cv_results[name] = {
                            'mean': float(mean_score),
                            'std': float(std_score),
                            'scores': [float(s) if not np.isnan(s) else 0.0 for s in scores]
                        }
                        self.logger.info(f"  {name}: {mean_score:.3f} (+/- {std_score:.3f})")
                    else:
                        self.logger.warning(f"  {name}: All scores were nan, skipping")
                        model_scores[name] = 0.0
                except Exception as e:
                    self.logger.warning(f"  {name} failed: {e}")
                    model_scores[name] = 0.0
        
        # Filter out models with zero scores before selecting best
        valid_model_scores = {k: v for k, v in model_scores.items() if v > 0}
        
        # Select best model or create ensemble
        if not valid_model_scores:
            self.logger.warning("No models trained successfully, using RandomForest as fallback")
            self.best_model_name = "random_forest"
            self.best_model = models["random_forest"]
        elif use_ensemble and len(valid_model_scores) >= 2:
            # Use top 3 models in voting ensemble, but only if they have at least 10% accuracy
            MIN_ENSEMBLE_ACCURACY = 0.10
            good_models = {k: v for k, v in valid_model_scores.items() if v >= MIN_ENSEMBLE_ACCURACY}
            
            if len(good_models) >= 2:
                top_models = sorted(good_models.items(), key=lambda x: x[1], reverse=True)[:3]
                estimators = [(name, models[name]) for name, _ in top_models]
                
                self.best_model = VotingClassifier(estimators=estimators, voting='soft')
                self.best_model_name = "ensemble"
                self.logger.info(f"Using ensemble of: {[n for n, _ in top_models]} (excluding models below {MIN_ENSEMBLE_ACCURACY:.0%} accuracy)")
            else:
                # Not enough good models for ensemble, use single best
                self.best_model_name = max(valid_model_scores, key=valid_model_scores.get)
                self.best_model = models[self.best_model_name]
                self.logger.info(f"Not enough models for ensemble (min {MIN_ENSEMBLE_ACCURACY:.0%}), using best: {self.best_model_name}")
        else:
            # Use single best model
            self.best_model_name = max(valid_model_scores, key=valid_model_scores.get)
            self.best_model = models[self.best_model_name]
            self.logger.info(f"Best model: {self.best_model_name} ({valid_model_scores[self.best_model_name]:.3f})")
        
        # Train best model on full dataset
        self.best_model.fit(X_scaled, y_encoded)
        
        # Calculate final metrics with walk-forward validation
        train_pred = self.best_model.predict(X_scaled)
        train_acc = accuracy_score(y_encoded, train_pred)
        
        # Top-3 accuracy if we have probabilities
        top3_acc = None
        try:
            if hasattr(self.best_model, 'predict_proba'):
                proba = self.best_model.predict_proba(X_scaled)
                k = min(3, len(self.strategy_names))
                top3_acc = float(top_k_accuracy_score(y_encoded, proba, k=k))
        except Exception:
            pass
        
        # Run walk-forward validation
        wf_scores = self._walk_forward_validation(X_scaled, y_encoded)
        
        # Calculate estimated CV accuracy for ensemble (average of component models)
        if self.best_model_name == "ensemble" and len(valid_model_scores) > 0:
            # Use average of models in the ensemble for CV estimate
            ensemble_cv_acc = np.mean([v for v in valid_model_scores.values() if v >= 0.10])
        else:
            ensemble_cv_acc = model_scores.get(self.best_model_name, train_acc)
        
        metrics = {
            'train_accuracy': float(train_acc),
            'test_accuracy': float(ensemble_cv_acc),
            'cv_mean_accuracy': float(ensemble_cv_acc),
            'walk_forward_accuracy': float(np.mean(wf_scores)) if wf_scores else None,
            'top3_accuracy': top3_acc,
            'n_samples': len(X),
            'n_features': len(self.feature_names),
            'n_strategies': len(strategy_names),
            'best_model': self.best_model_name,
            'all_model_scores': model_cv_results
        }
        
        self.is_trained = True
        self.training_metadata = {
            'trained_at': datetime.utcnow().isoformat(),
            'metrics': metrics,
            'model_type': 'classifier',
            'best_model_name': self.best_model_name,
            'training_symbols': self.training_symbols if hasattr(self, 'training_symbols') else []
        }
        
        self.logger.info(f"Training complete. Best: {self.best_model_name}, "
                        f"CV accuracy: {metrics['cv_mean_accuracy']:.1%}")
        
        return metrics
    
    def _walk_forward_validation(
        self, 
        X: np.ndarray, 
        y: np.ndarray,
        train_ratio: float = 0.7,
        step_size: int = None
    ) -> List[float]:
        """Perform walk-forward validation.
        
        Train on months 1-6, test on 7. Train on 1-7, test on 8. Etc.
        
        Args:
            X: Scaled feature matrix
            y: Encoded labels
            train_ratio: Initial training set ratio
            step_size: Step size for rolling window
            
        Returns:
            List of accuracy scores for each step
        """
        n_samples = len(X)
        if n_samples < 100:
            self.logger.warning("Insufficient data for walk-forward validation")
            return []
        
        initial_train_size = int(n_samples * train_ratio)
        if step_size is None:
            step_size = max(10, n_samples // 20)
        
        scores = []
        models = self._get_candidate_models()
        model = models.get(self.best_model_name, list(models.values())[0])
        
        for i in range(initial_train_size, n_samples - step_size, step_size):
            X_train = X[:i]
            y_train = y[:i]
            X_test = X[i:i + step_size]
            y_test = y[i:i + step_size]
            
            try:
                model_clone = type(model)(**model.get_params())
                model_clone.fit(X_train, y_train)
                pred = model_clone.predict(X_test)
                scores.append(accuracy_score(y_test, pred))
            except Exception as e:
                self.logger.warning(f"Walk-forward step failed: {e}")
        
        if scores:
            self.logger.info(f"Walk-forward validation: {np.mean(scores):.3f} "
                           f"(+/- {np.std(scores):.3f}) over {len(scores)} steps")
        
        return scores
    
    def predict(
        self,
        features: Dict[str, float]
    ) -> Tuple[str, float, Dict[str, float]]:
        """Predict best strategy for given market conditions.
        
        Args:
            features: Market condition features
            
        Returns:
            Tuple of (best_strategy_name, confidence_score, all_predictions)
        """
        if not self.is_trained:
            raise ValueError("Model not trained. Call train() first.")
        
        # Convert features to DataFrame
        feature_df = pd.DataFrame([features])
        
        # Ensure all required features are present
        for feat in self.feature_names:
            if feat not in feature_df.columns:
                feature_df[feat] = 0.0
        
        # Select only required features in correct order
        X = feature_df[self.feature_names]
        
        # Scale
        X_scaled = self.scaler.transform(X)
        
        # Get probabilities for each strategy
        if hasattr(self.best_model, 'predict_proba'):
            probabilities = self.best_model.predict_proba(X_scaled)[0]
        else:
            # Fallback for models without predict_proba
            pred = self.best_model.predict(X_scaled)[0]
            probabilities = np.zeros(len(self.strategy_names))
            probabilities[pred] = 1.0
        
        # Map probabilities to strategy names
        encoded_classes = self.label_encoder.classes_
        strategy_probs = {}
        for idx, prob in enumerate(probabilities):
            if idx < len(encoded_classes):
                strategy_name = encoded_classes[idx]
                strategy_probs[strategy_name] = float(prob)
        
        # Get best strategy
        best_idx = np.argmax(probabilities)
        best_strategy = encoded_classes[best_idx]
        confidence = float(probabilities[best_idx])
        
        return best_strategy, confidence, strategy_probs
    
    def get_feature_importance(self) -> Dict[str, float]:
        """Get feature importance scores.
        
        Returns:
            Dictionary of feature names to importance scores (JSON serializable)
        """
        if not self.is_trained or self.best_model is None:
            return {}
        
        try:
            if hasattr(self.best_model, 'feature_importances_'):
                importances = self.best_model.feature_importances_
                # Convert numpy types to native Python floats for JSON serialization
                return {name: float(val) for name, val in zip(self.feature_names, importances)}
            elif hasattr(self.best_model, 'coefs_'):
                # For MLP, use absolute mean of first layer weights
                importances = np.abs(self.best_model.coefs_[0]).mean(axis=1)
                return {name: float(val) for name, val in zip(self.feature_names, importances)}
        except Exception as e:
            self.logger.warning(f"Could not get feature importance: {e}")
        
        return {}
    
    def save(self, filename: Optional[str] = None) -> str:
        """Save model to disk.
        
        Args:
            filename: Optional custom filename
            
        Returns:
            Path to saved model
        """
        if filename is None:
            timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
            filename = f"strategy_selector_{self.best_model_name}_{timestamp}.joblib"
        
        filepath = self.model_dir / filename
        
        model_data = {
            'model': self.best_model,
            'scaler': self.scaler,
            'label_encoder': self.label_encoder,
            'feature_names': self.feature_names,
            'strategy_names': self.strategy_names,
            'is_trained': self.is_trained,
            'training_metadata': self.training_metadata,
            'model_type': 'classifier',
            'best_model_name': self.best_model_name
        }
        
        joblib.dump(model_data, filepath)
        self.logger.info(f"Model saved to {filepath}")
        return str(filepath)
    
    def load(self, filepath: str) -> bool:
        """Load model from disk.
        
        Args:
            filepath: Path to model file
            
        Returns:
            True if loaded successfully
        """
        try:
            if not os.path.exists(filepath):
                self.logger.error(f"Model file not found: {filepath}")
                return False
            
            model_data = joblib.load(filepath)
            
            self.best_model = model_data['model']
            self.scaler = model_data['scaler']
            self.label_encoder = model_data.get('label_encoder', LabelEncoder())
            self.feature_names = model_data['feature_names']
            self.strategy_names = model_data['strategy_names']
            self.is_trained = model_data['is_trained']
            self.training_metadata = model_data.get('training_metadata', {})
            self.best_model_name = model_data.get('best_model_name', 'unknown')
            
            self.logger.info(f"Model loaded from {filepath}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to load model: {e}")
            return False
    
    def get_latest_model_path(self) -> Optional[str]:
        """Get path to latest saved model.
        
        Returns:
            Path to latest model or None
        """
        model_files = list(self.model_dir.glob("strategy_selector_*.joblib"))
        if not model_files:
            return None
        
        # Sort by modification time
        latest = max(model_files, key=lambda p: p.stat().st_mtime)
        return str(latest)


# Backwards-compatible alias
class StrategySelectorModel(ModelSelector):
    """Backwards-compatible alias for ModelSelector.
    
    Now uses automated model selection instead of just RandomForest.
    """
    
    def __init__(self, model_type: str = "classifier"):
        """Initialize model.
        
        Args:
            model_type: Model type (only 'classifier' supported now)
        """
        super().__init__()
        self.model_type = model_type
        if model_type != "classifier":
            self.logger.warning("Only 'classifier' model type is supported. Using classifier.")
