#!/usr/bin/env python3
"""
Model Testing and Evaluation Script

This script loads the trained models and performs comprehensive testing:
1. Load saved models
2. Test on new data
3. Generate predictions
4. Visualize results
5. Cross-validation analysis
"""

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
import joblib
import torch
import torch.nn as nn
import xgboost as xgb
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.model_selection import cross_val_score
import warnings
import os

warnings.filterwarnings('ignore')

class DataPreprocessor:
    """Data preprocessing and feature engineering class."""
    
    def __init__(self):
        self.label_encoders = {}
        self.scaler = None
        self.imputer = None
        self.feature_names = None
        
    def prepare_features(self, df):
        """Prepare features for machine learning."""
        print("Preparing features...")
        
        # Create a copy to avoid modifying original data
        data = df.copy()
        
        # Remove records with zero or negative yield (invalid data)
        data = data[data['Yield'] > 0].copy()
        
        # Feature engineering
        data['Area_Production_Ratio'] = data['Area'] / (data['Production'] + 1e-6)
        data['Yield_Area_Interaction'] = data['Yield'] * data['Area']
        data['Production_Per_Area'] = data['Production'] / (data['Area'] + 1e-6)
        
        # Create season dummies
        season_dummies = pd.get_dummies(data['Season'], prefix='Season')
        data = pd.concat([data, season_dummies], axis=1)
        
        # Handle categorical variables
        categorical_cols = ['State', 'District', 'Crop']
        
        for col in categorical_cols:
            if col in data.columns:
                # Use label encoding for high cardinality features
                if col not in self.label_encoders:
                    from sklearn.preprocessing import LabelEncoder
                    self.label_encoders[col] = LabelEncoder()
                    data[f'{col}_encoded'] = self.label_encoders[col].fit_transform(data[col].astype(str))
                else:
                    # Handle unseen categories
                    unique_values = set(data[col].astype(str))
                    known_values = set(self.label_encoders[col].classes_)
                    new_values = unique_values - known_values
                    
                    if new_values:
                        # Add new categories to the encoder
                        all_values = list(known_values) + list(new_values)
                        self.label_encoders[col].classes_ = np.array(all_values)
                    
                    data[f'{col}_encoded'] = self.label_encoders[col].transform(data[col].astype(str))
        
        # Select features for modeling
        feature_cols = ['Crop_Year', 'Area', 'Production', 'Annual_Rainfall', 
                       'Fertilizer', 'Pesticide', 'State_encoded', 'Crop_encoded',
                       'Area_Production_Ratio', 'Yield_Area_Interaction', 
                       'Production_Per_Area'] + list(season_dummies.columns)
        
        # Add District_encoded if available
        if 'District_encoded' in data.columns:
            feature_cols.append('District_encoded')
        
        # Select only available columns
        available_cols = [col for col in feature_cols if col in data.columns]
        X = data[available_cols].copy()
        y = data['Yield'].copy()
        
        print(f"Selected features: {available_cols}")
        print(f"Dataset shape after preprocessing: {X.shape}")
        
        return X, y, data
    
    def transform(self, X):
        """Transform new data using fitted preprocessors."""
        # Handle missing values
        X_imputed = pd.DataFrame(
            self.imputer.transform(X), 
            columns=X.columns, 
            index=X.index
        )
        
        # Scale features
        X_scaled = pd.DataFrame(
            self.scaler.transform(X_imputed), 
            columns=X.columns, 
            index=X.index
        )
        
        return X_scaled

class PyTorchYieldPredictor(nn.Module):
    """PyTorch Neural Network for yield prediction (same as training script)."""
    
    def __init__(self, input_dim, hidden_dims=[256, 128, 64], dropout_rate=0.3):
        super(PyTorchYieldPredictor, self).__init__()
        
        layers = []
        prev_dim = input_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout_rate)
            ])
            prev_dim = hidden_dim
        
        # Output layer
        layers.append(nn.Linear(prev_dim, 1))
        
        self.model = nn.Sequential(*layers)
        
    def forward(self, x):
        return self.model(x)

class ModelTester:
    """Class for testing and evaluating trained models."""
    
    def __init__(self, models_dir='trained_models'):
        self.models_dir = models_dir
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Using device: {self.device}")
        
        self.models = {}
        self.preprocessor = None
        self.results = {}
        
    def load_models(self):
        """Load all trained models and preprocessor."""
        print("Loading trained models...")
        
        try:
            # Load preprocessor
            preprocessor_path = os.path.join(self.models_dir, 'preprocessor.pkl')
            if os.path.exists(preprocessor_path):
                self.preprocessor = joblib.load(preprocessor_path)
                print("✅ Preprocessor loaded")
            else:
                raise FileNotFoundError("Preprocessor not found")
            
            # Load Random Forest
            rf_path = os.path.join(self.models_dir, 'random_forest_model.pkl')
            if os.path.exists(rf_path):
                self.models['RandomForest'] = joblib.load(rf_path)
                print("✅ Random Forest model loaded")
            
            # Load XGBoost
            xgb_path = os.path.join(self.models_dir, 'xgboost_model.json')
            if os.path.exists(xgb_path):
                xgb_model = xgb.XGBRegressor()
                xgb_model.load_model(xgb_path)
                self.models['XGBoost'] = xgb_model
                print("✅ XGBoost model loaded")
            
            # Load PyTorch model
            pytorch_path = os.path.join(self.models_dir, 'pytorch_model.pth')
            if os.path.exists(pytorch_path):
                # We need to know the input dimension - get it from preprocessor
                # This is a bit tricky - we'll determine it from the data
                print("✅ PyTorch model path found (will load after determining input size)")
                
        except Exception as e:
            print(f"Error loading models: {e}")
            raise
    
    def load_pytorch_model(self, input_dim):
        """Load PyTorch model with known input dimension."""
        pytorch_path = os.path.join(self.models_dir, 'pytorch_model.pth')
        if os.path.exists(pytorch_path):
            pytorch_model = PyTorchYieldPredictor(input_dim).to(self.device)
            pytorch_model.load_state_dict(torch.load(pytorch_path, map_location=self.device))
            pytorch_model.eval()
            self.models['PyTorch'] = pytorch_model
            print("✅ PyTorch model loaded")
    
    def prepare_test_data(self, data_file='combined_crop_data.csv', sample_size=1000):
        """Prepare test data for evaluation."""
        print(f"Preparing test data from {data_file}...")
        
        # Load data
        df = pd.read_csv(data_file)
        
        # Sample data for testing if too large
        if len(df) > sample_size:
            df = df.sample(n=sample_size, random_state=42)
            print(f"Sampled {sample_size} records for testing")
        
        # Prepare features using the same preprocessor
        X, y, processed_data = self.preprocessor.prepare_features(df)
        
        # Transform using fitted preprocessor
        X_processed = self.preprocessor.transform(X)
        
        print(f"Test data shape: {X_processed.shape}")
        
        # Now we can load PyTorch model
        input_dim = X_processed.shape[1]
        self.load_pytorch_model(input_dim)
        
        return X_processed, y, processed_data
    
    def test_models(self, X_test, y_test):
        """Test all loaded models and calculate metrics."""
        print("\\n" + "="*50)
        print("TESTING MODELS")
        print("="*50)
        
        for model_name, model in self.models.items():
            print(f"\\nTesting {model_name}...")
            
            try:
                if model_name == 'PyTorch':
                    # PyTorch model prediction
                    X_tensor = torch.FloatTensor(X_test.values).to(self.device)
                    with torch.no_grad():
                        predictions = model(X_tensor).cpu().numpy().flatten()
                else:
                    # Sklearn/XGBoost prediction
                    predictions = model.predict(X_test)
                
                # Calculate metrics
                mse = mean_squared_error(y_test, predictions)
                rmse = np.sqrt(mse)
                mae = mean_absolute_error(y_test, predictions)
                r2 = r2_score(y_test, predictions)
                
                self.results[model_name] = {
                    'predictions': predictions,
                    'mse': mse,
                    'rmse': rmse,
                    'mae': mae,
                    'r2': r2
                }
                
                print(f"  RMSE: {rmse:.4f}")
                print(f"  MAE: {mae:.4f}")
                print(f"  R²: {r2:.4f}")
                
            except Exception as e:
                print(f"  ❌ Error testing {model_name}: {e}")
    
    def visualize_results(self, y_test):
        """Create visualizations of model performance."""
        print("\\nCreating visualizations...")
        
        # Create subplots for each model
        n_models = len(self.results)
        fig, axes = plt.subplots(2, n_models, figsize=(5*n_models, 10))
        
        if n_models == 1:
            axes = axes.reshape(-1, 1)
        
        for i, (model_name, results) in enumerate(self.results.items()):
            predictions = results['predictions']
            
            # Actual vs Predicted scatter plot
            axes[0, i].scatter(y_test, predictions, alpha=0.6)
            axes[0, i].plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
            axes[0, i].set_xlabel('Actual Yield')
            axes[0, i].set_ylabel('Predicted Yield')
            axes[0, i].set_title(f'{model_name} - Actual vs Predicted\\nR² = {results["r2"]:.4f}')
            
            # Residuals plot
            residuals = y_test - predictions
            axes[1, i].scatter(predictions, residuals, alpha=0.6)
            axes[1, i].axhline(y=0, color='r', linestyle='--')
            axes[1, i].set_xlabel('Predicted Yield')
            axes[1, i].set_ylabel('Residuals')
            axes[1, i].set_title(f'{model_name} - Residuals Plot')
        
        plt.tight_layout()
        plt.savefig('model_test_results.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print("✅ Visualization saved as model_test_results.png")
        
        # Create comparison chart
        self.plot_model_comparison()
    
    def plot_model_comparison(self):
        """Plot model comparison metrics."""
        comparison_data = []
        
        for model_name, results in self.results.items():
            comparison_data.append({
                'Model': model_name,
                'RMSE': results['rmse'],
                'MAE': results['mae'],
                'R²': results['r2']
            })
        
        comparison_df = pd.DataFrame(comparison_data)
        
        # Create comparison plots
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # RMSE comparison
        axes[0].bar(comparison_df['Model'], comparison_df['RMSE'], alpha=0.7, color='blue')
        axes[0].set_title('RMSE Comparison')
        axes[0].set_ylabel('RMSE')
        axes[0].tick_params(axis='x', rotation=45)
        
        # MAE comparison
        axes[1].bar(comparison_df['Model'], comparison_df['MAE'], alpha=0.7, color='orange')
        axes[1].set_title('MAE Comparison')
        axes[1].set_ylabel('MAE')
        axes[1].tick_params(axis='x', rotation=45)
        
        # R² comparison
        axes[2].bar(comparison_df['Model'], comparison_df['R²'], alpha=0.7, color='green')
        axes[2].set_title('R² Comparison')
        axes[2].set_ylabel('R² Score')
        axes[2].tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        plt.savefig('model_performance_comparison.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print("✅ Comparison chart saved as model_performance_comparison.png")
        
        # Print comparison table
        print("\\n" + "="*50)
        print("MODEL PERFORMANCE COMPARISON")
        print("="*50)
        print(comparison_df.to_string(index=False, float_format='%.4f'))
    
    def cross_validate_models(self, X, y, cv=5):
        """Perform cross-validation on models that support it."""
        print("\\n" + "="*50)
        print("CROSS-VALIDATION RESULTS")
        print("="*50)
        
        cv_results = {}
        
        for model_name, model in self.models.items():
            if model_name != 'PyTorch':  # Skip PyTorch for CV (more complex to implement)
                try:
                    print(f"\\nCross-validating {model_name}...")
                    cv_scores = cross_val_score(model, X, y, cv=cv, scoring='neg_mean_squared_error')
                    cv_rmse = np.sqrt(-cv_scores)
                    
                    cv_results[model_name] = {
                        'cv_rmse_mean': cv_rmse.mean(),
                        'cv_rmse_std': cv_rmse.std(),
                        'cv_scores': cv_rmse
                    }
                    
                    print(f"  CV RMSE: {cv_rmse.mean():.4f} ± {cv_rmse.std():.4f}")
                    
                except Exception as e:
                    print(f"  ❌ Error in cross-validation for {model_name}: {e}")
        
        return cv_results
    
    def generate_predictions_for_new_data(self, new_data_file=None):
        """Generate predictions for new data."""
        if new_data_file is None:
            print("\\nNo new data file provided for prediction.")
            return
        
        print(f"\\nGenerating predictions for {new_data_file}...")
        
        try:
            # Load new data
            new_df = pd.read_csv(new_data_file)
            
            # Prepare features
            X_new, _, _ = self.preprocessor.prepare_features(new_df)
            X_new_processed = self.preprocessor.transform(X_new)
            
            predictions_df = new_df.copy()
            
            # Generate predictions from each model
            for model_name, model in self.models.items():
                if model_name == 'PyTorch':
                    X_tensor = torch.FloatTensor(X_new_processed.values).to(self.device)
                    with torch.no_grad():
                        preds = model(X_tensor).cpu().numpy().flatten()
                else:
                    preds = model.predict(X_new_processed)
                
                predictions_df[f'Predicted_Yield_{model_name}'] = preds
            
            # Save predictions
            output_file = 'new_data_predictions.csv'
            predictions_df.to_csv(output_file, index=False)
            print(f"✅ Predictions saved to {output_file}")
            
            return predictions_df
            
        except Exception as e:
            print(f"❌ Error generating predictions: {e}")
    
    def run_comprehensive_test(self, data_file='combined_crop_data.csv', new_data_file=None):
        """Run comprehensive testing pipeline."""
        print("🧪 Starting Comprehensive Model Testing...")
        
        try:
            # Load models
            self.load_models()
            
            # Prepare test data
            X_test, y_test, processed_data = self.prepare_test_data(data_file)
            
            # Test models
            self.test_models(X_test, y_test)
            
            # Create visualizations
            self.visualize_results(y_test)
            
            # Cross-validation
            cv_results = self.cross_validate_models(X_test, y_test)
            
            # Generate predictions for new data if provided
            if new_data_file:
                self.generate_predictions_for_new_data(new_data_file)
            
            print("\\n🎉 Comprehensive testing completed successfully!")
            print("📊 Check the generated plots and results files.")
            
            return self.results, cv_results
            
        except Exception as e:
            print(f"❌ Error in testing pipeline: {e}")
            raise

def main():
    """Main function to run model testing."""
    tester = ModelTester()
    
    # Check if trained models exist
    if not os.path.exists('trained_models'):
        print("❌ No trained models found. Please run the training script first.")
        return
    
    # Run comprehensive testing
    results, cv_results = tester.run_comprehensive_test()
    
    return tester, results, cv_results

if __name__ == "__main__":
    tester, results, cv_results = main()