# Install packages
!pip install ucimlrepo

# Import libraries
from ucimlrepo import fetch_ucirepo
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import seaborn as sns
import warnings
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
from pyspark.sql.functions import col, count, when, sum as spark_sum
from pyspark.ml import Pipeline
from pyspark.ml.feature import (StringIndexer, OneHotEncoder, VectorAssembler, StandardScaler)
from pyspark.ml.classification import (LogisticRegression, RandomForestClassifier, GBTClassifier)
from pyspark.ml.evaluation import (BinaryClassificationEvaluator, MulticlassClassificationEvaluator)
from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
spark = SparkSession.builder.appName("PySpark ML").getOrCreate()
import time

# ── Plotting config ──────────────────────────────────────────────────────────
sns.set_theme(style = "whitegrid", palette = "muted")
plt.rcParams.update({"figure.dpi": 120, "figure.figsize": (10, 4)})
warnings.filterwarnings("ignore")

# Enable Apache Arrow optimization
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# Fetch dataset from UCI machine learning
adult = fetch_ucirepo(id = 2)

# Extract the data
df = adult.data.features
y = adult.data.targets

# metadata
print(f"Data Download Information: Dataset Name- {adult.metadata['name']} -- {adult.metadata['data_url']}")
print(f"Purpose - {adult.metadata['abstract']}")

# Normalise the target labels (strip whitespace, collapse '>50K.' -> '>50K')
y["income"] = y["income"].str.strip().str.replace(".", "", regex = False)

# Combine into one pandas DataFrame, then convert to PySpark
df = pd.concat([df, y], axis = 1)

# Convert the data to PySpark
pyspark_df = spark.createDataFrame(df)

# Utilities - function to generate output labels
def get_proba(preds):

    pdf = preds.select('label', 'probability').toPandas()
    pdf['prob_pos'] = pdf['probability'].apply(lambda v: float(v[1]))

    return pdf['label'].values, pdf['prob_pos'].values

print('Dataset Inspection..\n')
print("=== Schema ===")
pyspark_df.printSchema()

# Inspect the first 5 rows
pyspark_df.show(5)

# Summary statistics for numeric columns
numeric_cols = [f.name for f in pyspark_df.schema.fields if str(f.dataType) in ("IntegerType()", "LongType()", "DoubleType()", "FloatType()")]

print("\n=== Numeric Summary ===")
pyspark_df.select(numeric_cols).describe().show(truncate=False)

# Use pyspark syntax ro count nulls and missing
missing_counts = pyspark_df.select([spark_sum(when(col(c).cast("string") == "?", 1).otherwise(0)).alias(c) for c in pyspark_df.columns]).toPandas().T.rename(columns = {0: "missing_count"})

missing_counts["missing_pct"] = (missing_counts["missing_count"] / pyspark_df.count() * 100).round(2)
missing_counts = missing_counts[missing_counts["missing_count"] > 0].sort_values("missing_count", ascending = False)

print("Columns with missing values ('?'):")
print(missing_counts.to_string())

# Compute class distribution
target_dist = (pyspark_df.groupBy("income").count().withColumn("pct", F.round(F.col("count") / pyspark_df.count() * 100, 2)).orderBy("income").toPandas())
print("Target distribution:")
print(target_dist.to_string(index=False))

# Plot
fig, axes = plt.subplots(figsize = (10, 4))

palette = sns.color_palette("muted", 2)

# Bar chart
axes.bar(target_dist["income"], target_dist["count"], color=palette, edgecolor="white", linewidth = 1.2)

for i, row in target_dist.iterrows():

    axes.text(i, row["count"] + 200, f"{row['count']:,}\n({row['pct']}%)", ha = "center", fontsize = 8)

axes.set_title("Income Class Counts", fontweight="bold")
axes.set_ylabel("Count")
axes.yaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))

plt.suptitle("Target Variable: Annual Income", fontsize = 13, fontweight = "bold", y=1.01)
plt.tight_layout()
plt.show()

print("\nClass Ratio (<=50K : >50K):", round(target_dist.loc[0,'count'] / target_dist.loc[1,'count'], 2), ": 1")

num_features = ["age", "fnlwgt", "education-num", "capital-gain", "capital-loss", "hours-per-week"]

# Bring numeric columns to Pandas for plotting
num_pd = pyspark_df.select(num_features + ["income"]).toPandas()

# Distribution plots (histograms + KDE)
fig, axes = plt.subplots(2, 3, figsize=(18, 10))
axes = axes.flatten()

for i, col_name in enumerate(num_features):

    sns.histplot(num_pd[col_name], ax=axes[i], kde = True, bins = 40, color = sns.color_palette("bright")[i % 6])
    axes[i].set_title(f"Distribution of {col_name}", fontweight = "bold")
    axes[i].set_xlabel(col_name)
    axes[i].set_ylabel("Count")
    skew_val = num_pd[col_name].skew()
    axes[i].text(0.98, 0.92, f"Skew: {skew_val:.2f}", transform = axes[i].transAxes, ha = "right", fontsize = 8, bbox = dict(boxstyle = "round,pad=0.2", fc="white", alpha=0.7))

plt.suptitle("Numerical Feature Distributions", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()

# --- Box plots: distributions by income class ---
fig, axes = plt.subplots(2, 3, figsize=(18, 10))

axes = axes.flatten()

for i, col_name in enumerate(num_features):

    sns.boxplot(data=num_pd, x="income", y=col_name, ax=axes[i], palette = "bright", order = ["<=50K", ">50K"])
    axes[i].set_title(f"{col_name} by Income", fontweight="bold")
    axes[i].set_xlabel("")

plt.suptitle("Numerical Features Split by Income Class", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()

# Index the category fts
cat_features = ["workclass", "education", "marital-status", "occupation", "relationship", "race", "sex", "native-country"]
cat_pd = pyspark_df.select(cat_features + ["income"]).toPandas()

# Remove '?' placeholder rows for plotting
cat_pd = cat_pd.replace("?", np.nan).dropna()

fig, axes = plt.subplots(4, 2, figsize=(16, 22))
axes = axes.flatten()

for i, col_name in enumerate(cat_features):

    order = cat_pd[col_name].value_counts().index
    sns.countplot(data = cat_pd, y = col_name, ax = axes[i], order = order, palette = "bright")
    axes[i].set_title(f"{col_name}  (unique: {cat_pd[col_name].nunique()})", fontweight = "bold")
    axes[i].set_xlabel("Count")
    axes[i].set_ylabel("")

plt.suptitle("Categorical Feature Distributions", fontsize = 15, fontweight = "bold", y=1.002)
plt.tight_layout()
plt.show()

# Cardinality summary (useful for encoding strategy decisions)
card_df = pd.DataFrame({"column": cat_features,
                        "unique_values": [pyspark_df.select(c).distinct().count() for c in cat_features]}).sort_values("unique_values", ascending=False)

print("Cardinality of categorical features:")
print(card_df.to_string(index = False))
print("\nEncoding recommendation:")
print("  - Low cardinality (< 10): One-Hot Encoding via StringIndexer + OneHotEncoder")
print("  - High cardinality (native-country, 42 values): Target encoding or frequency grouping")

# Stacked bar: proportion of >50K earners per category
high_cardinality_skip = ["native-country"]  # too many categories for this chart
cat_for_bivar = [c for c in cat_features if c not in high_cardinality_skip]

fig, axes = plt.subplots(4, 2, figsize=(18, 21))
axes = axes.flatten()

for i, col_name in enumerate(cat_for_bivar):

    ct = pd.crosstab(cat_pd[col_name], cat_pd["income"], normalize="index") * 100
    ct = ct.reindex(ct[">50K"].sort_values(ascending=False).index)  # sort by % high earners
    ct.plot(kind="barh", stacked=True, ax=axes[i], color=sns.color_palette("bright", 2), edgecolor="white", linewidth=0.8)

    axes[i].set_title(f"{col_name}: Income Breakdown (%)", fontweight="bold")
    axes[i].set_xlabel("% of group")
    axes[i].set_ylabel("")
    axes[i].legend(title="Income", loc="lower right", fontsize=8)
    axes[i].axvline(x=25, color="grey", linestyle="--", linewidth=0.8, alpha=0.6)

# Hide unused subplot
if len(cat_for_bivar) < len(axes):

    for j in range(len(cat_for_bivar), len(axes)):

        axes[j].set_visible(False)

plt.suptitle("Income Distribution Within Each Categorical Feature", fontsize = 15, fontweight = "bold", y = 1.002)
plt.tight_layout()
plt.show()

# Correlation heatmap (numeric features only)
corr_pd = num_pd[num_features].copy()
corr_matrix = corr_pd.corr()

fig, ax = plt.subplots(figsize = (9, 7))
mask = np.triu(np.ones_like(corr_matrix, dtype=bool))
sns.heatmap(corr_matrix, mask=mask, annot=True, fmt=".2f", cmap = "coolwarm", center=0, linewidths=0.5, ax = ax, cbar_kws = {"shrink": 0.8})
ax.set_title("Pearson Correlation Matrix — Numeric Features", fontweight="bold", fontsize=13)
plt.tight_layout()
plt.show()

# Point-biserial correlation with target (numeric features vs income)
num_pd["income_binary"] = (num_pd["income"] == ">50K").astype(int)

target_corr = (num_pd[num_features + ["income_binary"]].corr()["income_binary"].drop("income_binary").sort_values(key = abs, ascending = False))

fig, ax = plt.subplots(figsize = (18, 3))
colors = ["#2E86AB" if v >= 0 else "#E84855" for v in target_corr.values]
bars = ax.barh(target_corr.index, target_corr.values, color=colors, edgecolor="white")
ax.axvline(0, color="black", linewidth=0.8)
ax.bar_label(bars, fmt="%.3f", padding=4, fontsize=9)
ax.set_title("Correlation of Numeric Features with Income (>50K)", fontweight="bold")
ax.set_xlabel("Pearson / Point-Biserial Correlation")
plt.tight_layout()
plt.show()

# Log1p-transform skewed capital columns (as identified in Key Insights)
pyspark_df = pyspark_df.withColumn('capital-gain', F.log1p(F.col('capital-gain')))
pyspark_df = pyspark_df.withColumn('capital-loss', F.log1p(F.col('capital-loss')))

print("Log1p transform applied to capital-gain and capital-loss.")

# Encode the target variable
pyspark_df = pyspark_df.withColumn('label', F.when(F.col('income') == '>50K', 1.0).otherwise(0.0))

# Drop the nonuseful columns and outcome variable
train_df, test_df = pyspark_df.drop('fnlwgt', 'education', 'income').randomSplit([0.8, 0.2], seed = 100)

train_df.printSchema()

# Split
print(f"Train rows: {train_df.count():,} | Test rows: {test_df.count():,}")

# Separate column types for model prep
cat_cols = ['workclass', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'native-country']
num_cols = ['age', 'education-num', 'capital-gain', 'capital-loss', 'hours-per-week']

# Init model compoinents
indexers = [StringIndexer(inputCol = c, outputCol = f'{c}_idx', handleInvalid = 'keep') for c in cat_cols]

# OneHotEncoder: convert indices to sparse binary vectors
encoders = [OneHotEncoder(inputCol = f'{c}_idx', outputCol = f'{c}_ohe') for c in cat_cols]

# VectorAssembler: combine all features into one vector
assembler_inputs = num_cols + [f'{c}_ohe' for c in cat_cols]
assembler = VectorAssembler(inputCols = assembler_inputs, outputCol = 'raw_features', handleInvalid = 'keep')

preprocessing_stages = indexers + encoders + [assembler]
print(f'Pipeline preprocessing stages: {len(preprocessing_stages)}')

# Init models
lr  = LogisticRegression(featuresCol = 'raw_features', labelCol = 'label', maxIter = 100)
rf  = RandomForestClassifier(featuresCol = 'raw_features', labelCol = 'label', seed = 100)
gbt = GBTClassifier(featuresCol = 'raw_features', labelCol = 'label', maxIter = 50, seed = 100)

pipeline_lr  = Pipeline(stages = preprocessing_stages + [lr])
pipeline_rf  = Pipeline(stages = preprocessing_stages + [rf])
pipeline_gbt = Pipeline(stages = preprocessing_stages + [gbt])

start = time.time()

print("Training Logistic Regression..")
model_lr  = pipeline_lr.fit(train_df)

print("Training Random Forest..")
model_rf  = pipeline_rf.fit(train_df)

print("Training Gradient Boosted Trees..")
model_gbt = pipeline_gbt.fit(train_df)

end = time.time()

print(f"All models trained- elapsed time {round((end - start), 3) / 60} minutes..")

# Get predictions
lrpreds  = model_lr.transform(test_df)
rfpreds  = model_rf.transform(test_df)
gbpreds = model_gbt.transform(test_df)

# Init scoring metrics
bin_eval = BinaryClassificationEvaluator(labelCol = 'label', rawPredictionCol = 'rawPrediction', metricName = 'areaUnderROC')
mc_eval = MulticlassClassificationEvaluator(labelCol = 'label', predictionCol = 'prediction')

# Compile the scores
results = {}

for name, preds in [("LR", lrpreds), ("RF", rfpreds), ("GBT", gbpreds)]:

    results[name] = {"AUC-ROC":   bin_eval.evaluate(preds),
                     "Accuracy":  mc_eval.setMetricName("accuracy").evaluate(preds),
                     "F1":        mc_eval.setMetricName("f1").evaluate(preds),
                     "Precision": mc_eval.setMetricName("weightedPrecision").evaluate(preds),
                     "Recall":    mc_eval.setMetricName("weightedRecall").evaluate(preds)}

scores = round(pd.DataFrame(results).T, 4)

scores

fig, axes = plt.subplots(1, 3, figsize = (18, 4))

for ax, (preds, name) in zip(axes, [(lrpreds, "Logistic Regression"), (rfpreds, "Random Forest"), (gbpreds, "Gradient Boosted Trees")]):

    cm = preds.groupBy('label', 'prediction').count().toPandas()
    cm_pivot = cm.pivot(index='label', columns='prediction', values='count').fillna(0).astype(int)
    sns.heatmap(cm_pivot, annot=True, fmt='d', cmap='Blues', ax=ax)
    ax.set_title(f'Confusion Matrix — {name}')
    ax.set_ylabel('Actual')
    ax.set_xlabel('Predicted')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 3, figsize = (18, 5))

for ax, (name, preds) in zip(axes, [("LR", lrpreds), ("RF", rfpreds), ("GBT", gbpreds)]):

    y_true, y_score = get_proba(preds)
    fpr, tpr, _ = roc_curve(y_true, y_score)
    roc_auc = auc(fpr, tpr)
    ax.plot(fpr, tpr, label=f"{name} (AUC = {roc_auc:.3f})")
    ax.plot([0,1],[0,1],'k--'); ax.set_xlabel('FPR'); ax.set_ylabel('TPR')
    ax.set_title(f'{name} ROC Curve');
    ax.legend();

plt.tight_layout();
plt.show()

# Precision-Recall curves — more informative than ROC under class imbalance
fig, axes = plt.subplots(1, 3, figsize = (18, 5))

for ax, (name, preds) in zip(axes, [("Logistic Regression", lrpreds),
                                     ("Random Forest",       rfpreds),
                                     ("Gradient Boosted Trees", gbpreds)]):

    y_true, y_score = get_proba(preds)
    precision, recall, _ = precision_recall_curve(y_true, y_score)
    ap = average_precision_score(y_true, y_score)
    baseline = y_true.mean()

    ax.plot(recall, precision, label=f"AP = {ap:.3f}")
    ax.axhline(baseline, color = "gray", linestyle = "--", linewidth = 0.9, label = f"Baseline ({baseline:.2f})")
    ax.set_xlabel("Recall")
    ax.set_ylabel("Precision")
    ax.set_title(f"PR Curve — {name}")
    ax.legend()

plt.suptitle("Precision-Recall Curves", fontweight = "bold", y = 1.02)
plt.tight_layout()
plt.show()

# Get ft importance
rf_model = model_rf.stages[-1]
importances = rf_model.featureImportances
feature_names = num_cols + [f'{c}_ohe' for c in cat_cols]

imp_df = pd.DataFrame({'feature': feature_names[:len(importances)],
                       'importance': importances.toArray()[:len(feature_names)]})
imp_df = imp_df.sort_values('importance', ascending = False).head(15)

plt.figure(figsize = (18, 7))
sns.barplot(data = imp_df, y = 'feature', x = 'importance', palette = 'viridis')
plt.title('Random Forest — Top Feature Importances')
plt.tight_layout(); plt.show()