"""
============================================================================
MODEL SELECTION METHODS IN PYTHON
Ridge, LASSO, Best Subset, Forward/Backward Selection
============================================================================
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import Ridge, Lasso, ElasticNet, LinearRegression
from sklearn.model_selection import cross_val_score, GridSearchCV, KFold
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score
from itertools import combinations
import warnings
warnings.filterwarnings('ignore')

# Set random seed for reproducibility
np.random.seed(42)

# Set plotting style
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = (10, 6)

# ============================================================================
# GENERATE SAMPLE ECONOMIC DATA
# ============================================================================

print("=" * 70)
print("GENERATING SAMPLE ECONOMIC DATA")
print("=" * 70)

n = 200  # Number of observations
p = 10   # Number of predictors

# Simulate economic indicators
X = np.random.randn(n, p)
feature_names = ['GDP_Growth', 'Inflation', 'Unemployment', 'Interest_Rate',
                 'Investment', 'Consumption', 'Exports', 'Imports',
                 'Money_Supply', 'Exchange_Rate']

# Generate response variable (e.g., Stock Market Returns)
# Only some predictors are truly relevant
true_coefficients = np.array([2.5, -1.8, -2.2, 0, 1.5, 0, 0, 0, 0, 0.8])
Y = X @ true_coefficients + np.random.randn(n) * 2

# Create DataFrame
df = pd.DataFrame(X, columns=feature_names)
df['Y'] = Y

# Split into train and test sets
train_size = int(0.7 * n)
train_indices = np.random.choice(n, train_size, replace=False)
test_indices = np.setdiff1d(np.arange(n), train_indices)

X_train, X_test = X[train_indices], X[test_indices]
Y_train, Y_test = Y[train_indices], Y[test_indices]

print(f"\nData prepared: {len(train_indices)} training observations, "
      f"{len(test_indices)} test observations\n")

# ============================================================================
# 1. BEST SUBSET SELECTION
# ============================================================================

print("=" * 70)
print("1. BEST SUBSET SELECTION")
print("=" * 70)

def best_subset_selection(X, y, max_features=None):
    """
    Perform best subset selection using exhaustive search.
    Returns best model for each subset size.
    """
    n_features = X.shape[1]
    if max_features is None:
        max_features = n_features
    
    results = []
    
    for k in range(1, max_features + 1):
        best_score = -np.inf
        best_subset = None
        
        # Try all combinations of k features
        for subset in combinations(range(n_features), k):
            X_subset = X[:, list(subset)]
            model = LinearRegression()
            
            # Use cross-validation to evaluate
            scores = cross_val_score(model, X_subset, y, 
                                    cv=5, scoring='r2')
            score = scores.mean()
            
            if score > best_score:
                best_score = score
                best_subset = subset
        
        results.append({
            'n_features': k,
            'features': best_subset,
            'cv_r2': best_score
        })
    
    return pd.DataFrame(results)

# Perform best subset selection
print("\nPerforming best subset selection (this may take a moment)...")
best_subset_results = best_subset_selection(X_train, Y_train, max_features=8)

print("\nBest subset for each model size:")
print(best_subset_results)

# Select best model
best_k = best_subset_results.loc[best_subset_results['cv_r2'].idxmax()]
print(f"\nBest model has {int(best_k['n_features'])} features")
print(f"Features: {[feature_names[i] for i in best_k['features']]}")

# ============================================================================
# 2. FORWARD STEPWISE SELECTION
# ============================================================================

print("\n" + "=" * 70)
print("2. FORWARD STEPWISE SELECTION")
print("=" * 70)

def forward_selection(X, y, feature_names):
    """
    Perform forward stepwise selection.
    """
    n_features = X.shape[1]
    selected_features = []
    remaining_features = list(range(n_features))
    
    results = []
    
    for step in range(n_features):
        best_score = -np.inf
        best_feature = None
        
        for feature in remaining_features:
            # Try adding this feature
            test_features = selected_features + [feature]
            X_subset = X[:, test_features]
            
            model = LinearRegression()
            scores = cross_val_score(model, X_subset, y, cv=5, scoring='r2')
            score = scores.mean()
            
            if score > best_score:
                best_score = score
                best_feature = feature
        
        # Add best feature
        selected_features.append(best_feature)
        remaining_features.remove(best_feature)
        
        results.append({
            'step': step + 1,
            'feature_added': feature_names[best_feature],
            'cv_r2': best_score,
            'features': [feature_names[i] for i in selected_features]
        })
        
        print(f"Step {step + 1}: Added {feature_names[best_feature]}, "
              f"CV R² = {best_score:.4f}")
    
    return pd.DataFrame(results)

# Perform forward selection
print("\nPerforming forward stepwise selection...")
forward_results = forward_selection(X_train, Y_train, feature_names)

print("\n" + forward_results.to_string(index=False))

# Select best model
best_forward = forward_results.loc[forward_results['cv_r2'].idxmax()]
print(f"\nBest model at step {int(best_forward['step'])}")
print(f"Selected features: {best_forward['features']}")

# ============================================================================
# 3. BACKWARD STEPWISE SELECTION
# ============================================================================

print("\n" + "=" * 70)
print("3. BACKWARD STEPWISE SELECTION")
print("=" * 70)

def backward_selection(X, y, feature_names):
    """
    Perform backward stepwise selection.
    """
    n_features = X.shape[1]
    selected_features = list(range(n_features))
    
    results = []
    
    # Start with all features
    model = LinearRegression()
    scores = cross_val_score(model, X, y, cv=5, scoring='r2')
    results.append({
        'step': 0,
        'feature_removed': 'None',
        'cv_r2': scores.mean(),
        'features': [feature_names[i] for i in selected_features]
    })
    
    for step in range(n_features - 1):
        best_score = -np.inf
        worst_feature = None
        
        for feature in selected_features:
            # Try removing this feature
            test_features = [f for f in selected_features if f != feature]
            if len(test_features) == 0:
                continue
                
            X_subset = X[:, test_features]
            
            model = LinearRegression()
            scores = cross_val_score(model, X_subset, y, cv=5, scoring='r2')
            score = scores.mean()
            
            if score > best_score:
                best_score = score
                worst_feature = feature
        
        # Remove worst feature
        selected_features.remove(worst_feature)
        
        results.append({
            'step': step + 1,
            'feature_removed': feature_names[worst_feature],
            'cv_r2': best_score,
            'features': [feature_names[i] for i in selected_features]
        })
        
        print(f"Step {step + 1}: Removed {feature_names[worst_feature]}, "
              f"CV R² = {best_score:.4f}")
    
    return pd.DataFrame(results)

# Perform backward selection
print("\nPerforming backward stepwise selection...")
backward_results = backward_selection(X_train, Y_train, feature_names)

# Select best model
best_backward = backward_results.loc[backward_results['cv_r2'].idxmax()]
print(f"\nBest model at step {int(best_backward['step'])}")
print(f"Selected features: {best_backward['features']}")

# ============================================================================
# 4. RIDGE REGRESSION
# ============================================================================

print("\n" + "=" * 70)
print("4. RIDGE REGRESSION (L2 Regularization)")
print("=" * 70)

# Standardize features (important for regularization)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# Ridge regression with cross-validation
alphas = np.logspace(-2, 4, 100)
ridge_cv = GridSearchCV(Ridge(), 
                        param_grid={'alpha': alphas},
                        cv=10, 
                        scoring='neg_mean_squared_error')
ridge_cv.fit(X_train_scaled, Y_train)

print(f"\nOptimal alpha: {ridge_cv.best_params_['alpha']:.4f}")

# Get best model
ridge_best = ridge_cv.best_estimator_

# Coefficients
ridge_coef = pd.DataFrame({
    'Feature': feature_names,
    'Coefficient': ridge_best.coef_
}).sort_values('Coefficient', key=abs, ascending=False)

print("\nRidge coefficients at optimal alpha:")
print(ridge_coef.to_string(index=False))

# Predictions on test set
ridge_pred = ridge_best.predict(X_test_scaled)
ridge_mse = mean_squared_error(Y_test, ridge_pred)
ridge_r2 = r2_score(Y_test, ridge_pred)

print(f"\nTest MSE: {ridge_mse:.4f}")
print(f"Test R²: {ridge_r2:.4f}")

# Plot coefficient paths
plt.figure(figsize=(10, 6))
ridge_models = []
coef_paths = []

for alpha in alphas:
    ridge_model = Ridge(alpha=alpha)
    ridge_model.fit(X_train_scaled, Y_train)
    coef_paths.append(ridge_model.coef_)

coef_paths = np.array(coef_paths)

for i in range(p):
    plt.plot(np.log10(alphas), coef_paths[:, i], label=feature_names[i])

plt.axvline(np.log10(ridge_cv.best_params_['alpha']), 
            color='red', linestyle='--', label='Optimal α')
plt.xlabel('log10(α)')
plt.ylabel('Coefficient Value')
plt.title('Ridge Regression: Coefficient Paths')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/ridge_coefficient_paths.png', dpi=300, bbox_inches='tight')
print("\nCoefficient paths plot saved!")

# Plot CV error vs alpha
plt.figure(figsize=(10, 6))
cv_scores = -ridge_cv.cv_results_['mean_test_score']
plt.plot(np.log10(alphas), cv_scores)
plt.axvline(np.log10(ridge_cv.best_params_['alpha']), 
            color='red', linestyle='--', label='Optimal α')
plt.xlabel('log10(α)')
plt.ylabel('Cross-Validation MSE')
plt.title('Ridge Regression: Cross-Validation Error')
plt.legend()
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/ridge_cv_error.png', dpi=300, bbox_inches='tight')
print("CV error plot saved!")

# ============================================================================
# 5. LASSO REGRESSION
# ============================================================================

print("\n" + "=" * 70)
print("5. LASSO REGRESSION (L1 Regularization)")
print("=" * 70)

# LASSO with cross-validation
lasso_cv = GridSearchCV(Lasso(max_iter=10000), 
                        param_grid={'alpha': alphas},
                        cv=10, 
                        scoring='neg_mean_squared_error')
lasso_cv.fit(X_train_scaled, Y_train)

print(f"\nOptimal alpha: {lasso_cv.best_params_['alpha']:.4f}")

# Get best model
lasso_best = lasso_cv.best_estimator_

# Coefficients
lasso_coef = pd.DataFrame({
    'Feature': feature_names,
    'Coefficient': lasso_best.coef_,
    'Selected': lasso_best.coef_ != 0
}).sort_values('Coefficient', key=abs, ascending=False)

print("\nLASSO coefficients at optimal alpha:")
print(lasso_coef.to_string(index=False))

n_selected = (lasso_best.coef_ != 0).sum()
print(f"\nNumber of selected features: {n_selected} out of {p}")

# Predictions on test set
lasso_pred = lasso_best.predict(X_test_scaled)
lasso_mse = mean_squared_error(Y_test, lasso_pred)
lasso_r2 = r2_score(Y_test, lasso_pred)

print(f"\nTest MSE: {lasso_mse:.4f}")
print(f"Test R²: {lasso_r2:.4f}")

# Plot coefficient paths
plt.figure(figsize=(10, 6))
coef_paths = []

for alpha in alphas:
    lasso_model = Lasso(alpha=alpha, max_iter=10000)
    lasso_model.fit(X_train_scaled, Y_train)
    coef_paths.append(lasso_model.coef_)

coef_paths = np.array(coef_paths)

for i in range(p):
    plt.plot(np.log10(alphas), coef_paths[:, i], label=feature_names[i])

plt.axvline(np.log10(lasso_cv.best_params_['alpha']), 
            color='red', linestyle='--', label='Optimal α')
plt.xlabel('log10(α)')
plt.ylabel('Coefficient Value')
plt.title('LASSO Regression: Coefficient Paths')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/lasso_coefficient_paths.png', dpi=300, bbox_inches='tight')
print("\nCoefficient paths plot saved!")

# Plot feature importance
plt.figure(figsize=(10, 6))
importance_df = lasso_coef.sort_values('Coefficient', key=abs, ascending=True)
colors = ['steelblue' if selected else 'lightgray' 
          for selected in importance_df['Selected']]
plt.barh(importance_df['Feature'], np.abs(importance_df['Coefficient']), 
         color=colors)
plt.xlabel('Absolute Coefficient Value')
plt.title('LASSO Variable Importance')
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/lasso_importance.png', dpi=300, bbox_inches='tight')
print("Variable importance plot saved!")

# ============================================================================
# 6. ELASTIC NET
# ============================================================================

print("\n" + "=" * 70)
print("6. ELASTIC NET (Combines Ridge and LASSO)")
print("=" * 70)

# Elastic Net with cross-validation
elastic_cv = GridSearchCV(ElasticNet(max_iter=10000),
                          param_grid={'alpha': alphas, 
                                     'l1_ratio': [0.1, 0.3, 0.5, 0.7, 0.9]},
                          cv=10,
                          scoring='neg_mean_squared_error')
elastic_cv.fit(X_train_scaled, Y_train)

print(f"\nOptimal alpha: {elastic_cv.best_params_['alpha']:.4f}")
print(f"Optimal l1_ratio: {elastic_cv.best_params_['l1_ratio']:.2f}")
print("(l1_ratio: 0 = Ridge, 1 = LASSO)")

# Get best model
elastic_best = elastic_cv.best_estimator_

# Coefficients
elastic_coef = pd.DataFrame({
    'Feature': feature_names,
    'Coefficient': elastic_best.coef_,
    'Selected': elastic_best.coef_ != 0
})

print("\nElastic Net coefficients:")
print(elastic_coef.to_string(index=False))

# Predictions
elastic_pred = elastic_best.predict(X_test_scaled)
elastic_mse = mean_squared_error(Y_test, elastic_pred)
elastic_r2 = r2_score(Y_test, elastic_pred)

print(f"\nTest MSE: {elastic_mse:.4f}")
print(f"Test R²: {elastic_r2:.4f}")

# ============================================================================
# 7. COMPARISON OF ALL METHODS
# ============================================================================

print("\n" + "=" * 70)
print("7. MODEL COMPARISON")
print("=" * 70)

# Fit OLS for comparison
ols_model = LinearRegression()
ols_model.fit(X_train_scaled, Y_train)
ols_pred = ols_model.predict(X_test_scaled)
ols_mse = mean_squared_error(Y_test, ols_pred)
ols_r2 = r2_score(Y_test, ols_pred)

# Create comparison DataFrame
comparison = pd.DataFrame({
    'Method': ['OLS (Full)', 'Ridge', 'LASSO', 'Elastic Net'],
    'Test_MSE': [ols_mse, ridge_mse, lasso_mse, elastic_mse],
    'Test_R2': [ols_r2, ridge_r2, lasso_r2, elastic_r2],
    'N_Variables': [p, p, n_selected, (elastic_best.coef_ != 0).sum()]
})

print("\nTest Set Performance:")
print(comparison.to_string(index=False))

best_method = comparison.loc[comparison['Test_MSE'].idxmin(), 'Method']
print(f"\nBest method by Test MSE: {best_method}")

# Visualization
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# MSE comparison
axes[0].bar(comparison['Method'], comparison['Test_MSE'], color='steelblue')
axes[0].set_ylabel('Test MSE')
axes[0].set_title('Model Comparison: Test MSE')
axes[0].tick_params(axis='x', rotation=45)

# Number of variables
axes[1].bar(comparison['Method'], comparison['N_Variables'], color='coral')
axes[1].set_ylabel('Number of Variables')
axes[1].set_title('Model Complexity')
axes[1].tick_params(axis='x', rotation=45)

plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/model_comparison.png', dpi=300, bbox_inches='tight')
print("\nModel comparison plot saved!")

# ============================================================================
# 8. PRACTICAL EXAMPLE: PREDICTING GDP GROWTH
# ============================================================================

print("\n" + "=" * 70)
print("8. PRACTICAL EXAMPLE: GDP GROWTH PREDICTION")
print("=" * 70)

# Generate GDP growth dataset
np.random.seed(123)
n_countries = 150

gdp_features = pd.DataFrame({
    'Inflation': np.random.normal(3, 2, n_countries),
    'Unemployment': np.random.normal(6, 3, n_countries),
    'Investment': np.random.normal(20, 5, n_countries),
    'Government_Spending': np.random.normal(18, 4, n_countries),
    'Exports': np.random.normal(25, 8, n_countries),
    'Imports': np.random.normal(23, 7, n_countries),
    'FDI': np.random.normal(5, 3, n_countries),
    'Debt_GDP': np.random.normal(60, 20, n_countries),
    'Population_Growth': np.random.normal(1.2, 0.8, n_countries),
    'Technology_Index': np.random.normal(50, 15, n_countries)
})

# True relationship (some variables don't matter)
gdp_features['GDP_Growth'] = (5 + 
    -0.5 * gdp_features['Inflation'] +
    -0.3 * gdp_features['Unemployment'] +
    0.4 * gdp_features['Investment'] +
    0.1 * gdp_features['FDI'] +
    np.random.normal(0, 1, n_countries))

# Split data
train_idx = np.random.choice(n_countries, int(0.7 * n_countries), replace=False)
test_idx = np.setdiff1d(np.arange(n_countries), train_idx)

X_gdp_train = gdp_features.drop('GDP_Growth', axis=1).iloc[train_idx].values
Y_gdp_train = gdp_features['GDP_Growth'].iloc[train_idx].values
X_gdp_test = gdp_features.drop('GDP_Growth', axis=1).iloc[test_idx].values
Y_gdp_test = gdp_features['GDP_Growth'].iloc[test_idx].values

# Standardize
scaler_gdp = StandardScaler()
X_gdp_train_scaled = scaler_gdp.fit_transform(X_gdp_train)
X_gdp_test_scaled = scaler_gdp.transform(X_gdp_test)

# Apply LASSO
gdp_lasso = GridSearchCV(Lasso(max_iter=10000),
                         param_grid={'alpha': np.logspace(-3, 2, 50)},
                         cv=10,
                         scoring='neg_mean_squared_error')
gdp_lasso.fit(X_gdp_train_scaled, Y_gdp_train)

# Results
gdp_coef = pd.DataFrame({
    'Feature': gdp_features.drop('GDP_Growth', axis=1).columns,
    'Coefficient': gdp_lasso.best_estimator_.coef_,
    'Selected': gdp_lasso.best_estimator_.coef_ != 0
})

print("\nLASSO coefficients for GDP growth model:")
print(gdp_coef.to_string(index=False))

# Predictions
gdp_pred = gdp_lasso.predict(X_gdp_test_scaled)
gdp_mse = mean_squared_error(Y_gdp_test, gdp_pred)
gdp_r2 = r2_score(Y_gdp_test, gdp_pred)

print(f"\nModel Performance:")
print(f"Test MSE: {gdp_mse:.4f}")
print(f"Test R²: {gdp_r2:.4f}")
print(f"Selected features: {(gdp_lasso.best_estimator_.coef_ != 0).sum()}")

# Plot actual vs predicted
plt.figure(figsize=(10, 6))
plt.scatter(Y_gdp_test, gdp_pred, alpha=0.6, color='steelblue')
plt.plot([Y_gdp_test.min(), Y_gdp_test.max()], 
         [Y_gdp_test.min(), Y_gdp_test.max()], 
         'r--', lw=2, label='Perfect Prediction')
plt.xlabel('Actual GDP Growth (%)')
plt.ylabel('Predicted GDP Growth (%)')
plt.title('LASSO: Actual vs Predicted GDP Growth')
plt.legend()
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/gdp_prediction.png', dpi=300, bbox_inches='tight')
print("\nGDP prediction plot saved!")

# ============================================================================
# FINAL SUMMARY
# ============================================================================

print("\n" + "=" * 70)
print("ANALYSIS COMPLETE!")
print("=" * 70)
print("\nGenerated plots:")
print("- ridge_coefficient_paths.png")
print("- ridge_cv_error.png")
print("- lasso_coefficient_paths.png")
print("- lasso_importance.png")
print("- model_comparison.png")
print("- gdp_prediction.png")
print("\nAll plots saved to /mnt/user-data/outputs/")