"""
Model inference utilities with SHAP integration.
"""
import numpy as np
import pandas as pd
import joblib
import shap
from typing import Dict, List, Optional, Tuple
import os
from preprocessing import align_schema_df, get_feature_columns


class JobFailurePredictor:
    """Wrapper for job failure prediction model with SHAP explainability."""
    
    def __init__(self, model_path: str = None,
                 background_path: str = None,
                 schema_path: str = None):
        """
        Initialize predictor.
        
        Args:
            model_path: Path to saved pipeline (defaults to env var or 'models/job_fail_pipeline_cpu.joblib')
            background_path: Path to SHAP background sample (defaults to env var or 'models/shap_background.npy')
            schema_path: Path to feature schema (defaults to env var or 'models/feature_schema.json')
        """
        # Support environment variables for flexible deployment
        model_base = os.getenv('MODEL_BASE_PATH', 'models')
        
        self.model_path = model_path or os.getenv(
            'JOB_FAIL_MODEL_PATH', 
            f'{model_base}/job_fail_pipeline_cpu.joblib'
        )
        self.background_path = background_path or os.getenv(
            'SHAP_BACKGROUND_PATH',
            f'{model_base}/shap_background.npy'
        )
        self.schema_path = schema_path or os.getenv(
            'FEATURE_SCHEMA_PATH',
            f'{model_base}/feature_schema.json'
        )
        
        if os.path.exists(self.model_path):
            self.pipeline = joblib.load(self.model_path)
            print(f"Loaded model from {self.model_path}")
        else:
            self.pipeline = None
            print(f"Warning: Model not found at {self.model_path}")
        
        self.explainer = None
        self.background = None
        self._load_shap_background()
    
    def _load_shap_background(self):
        """Load SHAP background sample if available."""
        if os.path.exists(self.background_path):
            try:
                self.background = np.load(self.background_path)
                # Use generic Explainer for compatibility
                if self.pipeline is not None:
                    self.explainer = shap.Explainer(
                        self.pipeline.predict_proba,
                        self.background,
                        feature_names=self._get_feature_names()
                    )
                print(f"Loaded SHAP background from {self.background_path}")
            except Exception as e:
                print(f"Warning: Could not load SHAP background: {e}")
                self.explainer = None
    
    def _get_feature_names(self) -> List[str]:
        """Get feature names from pipeline."""
        if self.pipeline is None:
            return []
        
        try:
            preprocess = self.pipeline.named_steps['preprocess']
            num_cols = get_feature_columns()['numeric']
            cat_cols = get_feature_columns()['categorical']
            
            # Get one-hot encoded names
            cat_encoder = preprocess.named_transformers_['cat']
            if hasattr(cat_encoder, 'get_feature_names_out'):
                cat_names = cat_encoder.get_feature_names_out(cat_cols).tolist()
            else:
                cat_names = [f"cat_{i}" for i in range(len(cat_cols))]
            
            return num_cols + cat_names
        except Exception:
            return []
    
    def predict(self, data: pd.DataFrame, explain: bool = False) -> Dict:
        """
        Predict job failure probability.
        
        Args:
            data: DataFrame with job features
            explain: Whether to compute SHAP explanations
        
        Returns:
            Dictionary with predictions and optional explanations
        """
        if self.pipeline is None:
            return {
                'fail_probability': 0.5,
                'risk_level': 'UNKNOWN',
                'error': 'Model not loaded'
            }
        
        try:
            # Align to schema
            data = align_schema_df(data, self.schema_path)
            
            # Get feature columns
            feature_cols = get_feature_columns()
            num_cols = feature_cols['numeric']
            cat_cols = feature_cols['categorical']
            
            X = data[num_cols + cat_cols].copy()
            
            # Predict
            proba = self.pipeline.predict_proba(X)[:, 1]
            fail_prob = float(proba[0]) if len(proba) == 1 else float(proba.mean())
            
            # Determine risk level
            if fail_prob >= 0.8:
                risk_level = 'CRITICAL'
            elif fail_prob >= 0.5:
                risk_level = 'MEDIUM'
            elif fail_prob >= 0.3:
                risk_level = 'LOW'
            else:
                risk_level = 'MINIMAL'
            
            result = {
                'fail_probability': fail_prob,
                'risk_level': risk_level
            }
            
            # Add SHAP explanation if requested
            if explain and self.explainer is not None:
                try:
                    # Preprocess to get encoded features
                    preprocess = self.pipeline.named_steps['preprocess']
                    X_encoded = preprocess.transform(X)
                    
                    # Compute SHAP values
                    shap_values = self.explainer(X_encoded)
                    
                    # Handle different SHAP output shapes
                    if hasattr(shap_values, 'values'):
                        sv = shap_values.values
                        if len(sv.shape) == 3:  # (n_samples, n_features, n_classes)
                            sv = sv[:, :, 1]  # Take positive class
                        elif len(sv.shape) == 2:
                            sv = sv
                        else:
                            sv = sv.flatten()
                    else:
                        sv = shap_values
                    
                    # Get feature names
                    feature_names = self._get_feature_names()
                    if len(sv.shape) == 1:
                        sv = sv.reshape(1, -1)
                    
                    # Get top drivers (absolute values)
                    if len(sv) > 0:
                        abs_sv = np.abs(sv[0])
                        top_indices = np.argsort(abs_sv)[::-1][:5]
                        
                        top_drivers = []
                        for idx in top_indices:
                            if idx < len(feature_names):
                                feature_name = feature_names[idx]
                                shap_val = float(sv[0, idx])
                                top_drivers.append({
                                    'feature': feature_name,
                                    'shap_value': shap_val,
                                    'effect': 'increase' if shap_val > 0 else 'decrease'
                                })
                        
                        result['top_drivers'] = top_drivers
                        
                        # Generate recommended actions
                        result['recommended_actions'] = self._generate_actions(
                            top_drivers, fail_prob
                        )
                except Exception as e:
                    print(f"Warning: SHAP explanation failed: {e}")
                    result['top_drivers'] = []
            
            return result
            
        except Exception as e:
            return {
                'fail_probability': 0.5,
                'risk_level': 'UNKNOWN',
                'error': str(e)
            }
    
    def _generate_actions(self, top_drivers: List[Dict], fail_prob: float) -> List[str]:
        """Generate recommended actions based on top drivers."""
        actions = []
        
        for driver in top_drivers[:3]:
            feature = driver['feature']
            effect = driver['effect']
            
            if 'failure_rate' in feature:
                actions.append("Monitor upstream dependencies and recent job history")
            elif 'duration' in feature:
                if effect == 'increase':
                    actions.append("Check for resource constraints or data volume spikes")
                else:
                    actions.append("Verify job completed successfully (unusually fast)")
            elif 'err_msg' in feature:
                actions.append("Review error logs and investigate root cause")
            elif 'job_nm' in feature or 'tasksgroup' in feature:
                actions.append("Check job configuration and dependencies")
        
        if fail_prob >= 0.8:
            actions.append("Consider immediate intervention or rerun")
        elif fail_prob >= 0.5:
            actions.append("Increase monitoring frequency")
        
        # Deduplicate
        return list(dict.fromkeys(actions))[:5]


class AnomalyDetector:
    """Wrapper for anomaly detection model."""
    
    def __init__(self, model_path: str = None,
                 scaler_path: str = None,
                 feature_path: str = None,
                 threshold_path: str = None):
        """
        Initialize anomaly detector.
        
        Args:
            model_path: Path to saved autoencoder (defaults to env var or 'models/anomaly_autoencoder_cpu.keras')
            scaler_path: Path to saved scaler (defaults to env var or 'models/anomaly_scaler.joblib')
            feature_path: Path to saved feature list (defaults to env var or 'models/anomaly_features.joblib')
            threshold_path: Path to saved threshold (defaults to env var or 'models/anomaly_threshold.joblib')
        """
        # Support environment variables for flexible deployment
        model_base = os.getenv('MODEL_BASE_PATH', 'models')
        
        self.model_path = model_path or os.getenv(
            'ANOMALY_MODEL_PATH',
            f'{model_base}/anomaly_autoencoder_cpu.keras'
        )
        self.scaler_path = scaler_path or os.getenv(
            'ANOMALY_SCALER_PATH',
            f'{model_base}/anomaly_scaler.joblib'
        )
        self.feature_path = feature_path or os.getenv(
            'ANOMALY_FEATURE_PATH',
            f'{model_base}/anomaly_features.joblib'
        )
        self.threshold_path = threshold_path or os.getenv(
            'ANOMALY_THRESHOLD_PATH',
            f'{model_base}/anomaly_threshold.joblib'
        )
        
        if os.path.exists(self.model_path):
            from tensorflow import keras
            import warnings
            warnings.filterwarnings('ignore', category=UserWarning)
            
            try:
                # Try loading with compile=False first (for inference only)
                self.model = keras.models.load_model(self.model_path, compile=False)
                print(f"Loaded autoencoder from {self.model_path}")
            except Exception as e:
                # If that fails, try with safe_mode=False (for Keras 3.x compatibility)
                try:
                    # Check if safe_mode parameter exists (Keras 3.x)
                    import inspect
                    load_model_sig = inspect.signature(keras.models.load_model)
                    if 'safe_mode' in load_model_sig.parameters:
                        self.model = keras.models.load_model(
                            self.model_path, 
                            compile=False,
                            safe_mode=False
                        )
                        print(f"Loaded autoencoder from {self.model_path} (with safe_mode=False)")
                    else:
                        raise e
                except Exception as e2:
                    # Last resort: try using tf.keras instead of keras
                    try:
                        import tensorflow as tf
                        self.model = tf.keras.models.load_model(
                            self.model_path,
                            compile=False
                        )
                        print(f"Loaded autoencoder from {self.model_path} (using tf.keras)")
                    except Exception as e3:
                        print(f"Error loading model. This might be a version compatibility issue.")
                        print(f"Error details: {str(e3)[:200]}")
                        print(f"Please ensure TensorFlow version matches the training environment.")
                        self.model = None
        else:
            self.model = None
            print(f"Warning: Model not found at {self.model_path}")
        
        if os.path.exists(self.scaler_path):
            self.scaler = joblib.load(self.scaler_path)
            print(f"Loaded scaler from {self.scaler_path}")
        else:
            self.scaler = None
        
        if os.path.exists(self.feature_path):
            self.feature_list = joblib.load(self.feature_path)
            print(f"Loaded feature list from {self.feature_path}")
        else:
            self.feature_list = get_feature_columns()['anomaly_numeric']
        
        if os.path.exists(self.threshold_path):
            self.global_threshold = joblib.load(self.threshold_path)
            print(f"Loaded threshold from {self.threshold_path}")
        else:
            self.global_threshold = 0.01
    
    def detect(self, features: Dict, threshold: Optional[float] = None) -> Dict:
        """
        Detect anomaly in feature vector.
        
        Args:
            features: Dictionary of feature values
            threshold: Optional custom threshold (overrides default)
        
        Returns:
            Dictionary with anomaly detection results
        """
        if self.model is None or self.scaler is None:
            # Ensure global_threshold exists
            threshold_value = getattr(self, 'global_threshold', 0.01)
            return {
                'reconstruction_error': 0.0,
                'is_anomaly': False,
                'threshold': float(threshold_value),
                'top_drivers': [],
                'error': 'Model not loaded'
            }
        
        try:
            # Build feature vector
            feature_vec = []
            for feat in self.feature_list:
                value = features.get(feat, 0.0)
                try:
                    feature_vec.append(float(value))
                except (ValueError, TypeError):
                    feature_vec.append(0.0)
            
            feature_vec = np.array(feature_vec).reshape(1, -1)
            
            # Scale
            feature_scaled = self.scaler.transform(feature_vec)
            
            # Reconstruct
            reconstructed = self.model.predict(feature_scaled, verbose=0)
            
            # Compute reconstruction error
            recon_error = np.mean((feature_scaled - reconstructed) ** 2)
            
            # Determine if anomaly
            thresh = threshold if threshold is not None else self.global_threshold
            is_anomaly = recon_error > thresh
            
            # Compute per-feature errors for top drivers
            per_feature_errors = np.square(feature_scaled - reconstructed).flatten()
            top_indices = np.argsort(per_feature_errors)[::-1][:3]
            
            top_drivers = []
            for idx in top_indices:
                if idx < len(self.feature_list):
                    top_drivers.append({
                        'feature': self.feature_list[idx],
                        'error': float(per_feature_errors[idx])
                    })
            
            return {
                'reconstruction_error': float(recon_error),
                'is_anomaly': bool(is_anomaly),
                'threshold': float(thresh),
                'top_drivers': top_drivers
            }
            
        except Exception as e:
            # Ensure global_threshold exists
            threshold_value = getattr(self, 'global_threshold', 0.01)
            return {
                'reconstruction_error': 0.0,
                'is_anomaly': False,
                'threshold': float(threshold_value),
                'top_drivers': [],
                'error': str(e)
            }