"""
Production SOH Prediction Function

predict_soh(battery_data_up_to_cycle_n) -> returns:
- Current SOH (point estimate)
- Future SOH at 10, 50, 100 cycles (with 95% prediction intervals)
- RUL to 80% SOH (with confidence range)
- Data quality score (completeness, noise level)
"""

import pandas as pd
import numpy as np
import pickle
import os
from typing import Dict, List, Tuple, Optional
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

try:
    import shap
    SHAP_AVAILABLE = True
except ImportError:
    SHAP_AVAILABLE = False
    shap = None

from feature_engineering_causal import engineer_features_causal
from feature_engineering_production import engineer_features_production, engineer_features_iterative


class ConformalPredictor:
    """
    Conformal Prediction for uncertainty quantification.
    Uses calibration set residuals to compute prediction intervals.
    """
    
    def __init__(self, alpha=0.05):
        """
        Parameters:
        -----------
        alpha : float
            Significance level (default 0.05 = 95% confidence)
        """
        self.alpha = alpha
        self.calibration_scores = None
    
    def calibrate(self, y_true, y_pred):
        """
        Calibrate using validation set residuals.
        
        Parameters:
        -----------
        y_true : array
            True SOH values
        y_pred : array
            Predicted SOH values
        """
        residuals = np.abs(y_true - y_pred)
        self.calibration_scores = residuals
    
    def predict_interval(self, y_pred, coverage=0.95):
        """
        Compute prediction intervals using conformal prediction.
        
        Parameters:
        -----------
        y_pred : array
            Point predictions
        coverage : float
            Desired coverage (default 0.95 = 95%)
        
        Returns:
        --------
        tuple: (lower, upper) prediction bounds
        """
        if self.calibration_scores is None:
            raise ValueError("Must calibrate before predicting intervals")
        
        alpha = 1 - coverage
        quantile = np.quantile(self.calibration_scores, 1 - alpha)
        
        lower = y_pred - quantile
        upper = y_pred + quantile
        
        return lower, upper


def calculate_data_quality(df: pd.DataFrame) -> Dict[str, float]:
    """
    Calculate data quality metrics.
    
    Parameters:
    -----------
    df : DataFrame
        Input battery data
    
    Returns:
    --------
    dict: Quality metrics (completeness, noise_level, etc.)
    """
    required_cols = ['cycle_index', 'voltage', 'current', 'temperature', 'capacity']
    
    completeness = 0.0
    for col in required_cols:
        if col in df.columns:
            completeness += (1 - df[col].isna().sum() / len(df)) / len(required_cols)
    
    # Noise level: coefficient of variation
    noise_level = 0.0
    if 'capacity' in df.columns:
        cv = df['capacity'].std() / df['capacity'].mean() if df['capacity'].mean() > 0 else 0
        noise_level = min(1.0, cv * 2)  # Normalize to 0-1
    
    quality_score = (completeness * 0.7 + (1 - noise_level) * 0.3)
    
    return {
        'completeness': completeness,
        'noise_level': noise_level,
        'quality_score': quality_score
    }


def predict_soh(
    battery_data: pd.DataFrame,
    model_path: str = './models/nasa_soh_model.pkl',
    pipeline_path: str = './models/feature_pipeline_causal.pkl',
    initial_capacity: Optional[float] = None,
    forecast_horizons: List[int] = [10, 50, 100]
) -> Dict:
    """
    Production SOH prediction function.
    
    Parameters:
    -----------
    battery_data : DataFrame
        BMS data up to current cycle (columns: cycle_index, voltage, current, temperature, capacity)
    model_path : str
        Path to trained model
    pipeline_path : str
        Path to feature pipeline
    initial_capacity : float, optional
        Initial battery capacity (if None, uses max of capacity column)
    forecast_horizons : list
        Cycles ahead to forecast (default [10, 50, 100])
    
    Returns:
    --------
    dict: Predictions with uncertainty intervals
    """
    # Load model and pipeline
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model not found: {model_path}")
    
    with open(model_path, 'rb') as f:
        model_data = pickle.load(f)
        model = model_data['model']
        feature_cols = model_data['feature_cols']
    
    # Check data quality
    quality = calculate_data_quality(battery_data)
    
    # Get initial capacity
    if initial_capacity is None:
        if 'capacity' in battery_data.columns:
            initial_capacity = battery_data['capacity'].max()
        else:
            raise ValueError("Capacity column required or provide initial_capacity")
    
    # Engineer causal features for production (handles missing SOH)
    # Use iterative approach if SOH-dependent features needed
    if 'soh' not in battery_data.columns:
        # Production prediction: SOH not available, use iterative prediction
        features = engineer_features_iterative(
            battery_data,
            initial_capacity,
            model,
            feature_cols,
            max_iterations=3
        )
    else:
        # Training/validation: SOH available
        features = engineer_features_causal(battery_data, initial_capacity, use_fixed_bounds=True)
    
    if features is None or len(features) == 0:
        raise ValueError("Feature engineering failed")
    
    # Ensure all required features are present
    missing_features = [f for f in feature_cols if f not in features.columns]
    if missing_features:
        # Fill missing features with 0 (shouldn't happen, but safety)
        for feat in missing_features:
            features[feat] = 0.0
    
    # Get current cycle features
    current_features = features[feature_cols].values[-1:]  # Last cycle
    current_cycle = features['cycle_index'].values[-1]
    
    # Current SOH prediction
    current_soh_pred = model.predict(current_features)[0]
    
    # Enhanced uncertainty: Conformal prediction + Tree variance
    # Load calibration data from model if available
    conformal_predictor = None
    if 'conformal_calibration_scores' in model_data:
        conformal_predictor = ConformalPredictor(alpha=0.05)
        conformal_predictor.calibration_scores = np.array(model_data['conformal_calibration_scores'])
    
    # Prediction intervals: Use conformal if available, else tree variance
    if hasattr(model, 'estimators_'):
        tree_preds = np.array([tree.predict(current_features)[0] for tree in model.estimators_])
        current_std_tree = np.std(tree_preds)
    else:
        current_std_tree = 0.01
    
    # Conformal prediction (guaranteed coverage)
    if conformal_predictor is not None:
        try:
            current_lower_conf, current_upper_conf = conformal_predictor.predict_interval(
                np.array([current_soh_pred]), coverage=0.95
            )
            current_lower = current_lower_conf[0]
            current_upper = current_upper_conf[0]
            # Use tighter of the two intervals (conformal usually more conservative)
            current_std = (current_upper - current_lower) / (2 * 1.96)
        except:
            # Fallback to tree variance
            current_std = current_std_tree
            current_lower = current_soh_pred - 1.96 * current_std
            current_upper = current_soh_pred + 1.96 * current_std
    else:
        # Use tree variance (empirical, but works well)
        current_std = current_std_tree
        current_lower = current_soh_pred - 1.96 * current_std
        current_upper = current_soh_pred + 1.96 * current_std
    
    # SHAP values for explainability
    shap_values_dict = {}
    if SHAP_AVAILABLE and hasattr(model, 'predict'):
        try:
            # Use TreeExplainer for Random Forest (fast)
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(current_features)
            
            # For regression, shap_values might be a list (per tree) or array
            if isinstance(shap_values, list):
                shap_values = np.array(shap_values).mean(axis=0) if len(shap_values) > 0 else shap_values[0]
            
            # Get SHAP values for the single prediction (last cycle)
            if len(shap_values.shape) > 1:
                shap_values_flat = shap_values[-1]  # Last cycle
            else:
                shap_values_flat = shap_values
            
            # Create feature importance dict (top 10)
            feature_shap_pairs = list(zip(feature_cols, shap_values_flat))
            feature_shap_pairs.sort(key=lambda x: abs(x[1]), reverse=True)
            
            shap_values_dict = {
                'top_features': [
                    {'feature': feat, 'shap_value': float(val), 'importance': float(abs(val))}
                    for feat, val in feature_shap_pairs[:10]
                ],
                'baseline': float(model_data.get('shap_baseline', current_soh_pred)),
                'prediction_contribution': float(sum(shap_values_flat))
            }
        except Exception as e:
            # SHAP failed, continue without it
            shap_values_dict = {'error': str(e)}
    
    # Future forecasts (simple extrapolation based on degradation rate)
    forecasts = {}
    
    # Calculate degradation rate from recent history
    if 'soh' in features.columns and len(features) > 5:
        recent_soh = features['soh'].values[-10:]
        degradation_rate = np.mean(np.diff(recent_soh)) if len(recent_soh) > 1 else -0.001
    else:
        degradation_rate = -0.001  # Default: 0.1% per cycle
    
    # Generate forecasts
    for horizon in forecast_horizons:
        # Simple linear extrapolation (can be improved with more sophisticated methods)
        forecast_soh = current_soh_pred + (degradation_rate * horizon)
        forecast_soh = max(0.0, min(1.0, forecast_soh))  # Clip to [0, 1]
        
        # Uncertainty grows with forecast horizon (sqrt scaling)
        forecast_std = current_std * np.sqrt(1 + horizon / 10)
        forecast_lower = forecast_soh - 1.96 * forecast_std
        forecast_upper = forecast_soh + 1.96 * forecast_std
        
        forecasts[f'soh_at_{horizon}_cycles'] = {
            'point_estimate': float(forecast_soh),
            'lower_95': float(max(0.0, forecast_lower)),
            'upper_95': float(min(1.0, forecast_upper))
        }
    
    # RUL calculation (cycles until 80% SOH)
    eol_threshold = 0.80
    if current_soh_pred <= eol_threshold:
        rul_cycles = 0
        rul_lower = 0
        rul_upper = 0
    else:
        if degradation_rate < 0:
            rul_cycles = (current_soh_pred - eol_threshold) / abs(degradation_rate)
            rul_lower = (current_lower - eol_threshold) / abs(degradation_rate)
            rul_upper = (current_upper - eol_threshold) / abs(degradation_rate)
        else:
            rul_cycles = 1000  # Very long if not degrading
            rul_lower = 500
            rul_upper = 1500
    
    # Health score (0-100)
    health_score = current_soh_pred * 100
    
    # Degradation rate (% per year, assuming ~365 cycles per year)
    degradation_rate_pct_per_year = abs(degradation_rate) * 365 * 100
    
    # Remaining life (years)
    cycles_per_year = 365
    remaining_life_years = rul_cycles / cycles_per_year if rul_cycles > 0 else 0
    
    result = {
        'current_soh': {
            'point_estimate': float(current_soh_pred),
            'lower_95': float(current_lower),
            'upper_95': float(current_upper)
        },
        'health_score': float(health_score),  # 0-100
        'remaining_life': {
            'cycles': float(rul_cycles),
            'years': float(remaining_life_years),
            'lower_95_cycles': float(max(0, rul_lower)),
            'upper_95_cycles': float(rul_upper)
        },
        'degradation_rate': {
            'per_cycle': float(degradation_rate),
            'per_year_pct': float(degradation_rate_pct_per_year)
        },
        'forecasts': forecasts,
        'data_quality': quality,
        'current_cycle': int(current_cycle),
        'confidence': float(1 - quality['noise_level']),  # Higher quality = higher confidence
        'shap_explanation': shap_values_dict if shap_values_dict else None
    }
    
    return result


if __name__ == "__main__":
    # Test with sample data
    print("Testing predict_soh() function...")
    
    # Load test data
    test_data = pd.read_csv('./data/nasa_converted/B0007.csv')
    
    # Take first 100 cycles
    test_data_subset = test_data[test_data['cycle_index'] <= 100]
    
    try:
        predictions = predict_soh(
            test_data_subset,
            model_path='./models/nasa_soh_model.pkl',
            pipeline_path='./models/feature_pipeline_causal.pkl'
        )
        
        print("\n" + "="*80)
        print("PREDICTION RESULTS")
        print("="*80)
        print(f"\nCurrent SOH: {predictions['current_soh']['point_estimate']:.4f}")
        print(f"  95% Interval: [{predictions['current_soh']['lower_95']:.4f}, {predictions['current_soh']['upper_95']:.4f}]")
        print(f"\nHealth Score: {predictions['health_score']:.2f}/100")
        print(f"\nRemaining Life: {predictions['remaining_life']['years']:.2f} years")
        print(f"  (95% Interval: {predictions['remaining_life']['lower_95_cycles']:.0f} - {predictions['remaining_life']['upper_95_cycles']:.0f} cycles)")
        print(f"\nDegradation Rate: {predictions['degradation_rate']['per_year_pct']:.2f}%/year")
        print(f"\nData Quality Score: {predictions['data_quality']['quality_score']:.2f}")
        print(f"  Completeness: {predictions['data_quality']['completeness']:.2%}")
        print(f"  Noise Level: {predictions['data_quality']['noise_level']:.2%}")
        
        print(f"\nForecasts:")
        for key, forecast in predictions['forecasts'].items():
            cycles = key.replace('soh_at_', '').replace('_cycles', '')
            print(f"  {cycles} cycles ahead: {forecast['point_estimate']:.4f}")
            print(f"    (95% Interval: [{forecast['lower_95']:.4f}, {forecast['upper_95']:.4f}])")
        
        print("\n✓ Prediction function working correctly")
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        print("Note: Model must be trained first using production_pipeline.py")