# ML Program 6 - Clustering: KMeans vs Gaussian Mixture (EM)
from sklearn.cluster import KMeans
from sklearn.mixture import GaussianMixture
from sklearn.preprocessing import StandardScaler
from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt
import numpy as np

X, y_true = make_blobs(
    n_samples=500,
    centers=[(0, 0), (3, 3), (0, 4)],
    cluster_std=[1.5, 1.5, 1.5],
    random_state=42
)

plt.figure(figsize=(15, 5))
plt.subplot(1, 3, 1)
plt.scatter(X[:, 0], X[:, 1], c=y_true, s=40)
plt.title("Real Data")

# ------------------------------------------------------------------

kmeans = KMeans(n_clusters=3, random_state=42)
k_pred = kmeans.fit_predict(X)

plt.subplot(1, 3, 2)
plt.scatter(X[:, 0], X[:, 1], c=k_pred, s=40)
plt.title("KMeans (fails on overlap)")

# ------------------------------------------------------------------

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

gmm = GaussianMixture(n_components=3, random_state=42)
gmm_pred = gmm.fit_predict(X_scaled)

plt.subplot(1, 3, 3)
plt.scatter(X[:, 0], X[:, 1], c=gmm_pred, s=40)
plt.title("GMM (handles overlap better)")

plt.show()
