life_expectancy_model.py

"""
Life Expectancy Prediction Model
Enhanced model evaluation with comprehensive visualization and metrics
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

# Set style for better plots
plt.style.use('default')
sns.set_palette("husl")

def load_and_clean_data():
    """
    Load dataset with enhanced cleaning and outlier removal
    """
    try:
        data = pd.read_csv("life_expectancy_data.csv")
        print("Loaded data from CSV file")
    except FileNotFoundError:
        print("Creating sample dataset")
        np.random.seed(42)
        n_countries = 150
        
        # Generate sample data
        gdp = np.random.lognormal(8, 1.5, n_countries)
        schooling = np.random.normal(8, 3, n_countries)
        schooling = np.clip(schooling, 2, 15)
        co2 = np.random.lognormal(1, 1, n_countries)
        population = np.random.lognormal(15, 2, n_countries)
        
        life_exp = (50 + 0.003 * gdp + 1.5 * schooling - 0.5 * co2 + 
                   np.random.normal(0, 3, n_countries))
        life_exp = np.clip(life_exp, 45, 85)
        
        countries = [f"Country_{i+1}" for i in range(n_countries)]
        
        data = pd.DataFrame({
            'Country': countries,
            'GDP_per_capita': gdp,
            'Schooling_years': schooling,
            'CO2_emissions': co2,
            'Population': population,
            'Life_expectancy': life_exp
        })
    
    print(f"Dataset shape: {data.shape}")
    print(f"Missing values before cleaning: {data.isnull().sum().sum()}")
    
    # Remove rows with missing values
    data_clean = data.dropna()
    
    # Enhanced outlier removal using IQR method
    def remove_outliers(df, column):
        Q1 = df[column].quantile(0.25)
        Q3 = df[column].quantile(0.75)
        IQR = Q3 - Q1
        lower = Q1 - 1.5 * IQR
        upper = Q3 + 1.5 * IQR
        return df[(df[column] >= lower) & (df[column] <= upper)]
    
    # Apply outlier removal to key variables
    original_size = len(data_clean)
    for col in ['GDP_per_capita', 'Life_expectancy']:
        if col in data_clean.columns:
            data_clean = remove_outliers(data_clean, col)
    
    print(f"Removed {original_size - len(data_clean)} outliers")
    print(f"Dataset shape after cleaning: {data_clean.shape}")
    return data_clean

def explore_data(data):
    """
    Comprehensive exploratory data analysis with correlation heatmap
    """
    print("\nEXPLORATORY DATA ANALYSIS")
    print("="*30)
    
    # Basic statistics
    numeric_cols = data.select_dtypes(include=[np.number]).columns
    print("\nSummary Statistics:")
    print(data[numeric_cols].describe().round(2))
    
    # Correlation matrix
    print("\nCorrelation Matrix:")
    corr_matrix = data[numeric_cols].corr()
    print(corr_matrix.round(3))
    
    # Enhanced visualization with correlation heatmap
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    fig.suptitle('Comprehensive Exploratory Data Analysis', fontsize=16, fontweight='bold')
    
    # 1. Correlation heatmap
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0, ax=axes[0,0])
    axes[0,0].set_title('Correlation Matrix')
    
    # 2. Life expectancy distribution
    axes[0,1].hist(data['Life_expectancy'], bins=20, alpha=0.7, color='skyblue')
    axes[0,1].set_title('Life Expectancy Distribution')
    axes[0,1].set_xlabel('Life Expectancy (years)')
    axes[0,1].set_ylabel('Frequency')
    
    # 3. GDP vs Life Expectancy
    axes[0,2].scatter(data['GDP_per_capita'], data['Life_expectancy'], alpha=0.6)
    axes[0,2].set_title('GDP vs Life Expectancy')
    axes[0,2].set_xlabel('GDP per Capita')
    axes[0,2].set_ylabel('Life Expectancy')
    
    # 4. Schooling vs Life Expectancy
    axes[1,0].scatter(data['Schooling_years'], data['Life_expectancy'], alpha=0.6, color='orange')
    axes[1,0].set_title('Education vs Life Expectancy')
    axes[1,0].set_xlabel('Years of Schooling')
    axes[1,0].set_ylabel('Life Expectancy')
    
    # 5. CO2 vs Life Expectancy
    axes[1,1].scatter(data['CO2_emissions'], data['Life_expectancy'], alpha=0.6, color='red')
    axes[1,1].set_title('CO2 vs Life Expectancy')
    axes[1,1].set_xlabel('CO2 Emissions')
    axes[1,1].set_ylabel('Life Expectancy')
    
    # 6. Population distribution
    axes[1,2].hist(np.log10(data['Population']), bins=20, alpha=0.7, color='green')
    axes[1,2].set_title('Population Distribution (Log Scale)')
    axes[1,2].set_xlabel('Log10(Population)')
    axes[1,2].set_ylabel('Frequency')
    
    plt.tight_layout()
    plt.show()

def build_regression_model(data):
    """
    Build regression model with comprehensive evaluation
    """
    print("\nBUILDING REGRESSION MODEL")
    print("="*25)
    
    # Define features and target
    feature_columns = ['GDP_per_capita', 'Schooling_years', 'CO2_emissions', 'Population']
    X = data[feature_columns]
    y = data['Life_expectancy']
    
    print(f"Features: {feature_columns}")
    print(f"Dataset size: {len(X)} samples, {len(feature_columns)} features")
    
    # Split the data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # Feature scaling
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    # Train the model
    model = LinearRegression()
    model.fit(X_train_scaled, y_train)
    
    # Make predictions
    y_pred_train = model.predict(X_train_scaled)
    y_pred_test = model.predict(X_test_scaled)
    
    # Calculate comprehensive metrics
    train_r2 = r2_score(y_train, y_pred_train)
    test_r2 = r2_score(y_test, y_pred_test)
    train_mse = mean_squared_error(y_train, y_pred_train)
    test_mse = mean_squared_error(y_test, y_pred_test)
    test_mae = mean_absolute_error(y_test, y_pred_test)
    test_rmse = np.sqrt(test_mse)
    
    print(f"\nCOMPREHENSIVE MODEL PERFORMANCE:")
    print(f"Training R² Score: {train_r2:.4f}")
    print(f"Testing R² Score: {test_r2:.4f}")
    print(f"Training MSE: {train_mse:.2f}")
    print(f"Testing MSE: {test_mse:.2f}")
    print(f"Testing MAE: {test_mae:.2f}")
    print(f"Testing RMSE: {test_rmse:.2f}")
    
    # Model interpretation
    print(f"\nMODEL COEFFICIENTS (Scaled Features):")
    feature_importance = list(zip(feature_columns, model.coef_))
    feature_importance.sort(key=lambda x: abs(x[1]), reverse=True)
    
    for feature, coef in feature_importance:
        print(f"   {feature}: {coef:.4f}")
    print(f"   Intercept: {model.intercept_:.4f}")
    
    # Advanced visualization
    create_comprehensive_visualizations(y_test, y_pred_test, model, feature_columns)
    
    return model, scaler, feature_columns

def create_comprehensive_visualizations(y_true, y_pred, model, feature_names):
    """
    Create comprehensive model evaluation visualizations
    """
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('Comprehensive Model Evaluation Results', fontsize=16, fontweight='bold')
    
    # 1. Actual vs Predicted scatter plot
    axes[0,0].scatter(y_true, y_pred, alpha=0.6)
    axes[0,0].plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'r--', lw=2)
    axes[0,0].set_xlabel('Actual Life Expectancy')
    axes[0,0].set_ylabel('Predicted Life Expectancy')
    axes[0,0].set_title('Actual vs Predicted Values')
    
    # Add metrics to the plot
    r2 = r2_score(y_true, y_pred)
    mae = mean_absolute_error(y_true, y_pred)
    axes[0,0].text(0.05, 0.95, f'R² = {r2:.3f}\nMAE = {mae:.2f}', 
                   transform=axes[0,0].transAxes, 
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    # 2. Residuals plot
    residuals = y_true - y_pred
    axes[0,1].scatter(y_pred, residuals, alpha=0.6)
    axes[0,1].axhline(y=0, color='r', linestyle='--')
    axes[0,1].set_xlabel('Predicted Life Expectancy')
    axes[0,1].set_ylabel('Residuals')
    axes[0,1].set_title('Residual Plot')
    
    # 3. Feature importance (coefficients)
    coefficients = model.coef_
    colors = ['red' if x < 0 else 'blue' for x in coefficients]
    bars = axes[1,0].barh(feature_names, coefficients, color=colors)
    axes[1,0].set_xlabel('Coefficient Value')
    axes[1,0].set_title('Feature Coefficients (Scaled)')
    axes[1,0].axvline(x=0, color='black', linestyle='-', alpha=0.3)
    
    # 4. Residuals histogram with normal curve
    axes[1,1].hist(residuals, bins=15, alpha=0.7, color='lightcoral', density=True)
    axes[1,1].set_xlabel('Residuals')
    axes[1,1].set_ylabel('Density')
    axes[1,1].set_title('Distribution of Residuals')
    axes[1,1].axvline(x=0, color='r', linestyle='--')
    
    # Add normal curve overlay
    x_norm = np.linspace(residuals.min(), residuals.max(), 100)
    y_norm = ((1/np.sqrt(2*np.pi*residuals.var())) * 
              np.exp(-0.5*((x_norm-residuals.mean())/residuals.std())**2))
    axes[1,1].plot(x_norm, y_norm, 'b-', linewidth=2, label='Normal Distribution')
    axes[1,1].legend()
    
    plt.tight_layout()
    plt.show()

def make_predictions(model, scaler, feature_columns):
    """
    Make predictions with detailed scenario analysis
    """
    print("\nDETAILED PREDICTION ANALYSIS")
    print("="*28)
    
    # Enhanced scenarios with more realistic data
    scenarios = {
        "High-income Country (e.g., Switzerland)": [80000, 13, 4.5, 8500000],
        "Upper-middle-income (e.g., Brazil)": [15000, 8, 2.2, 210000000],
        "Lower-middle-income (e.g., India)": [7000, 6, 1.8, 1400000000],
        "Low-income Country (e.g., Chad)": [1500, 2, 0.1, 15000000]
    }
    
    print("Predicting life expectancy for different economic scenarios:\n")
    
    for scenario_name, values in scenarios.items():
        scenario_data = np.array(values).reshape(1, -1)
        scenario_scaled = scaler.transform(scenario_data)
        prediction = model.predict(scenario_scaled)[0]
        
        print(f"{scenario_name}:")
        for i, feature in enumerate(feature_columns):
            print(f"   {feature}: {values[i]:,}")
        print(f"   -> Predicted Life Expectancy: {prediction:.1f} years")
        print()

def main():
    """
    Main function with comprehensive analysis
    """
    print("LIFE EXPECTANCY PREDICTION MODEL - COMPREHENSIVE ANALYSIS")
    print("="*60)
    
    data = load_and_clean_data()
    explore_data(data)
    model, scaler, feature_columns = build_regression_model(data)
    make_predictions(model, scaler, feature_columns)
    
    print("="*60)
    print("COMPREHENSIVE ANALYSIS COMPLETE!")
    print("Check the enhanced plots above for detailed insights")
    print("Model ready for production use")
    
    return data, model

if __name__ == "__main__":
    data, model = main()