#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Oct  6 13:08:43 2024

@author: rahulkulkarni
"""

# Importing stuff
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.model_selection import KFold, GridSearchCV
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

import joblib
import warnings
warnings.filterwarnings('ignore')

from tqdm import tqdm
import time

# Load the data
data_path = '/Users/rahulkulkarni/Downloads/Data/Proportional_Data_CSV.csv'

try:
    df = pd.read_csv(data_path)
    print("Loaded the data!")
except FileNotFoundError:
    print("File not found, check the path!")
    exit()
except pd.errors.ParserError:
    print("Can't parse the file, something's wrong with it.")
    exit()

# Drop some columns we don't need
exclude_columns = ['LGA Name', 'Region', '2012']
existing_exclude = [col for col in exclude_columns if col in df.columns]
df = df.drop(columns=existing_exclude, errors='ignore')

# Handle missing values
target = 'EGM expenditure per capita'

if target not in df.columns:
    raise ValueError(f"Can't find {target} in the data.")

df = df.dropna(subset=[target])

numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
object_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()

if numeric_cols:
    df[numeric_cols] = df[numeric_cols].fillna(df[numeric_cols].median())

if object_cols:
    for col in object_cols:
        if df[col].isnull().sum() > 0:
            mode_val = df[col].mode()[0]
            df[col].fillna(mode_val, inplace=True)

# Get features and target
features = df.drop([target], axis=1, errors='ignore')
X = features.copy()
y = df[target].copy()

# Encode categorical variables
categorical_cols = X.select_dtypes(include=['object', 'category']).columns.tolist()

if categorical_cols:
    X = pd.get_dummies(X, columns=categorical_cols, drop_first=True)

feature_names = X.columns.tolist()

# Remove any features with 'LGA' in the name
X = X.loc[:, ~X.columns.str.contains('LGA', case=False)]
feature_names = X.columns.tolist()

# Print features
print("\nFeatures after encoding and removing 'LGA':")
for i, feature in enumerate(feature_names, 1):
    print(f"{i}. {feature}")

# Set up Random Forest with GridSearch and K-Fold
k_folds = 5
kf = KFold(n_splits=k_folds, shuffle=True, random_state=42)

param_grid = {
    'n_estimators': [100, 200, 300],
    'max_depth': [None, 10, 20, 30],
    'min_samples_split': [2, 5, 10],
    'min_samples_leaf': [1, 2, 4]
}

rf = RandomForestRegressor(random_state=42)

grid_search = GridSearchCV(
    estimator=rf,
    param_grid=param_grid,
    cv=kf,
    n_jobs=-1,
    scoring='r2',
    verbose=1
)

print("\nStarting Grid Search...")
start_time = time.time()
grid_search.fit(X, y)
end_time = time.time()
print(f"Grid Search done in {end_time - start_time:.2f} seconds.\n")

print("Best params:", grid_search.best_params_)
best_rf = grid_search.best_estimator_

# Evaluate with cross-validation
r2_scores = []
mae_scores = []
mse_scores = []
rmse_scores = []
feature_importances_folds = []
residuals_all_folds = []

print("\nEvaluating with K-Fold CV...\n")

for fold, (train_idx, test_idx) in enumerate(kf.split(X), 1):
    X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
    y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
    
    best_rf.fit(X_train, y_train)
    y_pred = best_rf.predict(X_test)
    
    residuals = y_test - y_pred
    residuals_all_folds.append({
        'Actual': y_test,
        'Predicted': y_pred,
        'Residual': residuals,
        'Fold': fold
    })
    
    mae = mean_absolute_error(y_test, y_pred)
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_test, y_pred)
    
    r2_scores.append(r2)
    mae_scores.append(mae)
    mse_scores.append(mse)
    rmse_scores.append(rmse)
    
    feature_importances_folds.append(best_rf.feature_importances_)
    
    print(f"Fold {fold}: R²={r2:.4f}, MAE={mae:.2f}, MSE={mse:.2f}, RMSE={rmse:.2f}")

mean_r2 = np.mean(r2_scores)
std_r2 = np.std(r2_scores)
mean_mae = np.mean(mae_scores)
mean_mse = np.mean(mse_scores)
mean_rmse = np.mean(rmse_scores)

print("\nPerformance:")
print(f"Mean R²: {mean_r2:.4f} ± {std_r2:.4f}")
print(f"Mean MAE: {mean_mae:.2f}")
print(f"Mean MSE: {mean_mse:.2f}")
print(f"Mean RMSE: {mean_rmse:.2f}")

# Feature Importance
fi_df = pd.DataFrame(feature_importances_folds, columns=feature_names)
feature_stats = fi_df.agg(['min', 'mean', 'max']).transpose().reset_index()
feature_stats.columns = ['Feature', 'Min', 'Mean', 'Max']

print("\nFeature Importances:")
print(feature_stats.sort_values(by='Mean', ascending=False).to_string(index=False))

top_15 = feature_stats.sort_values(by='Mean', ascending=False).head(15)

plt.figure(figsize=(12,8))
sns.barplot(x='Mean', y='Feature', data=top_15, color='blue')
plt.title('Top 15 Features')
plt.tight_layout()
plt.show()

print("\nTop 15 Features:")
print(top_15[['Feature', 'Mean']].to_string(index=False))

# Boxplot for feature importances
fi_melt = fi_df.melt(var_name='Feature', value_name='Importance')
feature_order = fi_df.mean().sort_values(ascending=False).index

plt.figure(figsize=(14,10))
sns.boxplot(x='Importance', y='Feature', data=fi_melt, order=feature_order, color='green')
plt.title('Feature Importances Boxplot')
plt.tight_layout()
plt.show()

# R² Scores Distribution
r2_df = pd.DataFrame(r2_scores, columns=['R2_Score'])

plt.figure(figsize=(10,6))
sns.histplot(r2_df['R2_Score'], bins=20, kde=True, color='orange')
plt.title('R² Scores Distribution')
plt.xlabel('R² Score')
plt.ylabel('Frequency')
plt.grid(True)
plt.tight_layout()
plt.show()

plt.figure(figsize=(8,6))
sns.boxplot(x=r2_df['R2_Score'], color='purple')
plt.title('R² Scores Boxplot')
plt.xlabel('R² Score')
plt.grid(True)
plt.tight_layout()
plt.show()

print("\nR² Scores Stats:")
print(r2_df['R2_Score'].describe())

# Feature Selection based on importance
mean_importance = fi_df.mean().sort_values(ascending=False)
print("\nMean Feature Importances:")
print(mean_importance)

current_features = mean_importance.index.tolist()
initial_num = len(current_features)
best_mean_r2 = -np.inf
best_features = current_features.copy()
best_iter = 0
mean_r2_per_iter = []
iter_count = 0

print("\nStarting Feature Selection...")

while len(current_features) >=1:
    iter_count +=1
    print(f"\nIteration {iter_count}: {len(current_features)} features")
    
    r2_fs = []
    for fold, (train_idx, test_idx) in enumerate(kf.split(X),1):
        X_train_fs, X_test_fs = X.iloc[train_idx][current_features], X.iloc[test_idx][current_features]
        y_train_fs, y_test_fs = y.iloc[train_idx], y.iloc[test_idx]
        
        best_rf.fit(X_train_fs, y_train_fs)
        y_pred_fs = best_rf.predict(X_test_fs)
        r2_fs.append(r2_score(y_test_fs, y_pred_fs))
    
    mean_r2_fs = np.mean(r2_fs)
    mean_r2_per_iter.append(mean_r2_fs)
    print(f"Mean R²: {mean_r2_fs:.4f}")
    
    if mean_r2_fs > best_mean_r2:
        best_mean_r2 = mean_r2_fs
        best_features = current_features.copy()
        best_iter = iter_count
    
    if len(current_features) >1:
        least_imp = mean_importance[current_features].idxmin()
        print(f"Removing {least_imp}")
        current_features.remove(least_imp)
    else:
        print("Only one feature left, stopping.")
        break

print("\nFeature Selection Done.")
print(f"Best Iteration: {best_iter} with R²: {best_mean_r2:.4f}")
print("Best Features:")
print(best_features)

# Summary of feature selection
summary = pd.DataFrame({
    'Iteration': range(1, len(mean_r2_per_iter)+1),
    'Num_Features': [initial_num -i for i in range(len(mean_r2_per_iter))],
    'Mean_R2': mean_r2_per_iter
})

print("\nFeature Selection Summary:")
print(summary.to_string(index=False))

# Retrain with best features
X_opt = X[best_features]

categorical_opt = X_opt.select_dtypes(include=['object', 'category']).columns.tolist()
if categorical_opt:
    X_opt = pd.get_dummies(X_opt, columns=categorical_opt, drop_first=True)

rf_opt = RandomForestRegressor(
    n_estimators=grid_search.best_params_.get('n_estimators',100),
    max_depth=grid_search.best_params_.get('max_depth',None),
    min_samples_split=grid_search.best_params_.get('min_samples_split',2),
    min_samples_leaf=grid_search.best_params_.get('min_samples_leaf',1),
    random_state=42
)
rf_opt.fit(X_opt, y)

# Evaluate again with best features
r2_opt = []
mae_opt = []
mse_opt = []
rmse_opt = []
residuals_opt = []

print("\nEvaluating with Best Features...")

for fold, (train_idx, test_idx) in enumerate(kf.split(X_opt),1):
    X_train_opt, X_test_opt = X_opt.iloc[train_idx], X_opt.iloc[test_idx]
    y_train_opt, y_test_opt = y.iloc[train_idx], y.iloc[test_idx]
    
    rf_opt.fit(X_train_opt, y_train_opt)
    y_pred_opt = rf_opt.predict(X_test_opt)
    
    residual = y_test_opt - y_pred_opt
    residuals_opt.append({
        'Actual': y_test_opt,
        'Predicted': y_pred_opt,
        'Residual': residual,
        'Fold': fold
    })
    
    mae_v = mean_absolute_error(y_test_opt, y_pred_opt)
    mse_v = mean_squared_error(y_test_opt, y_pred_opt)
    rmse_v = np.sqrt(mse_v)
    r2_v = r2_score(y_test_opt, y_pred_opt)
    
    r2_opt.append(r2_v)
    mae_opt.append(mae_v)
    mse_opt.append(mse_v)
    rmse_opt.append(rmse_v)
    
    print(f"Fold {fold}: R²={r2_v:.4f}, MAE={mae_v:.2f}, MSE={mse_v:.2f}, RMSE={rmse_v:.2f}")

mean_r2_o = np.mean(r2_opt)
std_r2_o = np.std(r2_opt)
mean_mae_o = np.mean(mae_opt)
mean_mse_o = np.mean(mse_opt)
mean_rmse_o = np.mean(rmse_opt)

print("\nPerformance with Best Features:")
print(f"Mean R²: {mean_r2_o:.4f} ± {std_r2_o:.4f}")
print(f"Mean MAE: {mean_mae_o:.2f}")
print(f"Mean MSE: {mean_mse_o:.2f}")
print(f"Mean RMSE: {mean_rmse_o:.2f}")

# Feature importance for optimal features
importances_opt = rf_opt.feature_importances_
fi_opt_df = pd.DataFrame({
    'Feature': X_opt.columns,
    'Importance': importances_opt
}).sort_values(by='Importance', ascending=False)

plt.figure(figsize=(12,8))
sns.barplot(x='Importance', y='Feature', data=fi_opt_df.head(15), color='red')
plt.title('Top 15 Features (Optimal Set)')
plt.tight_layout()
plt.show()

print("\nTop 15 Features (Optimal):")
print(fi_opt_df.head(15).to_string(index=False))

# Summary again
print("\nFeature Selection Summary:")
print(summary.to_string(index=False))

# Save the model
model_file = 'random_forest_model_optimal.pkl'
joblib.dump(rf_opt, model_file)
print(f"\nModel saved as {model_file}")


Based on the above code, explain in the form of an animation (it can be consecutive images in python) How a random forest regression method works. Base your explanation of the datasets used above using plausable values when the data doesn't exist, to show how the regression works in the above scenario. Write the code in python in full