"""Datasets and plots for the live page on choosing k (Homework 6, Problem 3).

The live page downloads this file into the in-browser Python session and runs
`from kmeans_tools import *`, so every exercise starts from the same helpers.
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

KS = range(1, 9)  # the values of k we try everywhere on the page


# ---------------------------------------------------------------- datasets

def groups(centers, m=50, spread=0.5, seed=0):
    """m Gaussian points of standard deviation `spread` around each center."""
    rng = np.random.default_rng(seed)
    centers = np.asarray(centers, dtype=float)
    return np.vstack([c + spread * rng.standard_normal((m, 2)) for c in centers])


def four_blobs(spread=0.5, seed=0):
    """Four groups at the corners of a square of side 7."""
    return groups([[0, 0], [7, 0], [0, 7], [7, 7]], spread=spread, seed=seed)


def pairs_of_blobs(seed=0):
    """Four groups arranged as two far-apart pairs."""
    return groups([[-10, 0], [-10, 3], [10, 0], [10, 3]], spread=0.4, seed=seed)


def uniform_square(n=200, seed=0):
    """n points spread uniformly over the unit square: no clusters at all."""
    return np.random.default_rng(seed).uniform(0, 1, size=(n, 2))


def two_rings(n=200, seed=0):
    """Two noisy concentric circles, as in the k-means lecture."""
    rng = np.random.default_rng(seed)
    radius = np.where(np.arange(n) < n // 2, 1.0, 3.0)
    angle = rng.uniform(0, 2 * np.pi, n)
    ring = np.column_stack([radius * np.cos(angle), radius * np.sin(angle)])
    return ring + 0.15 * rng.standard_normal((n, 2))


def units_example(seed=0):
    """Four groups; the first feature is measured in decimeters."""
    return groups([[0, 0], [6, 0], [0, 4], [6, 4]], spread=0.5, seed=seed)


def many_blobs(seed=0):
    """Ten groups of 30 points on a 5 x 2 grid."""
    centers = [[5 * i, 5 * j] for i in range(5) for j in range(2)]
    return groups(centers, m=30, spread=0.5, seed=seed)


# ---------------------------------------------------------------- plots

def elbow_and_silhouette(X, ks=KS):
    """k-means cost and silhouette score for each k (silhouette needs k >= 2)."""
    costs, sils = [], []
    for k in ks:
        model = KMeans(n_clusters=k, n_init=10, random_state=0).fit(X)
        costs.append(model.inertia_)
        sils.append(silhouette_score(X, model.labels_) if k >= 2 else np.nan)
    return np.array(costs), np.array(sils)


def report(X, title="", ks=KS, k_show=None):
    """Print and plot the k-means cost and the silhouette score for each k.

    The data panel is colored by the k-means clustering with k = k_show,
    or with the k of largest silhouette score when k_show is None.
    """
    ks = list(ks)
    costs, sils = elbow_and_silhouette(X, ks)
    if k_show is None:
        k_show = ks[int(np.nanargmax(sils))]
    labels = KMeans(n_clusters=k_show, n_init=10, random_state=0).fit(X).labels_

    fig, ax = plt.subplots(1, 3, figsize=(12, 3.6))
    ax[0].scatter(X[:, 0], X[:, 1], c=labels, cmap="tab10", s=10)
    ax[0].set_title(f"{title} (k = {k_show})")
    ax[0].set_aspect("equal", adjustable="datalim")
    ax[1].plot(ks, costs, "o-")
    ax[1].set_xlabel("k")
    ax[1].set_title("k-means cost")
    ax[1].set_ylim(bottom=0)
    ax[2].plot(ks, sils, "o-")
    ax[2].set_xlabel("k")
    ax[2].set_title("silhouette score")
    for a in ax[1:]:
        a.set_xticks(ks)
        a.grid(alpha=0.3)
    fig.tight_layout()
    plt.show()

    print(f"{title}\n  k        cost   silhouette")
    for k, c, s in zip(ks, costs, sils):
        print(f"{k:3d} {c:11.1f}   {'' if np.isnan(s) else f'{s:10.3f}'}")