"""
ML-based email classifier.

MOCK STATUS: This uses a mock Random Forest model for demonstration.
- IMPORTANT: This is NOT the production model
- IMPORTANT: This model is trained on synthetic/demo data only
- IMPORTANT: For production, you MUST train on real Enron dataset or your own emails
- The mock will work for testing but will NOT achieve 94-96% accuracy
- When you get home: retrain with LightGBM + Enron + real tuning

DO NOT use this mock model on production data.
DO NOT expect accurate classifications from this mock.
This is ONLY for framework testing and integration validation.
"""
import logging
import pickle
import warnings
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)

# Suppress warnings
warnings.filterwarnings('ignore')


class MLClassifier:
    """
    Wrapper for ML-based email classification.

    MOCK IMPLEMENTATION: Uses Random Forest on synthetic data.
    Replace this with real LightGBM model trained on Enron dataset.
    """

    def __init__(self, model_path: Optional[str] = None):
        """Initialize ML classifier."""
        self.model = None
        self.categories = []
        self.feature_names = []
        self.is_mock = False
        self.model_path = model_path or "src/models/pretrained/classifier.pkl"

        # Try to load pre-trained model
        logger.info(f"Checking for model at: {self.model_path}")
        if Path(self.model_path).exists():
            logger.info(f"Model file found, loading...")
            self._load_model(self.model_path)
        else:
            logger.warning(f"Pre-trained model not found at {self.model_path}, creating MOCK model for testing")
            self._create_mock_model()

    def _load_model(self, model_path: str) -> None:
        """Load pre-trained model from file."""
        try:
            logger.info(f"Loading ML model from: {model_path}")
            with open(model_path, 'rb') as f:
                model_data = pickle.load(f)

            self.model = model_data.get('model')
            self.categories = model_data.get('categories', [])
            self.feature_names = model_data.get('feature_names', [])
            self.is_mock = model_data.get('is_mock', False)

            if self.is_mock:
                logger.warning("USING MOCK ML MODEL - Not for production!")

            logger.info(f"ML model loaded: {len(self.categories)} categories")

        except Exception as e:
            logger.error(f"Error loading model: {e}")
            logger.warning("Falling back to MOCK model")
            self._create_mock_model()

    def _create_mock_model(self) -> None:
        """Create a mock Random Forest model for testing."""
        try:
            from sklearn.ensemble import RandomForestClassifier
        except ImportError:
            logger.error("scikit-learn required for mock model")
            self.model = None
            return

        logger.info("Creating MOCK Random Forest model for framework testing")
        logger.warning("=" * 80)
        logger.warning("MOCK MODEL WARNING")
        logger.warning("=" * 80)
        logger.warning("This is a MOCK model trained on synthetic data")
        logger.warning("It will NOT provide accurate classifications")
        logger.warning("Framework testing ONLY - do not use for production")
        logger.warning("When you get home: Retrain with LightGBM + Enron dataset")
        logger.warning("=" * 80)

        # Define categories
        self.categories = [
            'junk', 'transactional', 'auth', 'newsletters',
            'social', 'automated', 'conversational', 'work',
            'personal', 'finance', 'travel', 'unknown'
        ]

        # Create synthetic training data
        n_samples = 500
        n_features = 50  # This will be expanded by feature extractor

        X_mock = np.random.rand(n_samples, n_features)
        y_mock = np.random.randint(0, len(self.categories), n_samples)

        # Train mock model
        try:
            self.model = RandomForestClassifier(
                n_estimators=50,
                max_depth=15,
                random_state=42,
                n_jobs=-1
            )
            self.model.fit(X_mock, y_mock)

            self.feature_names = [f'feature_{i}' for i in range(n_features)]
            self.is_mock = True

            logger.info(f"Mock model created: {len(self.categories)} categories, {n_features} features")
            logger.warning("REMEMBER: This mock model is for framework testing only!")

        except Exception as e:
            logger.error(f"Error creating mock model: {e}")
            self.model = None

    def predict(self, features: np.ndarray) -> Dict[str, Any]:
        """
        Predict category for email features.

        Args:
            features: Feature vector (numpy array)

        Returns:
            {
                'category': str (predicted category),
                'confidence': float (0-1),
                'probabilities': Dict[category -> confidence],
                'is_mock': bool (True if using mock model),
                'warning': str (warning if using mock)
            }
        """
        if self.model is None:
            logger.error("Model not loaded")
            return {
                'category': 'unknown',
                'confidence': 0.0,
                'probabilities': {c: 0.0 for c in self.categories},
                'error': 'Model not available',
                'is_mock': False
            }

        try:
            # Ensure feature vector is correct shape
            if len(features.shape) == 1:
                features = features.reshape(1, -1)

            # Get probabilities - handle both LightGBM and sklearn models
            if hasattr(self.model, 'predict_proba'):
                # sklearn API (RandomForest, etc.)
                probs = self.model.predict_proba(features)[0]
            else:
                # LightGBM API (Booster object)
                probs = self.model.predict(features)[0]

            pred_class = np.argmax(probs)
            category = self.categories[pred_class]
            confidence = float(probs[pred_class])

            # Build probabilities dict
            prob_dict = {
                self.categories[i]: float(probs[i])
                for i in range(len(self.categories))
            }

            result = {
                'category': category,
                'confidence': confidence,
                'probabilities': prob_dict,
                'is_mock': self.is_mock
            }

            if self.is_mock:
                result['warning'] = 'Using mock model - inaccurate for production'

            return result

        except Exception as e:
            logger.error(f"Prediction error: {e}")
            return {
                'category': 'unknown',
                'confidence': 0.0,
                'probabilities': {c: 0.0 for c in self.categories},
                'error': str(e),
                'is_mock': self.is_mock
            }

    def predict_batch(self, features: np.ndarray) -> List[Dict[str, Any]]:
        """Predict for batch of feature vectors."""
        return [self.predict(f if f.ndim > 1 else f.reshape(1, -1)) for f in features]

    def save_model(self, output_path: str) -> bool:
        """Save current model to file."""
        if self.model is None:
            logger.error("No model to save")
            return False

        try:
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)

            model_data = {
                'model': self.model,
                'categories': self.categories,
                'feature_names': self.feature_names,
                'is_mock': self.is_mock,
                'warning': 'Mock model - for testing only' if self.is_mock else 'Production model'
            }

            with open(output_path, 'wb') as f:
                pickle.dump(model_data, f)

            logger.info(f"Model saved to: {output_path}")
            return True

        except Exception as e:
            logger.error(f"Error saving model: {e}")
            return False

    def get_info(self) -> Dict[str, Any]:
        """Get model information."""
        return {
            'is_loaded': self.model is not None,
            'is_mock': self.is_mock,
            'categories': self.categories,
            'n_categories': len(self.categories),
            'n_features': len(self.feature_names),
            'feature_names': self.feature_names[:10],  # First 10 for brevity
            'model_type': 'RandomForest' if self.is_mock else 'LightGBM (production)'
        }
