#!/usr/bin/env python3
"""
Quick comparison between original and cleaned dataset performance
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import r2_score, mean_squared_error
import warnings
warnings.filterwarnings('ignore')

def test_dataset_performance(df, dataset_name="Dataset"):
    """Test model performance on a given dataset"""
    print(f"\n{'='*50}")
    print(f"TESTING {dataset_name.upper()}")
    print(f"{'='*50}")
    
    # Basic stats
    print(f"Shape: {df.shape}")
    print(f"Missing values: {df.isnull().sum().sum():,}")
    
    # Prepare data - find yield column
    yield_cols = [col for col in df.columns if 'yield' in col.lower()]
    if not yield_cols:
        print("❌ No yield column found!")
        return None
    
    target_col = yield_cols[0]
    print(f"Target column: {target_col}")
    
    # Remove zero yields for modeling
    df_model = df[df[target_col] > 0].copy()
    print(f"Non-zero yield records: {len(df_model):,}")
    
    if len(df_model) < 100:
        print("❌ Insufficient non-zero yield data!")
        return None
    
    # Select features
    numerical_cols = df_model.select_dtypes(include=[np.number]).columns.tolist()
    numerical_cols = [col for col in numerical_cols if col not in [target_col, 'Crop_Year']]
    
    # Add categorical features with encoding
    categorical_cols = ['Crop', 'Season', 'State']
    feature_cols = []
    
    # Add numerical features
    for col in numerical_cols:
        if col in df_model.columns and df_model[col].notna().sum() > len(df_model) * 0.5:
            feature_cols.append(col)
    
    X = df_model[feature_cols].copy()
    
    # Handle missing values by filling with median
    for col in X.columns:
        X[col] = X[col].fillna(X[col].median())
    
    # Add encoded categorical features
    le_dict = {}
    for col in categorical_cols:
        if col in df_model.columns:
            le = LabelEncoder()
            X[col + '_encoded'] = le.fit_transform(df_model[col].fillna('Unknown'))
            le_dict[col] = le
    
    y = df_model[target_col]
    
    print(f"Features used: {len(X.columns)}")
    print(f"Target range: {y.min():.2f} to {y.max():.2f}")
    
    # Train model
    try:
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        model = RandomForestRegressor(n_estimators=50, random_state=42)  # Smaller for quick test
        model.fit(X_train, y_train)
        
        y_pred = model.predict(X_test)
        
        r2 = r2_score(y_test, y_pred)
        rmse = np.sqrt(mean_squared_error(y_test, y_pred))
        
        print(f"✅ R² Score: {r2:.4f}")
        print(f"✅ RMSE: {rmse:.2f}")
        
        return {
            'dataset': dataset_name,
            'records_used': len(df_model),
            'features': len(X.columns),
            'r2_score': r2,
            'rmse': rmse,
            'missing_values': df.isnull().sum().sum()
        }
        
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        return None

def main():
    print("🔍 COMPARING ORIGINAL VS CLEANED DATASET PERFORMANCE")
    print("="*80)
    
    # Load datasets
    try:
        original_df = pd.read_csv("/home/aiavid/Yeild_pred_SIH/data/combined_crop_data.csv")
        cleaned_df = pd.read_csv("/home/aiavid/Yeild_pred_SIH/data/combined_crop_data_cleaned.csv")
    except Exception as e:
        print(f"❌ Error loading datasets: {str(e)}")
        return
    
    # Test both datasets
    original_results = test_dataset_performance(original_df, "ORIGINAL DATASET")
    cleaned_results = test_dataset_performance(cleaned_df, "CLEANED DATASET")
    
    # Compare results
    if original_results and cleaned_results:
        print(f"\n{'='*80}")
        print("COMPARISON RESULTS")
        print(f"{'='*80}")
        
        print("\n| Metric | Original | Cleaned | Improvement |")
        print("|--------|----------|---------|-------------|")
        
        # Missing values
        orig_missing = original_results['missing_values']
        clean_missing = cleaned_results['missing_values']
        missing_improvement = ((orig_missing - clean_missing) / orig_missing * 100) if orig_missing > 0 else 0
        print(f"| Missing Values | {orig_missing:,} | {clean_missing:,} | {missing_improvement:.1f}% reduction |")
        
        # Records used
        orig_records = original_results['records_used']
        clean_records = cleaned_results['records_used']
        record_improvement = ((clean_records - orig_records) / orig_records * 100) if orig_records > 0 else 0
        print(f"| Usable Records | {orig_records:,} | {clean_records:,} | {record_improvement:+.1f}% |")
        
        # R² Score
        orig_r2 = original_results['r2_score']
        clean_r2 = cleaned_results['r2_score']
        r2_improvement = clean_r2 - orig_r2
        print(f"| R² Score | {orig_r2:.4f} | {clean_r2:.4f} | {r2_improvement:+.4f} |")
        
        # RMSE
        orig_rmse = original_results['rmse']
        clean_rmse = cleaned_results['rmse']
        rmse_improvement = ((orig_rmse - clean_rmse) / orig_rmse * 100) if orig_rmse > 0 else 0
        print(f"| RMSE | {orig_rmse:.2f} | {clean_rmse:.2f} | {rmse_improvement:.1f}% better |")
        
        print(f"\n🎯 **OVERALL ASSESSMENT:**")
        print(f"   📈 Model Accuracy Improved by {r2_improvement:+.4f} R² points")
        print(f"   📉 Prediction Error Reduced by {rmse_improvement:.1f}%")
        print(f"   🧹 Data Quality Improved: {missing_improvement:.1f}% fewer missing values")
        
        if clean_r2 > 0.9:
            print(f"   🏆 **EXCELLENT** model performance achieved!")
        elif clean_r2 > 0.8:
            print(f"   🥇 **VERY GOOD** model performance achieved!")
        elif clean_r2 > 0.7:
            print(f"   🥈 **GOOD** model performance achieved!")
        else:
            print(f"   📊 Model performance: **ACCEPTABLE**")

if __name__ == "__main__":
    main()