"""From-scratch clustering: k-means (Lloyd's algorithm) and a Gaussian mixture
model fit by Expectation-Maximization.

k-means makes HARD assignments to the nearest centroid and implicitly assumes
spherical, equal-size clusters (it minimizes within-cluster squared distance).
A GMM generalizes it: each cluster is a full Gaussian (its own mean AND
covariance -- so clusters can be elliptical and overlapping), and EM produces
SOFT assignments (a posterior probability of belonging to each cluster).  In
fact k-means is the limit of GMM-EM with shared spherical covariance shrunk to
zero, so this file builds the two as one family.

Both are validated against scikit-learn in the notebook.  The Bayesian relative
-- a Dirichlet-process mixture that INFERS the number of clusters instead of
fixing k -- lives in the nonparametrics arc; here k is chosen by the elbow (SSE)
or BIC.
"""
import numpy as np
from scipy.stats import multivariate_normal


def kmeans(X, k, n_iter=100, seed=0):
    """Lloyd's algorithm. Returns (labels, centroids, inertia)."""
    X = np.asarray(X, float); rng = np.random.default_rng(seed)
    C = X[rng.choice(len(X), k, replace=False)].copy()          # random data points as seeds
    for _ in range(n_iter):
        D = ((X[:, None, :] - C[None, :, :]) ** 2).sum(2)       # squared distances to each centroid
        lab = D.argmin(1)
        newC = np.array([X[lab == j].mean(0) if np.any(lab == j) else C[j] for j in range(k)])
        if np.allclose(newC, C):
            break
        C = newC
    inertia = ((X - C[lab]) ** 2).sum()                         # within-cluster SSE
    return lab, C, inertia


class GMM:
    """Gaussian mixture model fit by Expectation-Maximization (full covariance).

    E-step: responsibility r_ij = pi_j N(x_i | mu_j, Sig_j) / sum_l (...).
    M-step: pi_j, mu_j, Sig_j <- responsibility-weighted proportions/means/covs.
    Initialized from k-means; `reg` adds a ridge to each covariance for stability.
    """

    def __init__(self, k, n_iter=200, tol=1e-6, reg=1e-6, seed=0):
        self.k = k; self.n_iter = n_iter; self.tol = tol; self.reg = reg; self.seed = seed

    def fit(self, X):
        X = np.asarray(X, float); n, d = X.shape
        lab, C, _ = kmeans(X, self.k, seed=self.seed)           # k-means initialization
        self.mu = C.copy()
        self.Sig = np.array([np.cov(X[lab == j].T) + self.reg * np.eye(d)
                             if np.sum(lab == j) > 1 else np.eye(d) for j in range(self.k)])
        self.pi = np.array([max((lab == j).mean(), 1e-3) for j in range(self.k)]); self.pi /= self.pi.sum()
        self.loglik_ = []
        ll_old = -np.inf
        for _ in range(self.n_iter):
            R = np.column_stack([self.pi[j] * multivariate_normal.pdf(X, self.mu[j], self.Sig[j], allow_singular=True)
                                 for j in range(self.k)])       # E-step (unnormalized)
            ll = np.log(R.sum(1) + 1e-300).sum(); self.loglik_.append(ll)
            R = R / (R.sum(1, keepdims=True) + 1e-300)          # responsibilities
            Nk = R.sum(0) + 1e-10                               # M-step
            self.pi = Nk / n
            self.mu = (R.T @ X) / Nk[:, None]
            for j in range(self.k):
                Xc = X - self.mu[j]
                self.Sig[j] = (R[:, j, None] * Xc).T @ Xc / Nk[j] + self.reg * np.eye(d)
            if abs(ll - ll_old) < self.tol:
                break
            ll_old = ll
        self.n_params = self.k * (d + d * (d + 1) / 2) + (self.k - 1)   # means + covs + weights
        return self

    def predict_proba(self, X):
        X = np.asarray(X, float)
        R = np.column_stack([self.pi[j] * multivariate_normal.pdf(X, self.mu[j], self.Sig[j], allow_singular=True)
                             for j in range(self.k)])
        return R / (R.sum(1, keepdims=True) + 1e-300)

    def predict(self, X):
        return self.predict_proba(X).argmax(1)

    def bic(self, X):
        ll = self.loglik_[-1]
        return -2 * ll + self.n_params * np.log(len(X))